# Adaptive Order Policies for Masked Diffusion

arXiv: https://arxiv.org/abs/2606.00295 · HTML: https://mlsanity.com/paper/2606.00295/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 28 of its 41 references that are in the map.
- First posted on arXiv 29 May 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 28 of 28)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [CoQA: A Conversational Question Answering Challenge](https://mlsanity.com/paper/1808.07042/) (2018)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning](https://mlsanity.com/paper/2410.14157/) (2024)
- [Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking](https://mlsanity.com/paper/2505.24857/) (2025)
- [Riemannian Diffusion Models](https://mlsanity.com/paper/2208.07949/) (2022)
- [Diffusion Language Models Are Versatile Protein Learners](https://mlsanity.com/paper/2402.18567/) (2024)
- [GenMol: A Drug Discovery Generalist with Discrete Diffusion](https://mlsanity.com/paper/2501.06158/) (2025)
- [Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models](https://mlsanity.com/paper/2509.23653/) (2025)
- [Learning Unmasking Policies for Diffusion Language Models](https://mlsanity.com/paper/2512.09106/) (2025)
- [Learning-Order Autoregressive Models with Application to Molecular Graph Generation](https://mlsanity.com/paper/2503.05979/) (2025)
- [Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies](https://mlsanity.com/paper/2510.05725/) (2025)
- [Planner Aware Path Learning in Diffusion Language Models Training](https://mlsanity.com/paper/2509.23405/) (2025)
- [DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking](https://mlsanity.com/paper/2603.01367/) (2026)
- [Guided Star-Shaped Masked Diffusion](https://mlsanity.com/paper/2510.08369/) (2025)

## Abstract (arXiv)

Masked diffusion models have seen great success in capturing data distributions over discrete sequences in domains such as text and proteins. These models generate data by iteratively unmasking tokens starting from a fully masked sequence, with the unmasking order typically chosen at random or using a heuristic based on denoiser probabilities. In this work, we propose a scheme for learning the unmasking order using an additional lightweight policy network on top of a diffusion model. Our proposed loss reweights terms in the masked diffusion loss according to policy probabilities, and results in a policy that prefers positions where the denoiser is more likely to be correct. We study this loss in two settings: (i) training solely the policy while using a frozen pre-trained denoiser, and (ii) training the policy and denoiser jointly with the weighted loss to allow for mutual adaptation. We demonstrate that our approach outperforms common heuristics on problems that are sensitive to token ordering, such as combinatorial problems, proteins as well as various coding and language tasks.