arXiv preprint ·

Towards Unbiased On-Policy Distillation for Block Diffusion Language Models

Zaiquan Yang, Fei Wei, Yong Wang, Yudong Han, Yiyu Li, Zhuofan Zong, Gerhard Petrus Hancke, Xiangxiang Chu and Rynson WH Lau

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Computation & Language / NLP

Abstract

On-policy distillation (OPD) has emerged as an effective post-training paradigm for language models, with recent efforts extending it to block diffusion language models (BDLMs). However, existing studies focus almost exclusively on small block sizes, leaving distillation into student models with larger blocks underexplored. In this work, we investigate this regime and reveal two critical optimization biases that induce severe training instability. First, mismatched block boundaries between teacher and student cause \textbf{\textit{context misalignment}}, providing distorted supervisory signals that misguide student decoding. Second, even under aligned contexts, an \textbf{\textit{intrinsic optimization bias}} in OPD, where the student tends to rapidly absorb high-support signals while lagging on low-support updates, drives a premature confidence surge that traps weaker students in catastrophic overconfidence collapse. To resolve these, we propose \mbox{\textbf{Un-OPD}}, an unbiased on-policy distillation framework with two novelties for stabilizing BDLM training. First, Un-OPD introduces a boundary-aware step filtering strategy that eliminates context-misaligned decoding steps. Second, Un-OPD proposes moderating optimization intensity at high-support positions via a support-rebalanced confidence calibration, thereby bypassing overconfidence collapse. Beyond stability, we also introduce a rollout reuse mechanism to reduce rollout generation overhead. Extensive experiments on math reasoning and code generation benchmarks show that Un-OPD consistently stabilizes training and delivers superior performance while reducing wall-clock training time by approximately half.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Computation & Language / NLP

Get the few papers worth your time in Computation & Language / NLP every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics