sudoku-dit

SudokuDiT

A compact Diffusion Transformer (1.28 M params) that solves Sudoku as a masked discrete diffusion (MDLM-style) denoiser: it fills a puzzle by iteratively un-masking the cells it is most confident about, MaskGIT-style.

• Architecture: DiT with adaLN-Zero conditioning — hidden=128, heads=4, blocks=4; per-cell token + 2-D positional + 3x3-box embeddings, plus a timestep.
• 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")     # downloads model.safetensors + config.json
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 (recommended)
# fixed-budget alternative: solver.solve(x, num_steps=81)

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

Performance

Held-out validation puzzles, adaptive decoder (tau = 0.999):

Solve rate is 100 % up to ~39 blanks and ~98 % on the hard 50+ blank tail (where the adaptive decoder spends more steps). exact_match is lower than valid only because the dataset's high-blank puzzles aren't always uniquely solvable, so the model may return a different valid grid.

Training

• Data: Ritvik19/Sudoku-Dataset (~17 M puzzles), tokenized on the fly.
• Objective: masked cross-entropy over masked cells, conditional (given clues are never masked), with a linear masking schedule.
• ~50 k steps, AdamW, batch 256. Kept deliberately small — larger models sit at the entropy floor far longer before the loss breaks through.

Limitations

• Not a guaranteed solver: a small fraction of very hard (60+ blank) boards come out invalid.
• Trained only on standard 9x9 Sudoku from the dataset above.