Hesitation-Aware On-Policy Distillation for Diffusion Language Models

2026-09-27

AI Summary

This study addresses the limitation of existing diffusion language model distillation methods that overlook signals from uncommitted, or "hesitant," positions. To this end, it proposes a hesitation-aware online policy distillation approach built upon a trajectory-level framework. By introducing a hindsight weighting mechanism, this work pioneers the identification and upweighting of highly informative hesitant proposals. This expands the supervision scope without requiring additional forward passes, enabling precise alignment between all masked positions and the teacher distribution to optimize the denoising process. Experimental results demonstrate that the proposed method achieves state-of-the-art average performance on mathematical and code generation benchmarks, significantly improving both decoding efficiency and single-step commitment rates.

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.

Problem

• Diffusion Language Models
• Knowledge Distillation
• On-Policy Distillation
• Hesitation Proposals
• Token Unmasking

Innovation

• Diffusion Language Models
• On-Policy Distillation
• Hesitation-Aware
• Iterative Unmasking
• Knowledge Distillation