Skip to content
Preprint

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

Jul 2026 · 0 citations · 41 references
Computer Science

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.

View source

Similar papers

Preprint Jul 2026

dOPSD: On-Policy Self-Distillation for Diffusion Language Models

dOPSD derives the teacher's privilege directly from the student's own denoising trajectory, evaluating masked positions using later, more-decoded steps of that same trajectory rather than an external label, so the teacher's advantage emerges from the model's own decoding process.

P. Dat, Qi Li, Xinchao Wang · 0 citations
Preprint Aug 2026

On-Policy Self-Distillation in Diffusion Models

The results support on-policy self-distillation as an efficient and analyzable approach to diffusion post-training by converting image-level reward guidance into explicit and continually refreshed intermediate supervision, thereby opening a path toward more efficient and diagnosable alignment.

Weina Zhou, Xiongwei Zhu, Lingdong Kong et al. · 0 citations
Preprint Aug 2026

Adaptive Supervised Anchoring for On-Policy Self-Distillation

Context quality is identified as a central bottleneck in on-policy self-distillation and the value of separating rollout-conditioned guidance from canonical supervision is demonstrated, demonstrating the value of separating rollout-conditioned guidance from canonical supervision.

Meilin Yang, Zixuan Ding, Jianhao Nie et al. · 0 citations
Preprint Aug 2026

Self-OPD: On-Policy Distillation for Flow Matching Models without Teacher

On-policy distillation (OPD), which leverages a pre-trained, specialized teacher model to provide dense supervisory signals, has achieved significant success in Large Language Models (LLMs) and has recently been adapted to flow matching models. However, this paradigm suffers from two major issues: First, training a separate, task-specific teacher for every new objective incurs high computational costs. Second, the discrepancy between teacher and student distributions often leads to compounding errors along the generation trajectory. In this paper, we introduce \textbf{Self-OPD}, a teacher-free OPD framework for flow matching models that turns the student's own self-exploration into step-wise supervision. At each timestep, Self-OPD branches the deterministic next-state prediction into $K$ stochastic SDE candidates, rolls them out with the ODE sampler, and compares their rewards against a deterministic self-reference baseline to obtain normalized advantages. The velocity field is optimized with an all-branch pull-push objective, where high-advantage branches attract the student and low-advantage branches repel it under direction-aware attenuation and SDE-variance normalization. For multi-objective alignment, Self-OPD fuses normalized scores at the reward level, avoiding direct gradient conflict. Experiments on single and mixed reward benchmarks show that Self-OPD outperforms prior RL and OPD methods without task-specific teachers.

Shiyi Zhang, Mushui Liu, Yunze Tong et al. · 0 citations
Preprint Jul 2026

Mask-Aware Policy Gradients for Diffusion Language Models

This work observes that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask, and formalizes this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term.

Haran Raajesh, Kulin Shah, Adam R. Klivans et al. · 1 citation
Preprint Aug 2026

DASH: Divergence-Adaptive Supervision Horizons for On-Policy Self-Distillation of Reasoning Models

DASH maps the gap between each local distillation signal and the sequence-level mean to an adaptive propagation gate and then uses these gates to control backward multi-step aggregation and improves over matched vanilla OPSD reruns on every benchmark at all three model scales.

Zhi-Yan Hou, Xinyu Tang, Hongyan An et al. · 1 citation