Search papers, labs, and topics across Lattice.
Masked diffusion language models (dLMs) present a parallel alternative to autoregressive models but struggle with credit assignment due to the disproportionate impact of early denoising steps. Pivot-SD addresses this by employing an offline self-distillation framework that selectively trains on high-impact "pivot" tokens identified by an information-gain metric. This targeted approach, using cross-entropy for successful pivots and targeted unlikelihood for failed ones, significantly improves performance on math and code tasks with minimal data.
A few critical denoising steps in diffusion language models dictate the entire output, and training only these "pivots" dramatically boosts performance.
Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning. However, they face a distinct credit-assignment challenge, since a few commitments during denoising sharply reduce the uncertainty over the remaining masked positions and shape much of the response. Most post-training recipes for dLMs do not use this signal to decide which tokens to train on: they typically train on the final text or assign rewards to whole denoising steps, rather than selecting the individual commitments that shape the response. We introduce Pivot-SD, an efficient offline self-distillation framework that supervises only these high-impact commitments (pivots). Pivot-SD selects pivots using an information-gain metric measuring uncertainty reduction over the remaining masked positions. Pivots from successful trajectories are trained with cross-entropy, and pivots from failed trajectories with targeted unlikelihood, leaving the rest of the failed trajectory untouched. Using only 200 questions and four rollouts each, Pivot-SD improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines across math and code benchmarks.
