Hesitation-Aware On-Policy Distillation for Diffusion Language Models
cs.CL, cs.AI
Submitted: 2026-09-27
Updated: 2026-09-27
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion large language models (dLLMs) generate text by iterative unmasking.
Terminology
Abstract
Diffusion large language models (dLLMs) generate text by iterative unmasking. At each denoising step, a dLLM proposes a token at every masked position, but the decoder commits only a confident subset of these proposals. Trace-based on-policy distillation (TOPD) builds on this process by matching the student to a stronger teacher, yet only at the committed positions. We argue that this discards much of the useful signal, which resides in the uncommitted proposals, where the student has made a prediction but is not yet confident enough to commit it. We call these proposals hesitations. In our pilot study on an SDAR-4B student, hesitations make up only 24% of supervisable state-position pairs but carry 66% of the teacher-student divergence. To exploit this signal, we propose Hesitation-Aware On-Policy Distillation (HOPD), which extends teacher distribution matching to every masked position of each denoising step. Because hesitations are not equally informative, we further allocate supervision using hindsight from the completed trajectory, placing more weight on positions whose proposal was later disagreed with the final token and on blocks where first-step proposals rarely survive. Since both models already produce distributions at all masked positions, HOPD requires no additional forward passes over TOPD. The only extra cost is evaluating the loss at more positions. With SDAR-1.7B and SDAR-4B students distilled from TraDo-8B-Instruct, HOPD achieves the best average score among the evaluated methods on five math and coding benchmarks, under both static and dynamic decoding and at both scales. It also speeds up decoding. On SDAR-4B, the HOPD student hesitates less and commits 11% more tokens per denoising step than TOPD, while reaching higher accuracy.
Sources
- LLaDA2.0: Scaling Up Diffusion Language Models to 100B
- SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training
- Training Verifiers to Solve Math Word Problems
- dOPSD: On-Policy Self-Distillation for Diffusion Language Models
- Re-evaluating Confidence Remasking in Masked Diffusion Language Models
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models
- Measuring Mathematical Problem Solving With the MATH Dataset
- Distilling the Knowledge in a Neural Network
- INTELLECT-1 Technical Report
- Trace-Based On-Policy Distillation for Masked Diffusion Language Models
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models
- Deferred Commitment Decoding for Diffusion Language Models
- Dream 7B: Diffusion Large Language Models
- Simple Policy Gradients for Reasoning with Diffusion Language Models
- DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning
- Stabilizing Reinforcement Learning for Diffusion Language Models
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering