aiqbench analysis
This research introduces a lightweight policy network to optimize the token unmasking order in masked diffusion models, moving beyond random or heuristic-based approaches. By reweighting the training loss based on denoiser confidence, the method improves generation quality for discrete sequence tasks like protein folding and code synthesis. The study validates that jointly adapting the policy and denoiser yields superior performance over frozen pre-trained baselines on ordering-sensitive problems.
Adaptive Order Policies for Masked Diffusion
arXiv:2606.00295v2 Announce Type: replace Abstract: 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.
Abstract: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.
Machine Learning (cs.LG)
arXiv:2606.00295 [cs.LG]
(or arXiv:2606.00295v2 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2606.00295 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 262 words
Related coverage
Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling
arXiv · Aug 27, 2026
$D^{2}R^{2}$: Discrete Diffusion with Regulation Reinforcement for Single-Cell Perturbation Prediction
arXiv · Aug 18, 2026
Build a Physical AI model factory with NVIDIA Cosmos 3 on SageMaker HyperPod
AWS · Sep 4, 2026
Improving CoT Monitorability of Evaluation Awareness via Verbalization Training
LessWrong · Oct 2, 2026