Tiny Recursive Model

This repo is a modular PyTorch implementation of a Tiny Recursive Model (TRM) for Sudoku reasoning experiments. It follows the paper-style idea from "Less is More: Recursive Reasoning with Tiny Networks": a small shared network recursively improves an answer state y and latent reasoning state z conditioned on the input puzzle embedding x.

The code is set up for training functionality and smoke verification. It does not run full paper-scale training by default.

What It Does

The model learns a supervised Sudoku mapping from an 81-token puzzle to an 81-token solution. Puzzles use 0 for blanks and digits 1..9 for clues. The model predicts 9 logits for each cell.

Core TRM flow:

1. Embed the Sudoku puzzle as x.
2. Initialize answer state y and latent reasoning state z.
3. Recursively update z = net(x + y + z) for n latent steps.
4. Refine y = net(y + z).
5. Repeat this for T recursion cycles, with earlier cycles run without gradients and the final cycle trained.
6. Train with deep supervision, digit cross entropy, halt prediction BCE, and EMA parameter tracking.

Repository Structure

.
├── TRM.py                         # Compatibility imports for older single-file usage
├── main.py                        # CLI wrapper for training
├── pyproject.toml                 # Package metadata and dependencies
├── datasets/
│   └── sudoku_tiny.csv            # Tiny bundled Sudoku dataset for smoke tests
├── docs/
│   └── assets/                    # Generated README graph PNGs
├── trm/
│   ├── config.py                  # ModelConfig and TrainConfig
│   ├── layers.py                  # RMSNorm, SwiGLU, MLP-mixer, recursive network
│   ├── model.py                   # TinyRecursiveModel and TRMState
│   ├── data.py                    # Sudoku CSV parser and Dataset
│   ├── augment.py                 # Sudoku-valid augmentation utilities
│   ├── losses.py                  # Digit CE and halt BCE
│   ├── metrics.py                 # Cell and puzzle accuracy
│   ├── ema.py                     # EMA helper
│   └── train.py                   # Deep-supervision training loop and CLI
├── tests/                         # Unit and smoke tests
├── experiments/
│   ├── compare_constraint_loss.py  # Baseline-vs-improved training comparison
│   ├── plot_constraint_comparison.py
│   └── results/                   # CSV metrics used by plots
└── benchmarks/
    └── benchmark_sudoku.py        # Lightweight timing benchmark

Setup

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

If you already have the local environment active, just run the install command.

Dataset Format

CSV files need two columns:

puzzle,solution
530070000600195000...,534678912672195348...

Rules:

• puzzle must be 81 characters.
• solution must be 81 characters.
• puzzle uses 0 or . for blanks.
• solution uses only 1..9.
• All puzzle clues must match the solution.

The included datasets/sudoku_tiny.csv is only for verifying the pipeline. Replace it with a larger Sudoku dataset for real training.

Smoke Train

Run a tiny CPU-friendly training check:

python -m trm.train --max-steps 2 --batch-size 2 --d-model 64 --token-mlp-dim 64 --channel-mlp-dim 128 --max-supervision-steps 2

Use a custom dataset:

python -m trm.train --dataset-path datasets/your_sudoku.csv --max-steps 10

For a larger experiment, increase --d-model, --n-latent-steps, --recursion-cycles, --max-supervision-steps, and --max-steps. Paper-like Sudoku recursion uses --n-latent-steps 6 --recursion-cycles 3, but full training requires a real dataset and substantial compute.

Improved Training Objective

The base implementation uses the original supervised objective: digit cross entropy plus halt BCE. This repo also includes an optional Sudoku-aware auxiliary loss that keeps the same TRM architecture but encourages predicted probabilities to satisfy Sudoku row, column, and 3x3 box constraints.

Enable it with:

python -m trm.train --max-steps 100 --sudoku-constraint-weight 0.5

Compare the original CE-only objective against the improved objective:

python experiments/compare_constraint_loss.py --max-steps 800 --batch-size 3 --constraint-weight 0.5 --learning-rate 0.0001 --weight-decay 0

Generate the experiment plots:

python experiments/plot_constraint_comparison.py

The generated metrics are stored in experiments/results/constraint_comparison_metrics.csv, and the README figures are stored in docs/assets/.

Tiny Experiment Results

Both runs use the same tiny TRM architecture and seed. The only difference is the objective: original CE-only training versus CE plus the Sudoku constraint auxiliary loss. Open this README in Markdown preview or on GitHub to see the actual graph PNGs rendered below.

Validation digit loss comparison

The improved objective produced a lower final validation digit loss on the tiny augmented comparison run: 5.8145 versus 5.8911.

Metric deltas from improved objective

Positive bars mean the improved objective beat the original objective. In this run, the improvement shows up in validation digit loss while cell accuracy stays tied on the tiny dataset.

Cell accuracy after tiny run

The cell-accuracy plot is included as a sanity check: the tiny 3-example dataset is too small to claim real Sudoku quality, but it confirms both variants were evaluated on the same setup.

Tests

pytest

The tests cover:

• Sudoku parsing and clue preservation.
• Model forward shapes.
• Loss backpropagation.
• One-step training smoke run.

Benchmarks

python benchmarks/benchmark_sudoku.py

The benchmark runs a small recursive forward pass and a two-step train smoke benchmark. It is meant to confirm runtime behavior, not report final model quality.