trm-sudoku-extreme

Tiny Recursive Model (TRM) trained from scratch to solve Sudoku-Extreme puzzles. One small network recursively refines a latent state and an answer, from Less is More: Recursive Reasoning with Tiny Networks. trm.py in this repo is a single-file port of the official code and doubles as the loader.

🚀 Usage

Solve one puzzle (. for blanks):

uv run https://huggingface.co/jgalego/trm-sudoku-extreme/resolve/main/trm.py solve .9...12...3..284.6.6.....8..7....14..2..5.7....3.....2.1.9........7.5....872.6..4

From Python, with trm.py on the path:

from trm import load, solve_grids

model = load("jgalego/trm-sudoku-extreme")
print(solve_grids(model, [".9...12...3..284.6.6.....8..7....14..2..5.7....3.....2.1.9........7.5....872.6..4"]))

🏋️ Training

Model	TRM-MLP (MLP over positions), 5.0M parameters
Recursion	T=3, n=6, 16 supervision steps, 2 layers, hidden 512
Data	1,000 training puzzles, 1,000 augmentations each (digit, row, column and transpose symmetries)
Steps	65,104 of 65,104, batch size 768
Optimizer	AdamATan2, lr 0.0001, betas (0.9, 0.95), weight decay 1.0, 2000 warmup steps; sign SGD for the puzzle embedding
EMA	0.999
Hardware	NVIDIA A10G from step 0; NVIDIA RTX PRO 6000 Blackwell Server Edition from step 21,000, 9.6 h in total

📊 Results

Sudoku-Extreme test split, 16 supervision steps without early halting. Exact accuracy counts fully solved puzzles; cell accuracy counts all 81 cells, givens included.

Puzzles	Step	Exact accuracy	Cell accuracy
422,786 (full)	65,104	86.8%	95.1%

Reported in the paper (Table 4) on the full test set (423k puzzles):

Model	Params	Exact accuracy
TRM-MLP	5M	87.4%
TRM-Att	7M	74.7%
HRM	27M	55.0%
Direct prediction	27M	0.0%

⚠️ Limitations

• Solves 9x9 Sudoku only; other grid sizes or rule variants need retraining.
• Answers are not checked against the Sudoku rules. Verify a solution before relying on it.
• Trained on a small seeded subset of puzzles; accuracy varies with the seed and run length.