Trace-Based On-Policy Distillation for Masked Diffusion Language Models

2026-07-18

AI Summary

This work addresses the scarcity of efficient, dense supervision signals for inference-oriented post-training of diffusion language models by proposing Trajectory-aligned Online Policy Distillation (TOPD). The framework uniquely integrates online policy distillation with diffusion language models, leveraging the target model’s own denoising trajectories to provide token-level supervision from a teacher model’s token distribution at partially denoised states. By employing reverse KL divergence as the loss function, TOPD enables reward-free transfer of reasoning capabilities. Evaluated on the MATH500 benchmark, the SDAR-4B-Chat model achieves accuracy comparable to reinforcement learning approaches—improving static and dynamic settings by +5.7 and +4.5 points, respectively—while reducing rollout iterations by a factor of four and enhancing computational efficiency by approximately 96×.

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.

Problem

• masked diffusion language models
• reasoning-oriented post-training
• on-policy distillation
• teacher-supervised learning
• denoising trajectory

Innovation

• trace-based distillation
• on-policy learning
• masked diffusion language models
• Reverse-KL objective
• reasoning distillation