Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
Pivot-SD trains masked diffusion LLMs only on high-impact token commitments, improving LLaDA-8B-Instruct math and code results with just 200 questions.
Pivot-SD is an offline self-distillation framework for masked diffusion language models that supervises only 'pivots' — denoising commitments selected by an information-gain metric measuring uncertainty reduction over remaining masked positions. Pivots from successful trajectories get cross-entropy training while failed-trajectory pivots get targeted unlikelihood loss, leaving the rest untouched. Using only 200 questions with four rollouts each, it improves LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines on math and code benchmarks.
- Selects high-impact 'pivot' commitments via uncertainty-reduction information gain
- Successful pivots trained with cross-entropy; failed pivots with unlikelihood loss
- Needs only 200 questions with four rollouts each
- Beats full-sequence SFT and budget-matched diffusion RL on LLaDA-8B-Instruct
Full article163 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2610.03665