File size: 4,292 Bytes
0f956c4 2716eba 0f956c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | ---
license: apache-2.0
pipeline_tag: other
library_name: pytorch
tags:
- maze
- path-planning
- masked-diffusion
- diffusion-transformer
- dit
- reasoning
datasets:
- sapientinc/maze-30x30-hard-1k
---
# 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`](https://huggingface.co/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
```python
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):
| metric | value |
|--------|-------|
| valid_shortest (**the honest metric**) | **52.4 %** |
| valid simple S->G path | 61.8 % |
| + restart sampling, k = 32 | **57.2 %** valid_shortest |
| exact match vs reference | 6.9 % |
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`](https://huggingface.co/datasets/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.
|