Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
The authors introduce a new way to teach diffusion-based large language models (dLLMs) how to reason better without using complicated reward signals. Their method, called trace-based on-policy distillation (TOPD), trains the dLLM by having it learn from a teacher model on the same intermediate steps it naturally goes through while generating text. This helps the model improve more efficiently and accurately in math reasoning tasks compared to previous reinforcement learning methods, using much less computation. Their approach balances dense guidance from the teacher with training that matches the model’s own behavior.
Authors
Haolin Ren, Ziyang Huang, Chenhao Yuan, Jun Zhao, Kang Liu
Abstract
Diffusion large language models (dLLMs) are a promising alternative to autoregressive generation. However, reasoning-oriented post-training for dLLMs remains challenging. Supervised fine-tuning (SFT) for dLLMs requires dense but often off-policy masked states, while reinforcement learning (RL) relies on sparse rewards or value modeling. This paper proposes \textbf{trace-based on-policy distillation (TOPD)}, a teacher-supervised framework that transfers reasoning ability to a target dLLM without reward estimation. The key idea is to supervise a dLLM on its own denoising trajectory, focusing on the trace-aligned token decisions that form the final response. Specifically, TOPD samples on-policy diffusion trajectories from the target dLLM, obtains teacher token distributions from a teacher model on the corresponding partially denoised states, and updates the target dLLM with a token-level Reverse Kullback-Leibler (Reverse-KL) objective. This design preserves dense teacher supervision while aligning training with the model's own denoising states. On mathematical reasoning benchmarks, TOPD enables SDAR-4B-Chat to match the MATH500 accuracy of its RL-trained counterpart TraDo-4B-Instruct, with gains of +5.7 under static evaluation and +4.5 under dynamic evaluation. Compared with the RL-trained counterpart, TOPD achieves this with 4$\times$ fewer rollout rounds, corresponding to an estimated 96.0$\times$ to-accuracy model-compute speedup.