Research map / Probabilistic and scientific modeling / Molecules and proteins
Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows masked diffusion models train on computationally hard infilling subproblems, and that adaptive (top-probability-margin) unmasking order at inference lets pretrained MDMs solve Sudoku/Zebra puzzles and boost LLaDA 8B on coding/math, beating teacher-forced ARMs on logic puzzles.
Problem
Masked diffusion models must learn exponentially many infilling subproblems during training, and vanilla random-order inference evaluates poorly trained marginals; it is unclear whether MDMs' inference flexibility can outweigh training complexity.
Approach
Instead of retraining, choose which positions to unmask at inference adaptively using model confidence (top probability or top-probability margin), so the sampler sidesteps hard subproblems; theory based on planted CSP hardness shows why training on all masking problems is hard for order-agnostic learners.
Key results
- On Sudoku, top-probability-margin adaptive MDM inference reaches 89.49% accuracy versus 6.88% for vanilla MDM inference, above the ARM with ordering at 87.18%.
- On Zebra puzzles, adaptive MDM inference reaches 98.3% (top prob margin) and 98.5% (top prob) versus 76.9% vanilla and 91.17% for ARM with ordering.
- For LLaDA 8B, top-probability-margin inference improves HumanEval-Single from 31.8% to 33.5% and Math from 28.5% to 34.3% over vanilla inference.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace vanilla inference with adaptive inference | MDM · L&O-NAE-SAT | accuracy: 78.06% -> 93.76%; 75.70% -> 93.54%; 74.60% -> 92.21%; 67.94% -> 90.01%; 62.84% -> 88.91% | Helped |
| use Top probability margin ordering instead of vanilla decoding | 6M MDM · Sudoku | accuracy: 6.88% -> 89.49% | Helped |
| use Top probability ordering instead of vanilla decoding (ablation) | 6M MDM · Sudoku | accuracy: 6.88% -> 18.51% | Helped |
| use Top probability margin ordering instead of Top probability ordering (ablation) | 6M MDM · Sudoku | accuracy: 18.51% -> 89.49% | Helped |
| use Top probability ordering instead of vanilla decoding (ablation) | 19M MDM · Zebra puzzle | accuracy: 76.9 % -> 98.5 % | Helped |
| use Top probability margin ordering instead of vanilla decoding | 19M MDM · Zebra puzzle | accuracy: 76.9 % -> 98.3 % | Helped |
| use adaptive inference instead of vanilla inference | LLaDa 8B · HumanEval-Multi | accuracy: 16.5% -> 20.8% (Top probability); 16.5% -> 25.4% (Top prob. margin) | Helped |
| use Top probability margin instead of Top probability ordering (ablation) | LLaDa 8B · Math | accuracy: 31.3% -> 34.3% | Helped |
| use Top probability ordering instead of vanilla inference (ablation) | LLaDa 8B · MMLU | accuracy: 33.2% -> 36.5% | Helped |
| use Top probability margin instead of vanilla inference | LLaDa 8B · MMLU | accuracy: 33.2% -> 35.4% | Helped |
| use Top probability ordering instead of vanilla inference (ablation) | LLaDa 8B · ROCStories | accuracy: 21.23% -> 21.10% | No clear effect |
| use Top probability margin instead of vanilla inference | LLaDa 8B · ROCStories | accuracy: 21.23% -> 21.41% | Helped |
Surprising findings
- Adaptive MDM inference with no retraining outperforms ARMs explicitly trained with ground-truth decoding order on Sudoku and Zebra puzzles.
- Top-probability ordering can fail where top-probability margin succeeds: 18.51% vs 89.49% on Sudoku with the same 6M model.
Limitations
- Adaptive strategies are simple heuristics (top probability and margin) and the authors note better ordering strategies are needed.
- Hardness results rely on conjectures (1RSB cavity prediction) and average-case assumptions rather than unconditional proofs.
- Sudoku dataset differs from prior work, making direct comparison with previously reported near-100% results difficult.
Experimental setup
Datasets: SlimPajama Sudoku puzzle dataset Zebra puzzle dataset L&O-NAE-SAT synthetic distribution
Benchmarks: Sudoku Zebra puzzle HumanEval MMLU Math ROCStories SlimPajama likelihood
Metrics: accuracy generative perplexity entropy log-likelihood scaling
Baselines: vanilla MDM inference ARM without ordering ARM with ordering LLaDA vanilla inference top probability
Builds on: masked diffusion models score-entropy loss SEDD LLaDA
Closest papers by meaning
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.92)
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.91)
- Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.91)
- Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.90)
- From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.90)
- A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Curse (NeurIPS 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.