maze-dit

MazeDiT-30x30

A 1.38 M-parameter Diffusion Transformer that solves 30x30 maze path-planning as masked discrete diffusion. It labels every open cell as on-path or off-path, most confident first — it never traces a route.

A 131-cell shortest path recovered in 40 adaptive steps. Walls are near-black, cells not yet labelled stay slate, and path cells are tinted by the model's step-0 confidence (blue = unsure -> green = sure).

Same recipe as tchauffi/sudoku-dit, with the 9-digit vocabulary swapped for 3 tokens (wall / open / path).

• Architecture: GridDiT — DiT with adaLN-Zero conditioning, hidden=128, heads=4, blocks=4; per-cell token + 2-D positional embeddings, plus a timestep. No Sudoku box embedding.
• Code, training and an interactive web demo: https://github.com/tchauffi/nonet

Input / output

A maze is 900 tokens, row-major: 0 = [MASK] (an open cell to label), 1 = wall, 2 = open/off-path, 3 = path. The question marks walls and the two endpoints (start and goal, both token 3) and masks everything else; the solver clamps the givens.

Usage

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

solver = load_maze_solver("tchauffi/maze-dit")     # config records the cosine-high schedule
question = torch.tensor([[...]])         # (1, 900), see the encoding above
pred = solver.solve(question, num_steps=900, conf_threshold=0.999)

scripts/eval_maze30.py in the repo downloads the benchmark and reproduces the table below.

The metric matters here

Shortest paths on these mazes are massively non-unique — a median of ~10^7 distinct optimal routes per maze. Grading against the dataset's single reference answer therefore measures tie-break mimicry, not solving. We report valid_shortest: the prediction is a valid simple start-to-goal path and its length equals the BFS optimum. Both are checkable from the question alone, without the reference.

Performance

Full 1,000-maze test split, adaptive decoder (tau = 0.999):

The gap between 61.8 % valid and 6.9 % exact is the degeneracy above: the model routinely finds a correct path that is not the one the generator happened to emit.

For reference, HRM reports 74.5 % on this benchmark at 27 M parameters (~20x larger) under the exact-match protocol, which the degeneracy finding makes hard to compare directly.

Restart sampling climbs slowly here (52.4 -> 57.2 % over 32 attempts) and is still unplateaued — nothing like the near-doubling the same trick gives on Sudoku-Extreme. That says the residual failures are a systematic data ceiling, not decoding luck.

Training

• Data: the 1,000 training mazes provided by sapientinc/maze-30x30-hard-1k, plus dihedral x8 augmentation (the square's symmetry group).
• Objective: masked cross-entropy over masked cells, conditional (walls and endpoints are never masked), cosine-high masking schedule.
• ~18 k steps at the selected checkpoint, AdamW, batch 256, cosine lr decay.

Limitations

• Checkpoint selection used a 256-maze subset of the same test split; no separate validation split ships with the benchmark (HRM's protocol shares this).
• Memorization collapse: dihedral x8 over 1,000 mazes yields only ~8 k effective boards. Past ~20 k steps the model memorizes them — a 60 k-step run drives train cell-accuracy to 1.000 and test-valid to 0 %. This checkpoint is the step-18k peak; longer training is strictly worse.
• Trained only on 30x30 mazes of this generator; the released weights regenerate pos_embed for the grid size, but nothing about other sizes is tested.