Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models

Pivot-SD. Efficient Self-Distillation for Masked Diffusion Language Models. By Seo Hyun Kim and five others, from KAIST AI, the University of Toronto, and one other institution. Published on the second of October. No code repository is linked to it.

Masked diffusion language models generate in parallel, offering an alternative to autoregressive models for complex reasoning. They come with their own credit-assignment problem. During denoising, a few commitments sharply reduce the uncertainty over the positions still masked, and in doing so shape much of the final response. Most post-training recipes for these models ignore that signal when choosing what to train on. They train on the final text, or they assign rewards to whole denoising steps, instead of picking out the individual commitments that actually decided the answer.

Pivot-SD is an efficient, offline self-distillation framework that supervises only those high-impact commitments, which it calls pivots. Self-distillation means the model learns from its own rollouts. Offline means those rollouts are generated first and trained on afterwards. And supervising only the pivots means the training signal is concentrated on the few decisions that shaped each response, rather than spread across every token or every step.

Pivots are chosen with an information-gain metric: how much a commitment reduces uncertainty over the remaining masked positions. The trajectories are then split by outcome. Pivots from successful trajectories are trained with cross-entropy, reinforcing them. Pivots from failed trajectories are trained with targeted unlikelihood, pushing the model away from them, and the rest of a failed trajectory is left untouched.

The data requirement is small: two hundred questions, with four rollouts each. With that, Pivot-SD improves LLaDA-8B-Instruct over supervised fine-tuning on full sequences, and over diffusion reinforcement learning baselines matched for budget, across math and code benchmarks.

The method's bet is about where the learning signal lives. If a handful of commitments decide a response, then training on those and leaving the rest alone should be both cheaper and sharper, and penalising only the pivots of a failure avoids punishing tokens that were never the problem.