sudoku-dit-extreme-1k

SudokuDiT-Extreme-1k

A 1.28 M-parameter Diffusion Transformer trained on only 1,000 Sudoku-Extreme puzzles — the same small-data regime HRM and TRM report on. It solves Sudoku as a masked discrete diffusion denoiser: fill the cells you are most confident about, re-read the board, repeat.

This model solving a 51-blank board over 24 reveal steps. Given clues stay grey; filled cells are tinted by the model's step-0 confidence (blue = unsure -> green = sure).

• Architecture: DiT with adaLN-Zero conditioning — hidden=128, heads=4, blocks=4; per-cell token + 2-D positional + 3x3-box embeddings, plus a timestep.
• Sibling model: tchauffi/sudoku-dit, same architecture trained on a 3.8 M-board stream of ordinary Sudoku.
• Code, training and an interactive web demo: https://github.com/tchauffi/nonet

Input / output

A board is 81 tokens, row-major: 0 = empty / [MASK], 1..9 = digits. The solver clamps the given clues and only fills the blanks.

Usage

import torch
from nonet.hub import load_solver   # pip install git+https://github.com/tchauffi/nonet

solver = load_solver("tchauffi/sudoku-dit-extreme-1k")   # the config records the cosine-high schedule
puzzle = "402609001000000000700002084000080500030005007004060930008021406005000000147030000"
x = torch.tensor([[int(c) for c in puzzle]])      # (1, 81), 53 blanks
solution = solver.solve(x, conf_threshold=0.999)  # adaptive reveal

"".join(map(str, solution[0].tolist()))
# 482659371613874295759312684971483562836295147524167938398721456265948713147536829

This is also the default model for load_solver() with no arguments.

Performance

Held-out Sudoku-Extreme test puzzles, adaptive decoder (tau = 0.999). Because every solution is exactly checkable, a wrong answer can simply be retried: attempt 1 is the greedy decode, later attempts resample at temperature 1 (randomness enters only at the forced below-tau reveals). k is the attempt budget.

For reference, HRM reports 55.0 % on this benchmark at 27 M parameters (~21x larger). The comparison is not compute-matched: our 60.4 % spends 64 decode attempts per puzzle, and the curve is still unsaturated at k = 64. The honest reading is that a tiny diffusion denoiser plus a verifier covers this regime, not that it is cheaper per puzzle.

Training

• Data: 1,000 puzzles from sapientinc/sudoku-extreme, plus the Sudoku symmetry group sampled per example per batch (digit relabeling, transpose, band/row and stack/column permutations) — validity-preserving, the same family HRM/TRM use.
• Objective: masked cross-entropy over masked cells, conditional (given clues are never masked), cosine-high masking schedule (mass where deduction is hard).
• 30 k steps, AdamW, batch 256, lr 3e-4.

Augmentation is load-bearing. The identical run on the raw 1,000 boards without symmetry augmentation memorizes them and collapses to 0.1 % test-valid. With augmentation the effective board count is large enough to generalize.

Limitations

• Not a guaranteed solver: the hardest rating bands remain below 50 % even at k = 64.
• The headline number depends on a verifier and a 64x decode budget; single-shot accuracy is 22.7 %.
• Trained only on standard 9x9 Sudoku.