用自洽轨迹蒸馏让扩散语言模型学会推理,无需强化学习。
Trace-Based On-Policy Distillation for Masked Diffusion Language Models

- 以模型自身去噪轨迹为监督信号,实现密集且对齐的训练。
- 在数学推理任务上达到与强化学习相当的准确率,提升5.7%(静态)和4.5%(动态)。
- 相比强化学习,推理轮次减少4倍,计算效率提升96倍,适合高效训练。
扩散大语言模型(dLLMs)是自回归生成的有前景替代方案。然而,面向推理的后训练仍具挑战:监督微调(SFT)需要密集但常偏离策略的掩码状态,而强化学习(RL)依赖稀疏奖励或价值建模。本文提出一种基于轨迹的在线策略蒸馏(TOPD),无需奖励估计即可将推理能力迁移到目标dLLM。核心思想是监督目标dLLM在其自身去噪轨迹上的决策,聚焦形成最终回答的对齐标记序列。具体而言,从目标dLLM采样在线扩散轨迹,从教师模型获取对应部分去噪状态下的标记分布,并使用标记级反KL(Reverse-KL)目标更新目标模型。该设计保持密集教师监督,同时使训练与模型自身去噪状态对齐。在数学推理基准测试中,TOPD使SDAR-4B-Chat在MATH500上的准确率媲美其强化学习训练的对手TraDo-4B-Instruct,静态评估下提升+5.7,动态评估下提升+4.5。相较于强化学习训练对手,TOPD仅需4×更少的回滚轮次,相当于模型计算速度提升96.0×至准确率水平。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。