Masked diffusion language models (dLMs) present a parallel alternative to autoregressive models for complex reasoning, yet they encounter a credit-assignment problem: a few commitments during denoising sharply reduce uncertainty over remaining masked positions and disproportionately shape the final output. Standard post-training recipes for dLMs typically train on the full final text or assign rewards to entire denoising steps, missing the individual high-impact tokens that drive performance. Pivot-SD addresses this by introducing an efficient offline self-distillation framework that supervises only these pivotal commitments.
Pivot-SD selects pivots using an information-gain metric that quantifies uncertainty reduction across the remaining masked positions. Commitments from successful trajectories are distilled with cross-entropy loss, while those from failed trajectories receive targeted unlikelihood training, leaving the rest of each failed trajectory untouched. This selective supervision ensures that only the most informative tokens are updated, preserving the diversity of the original diffusion process.
The method achieves significant improvements using minimal data: just 200 questions with four rollouts each are sufficient to improve LLaDA-8B-Instruct over full-sequence SFT and budget-matched diffusion RL baselines. The gains appear across math and code benchmarks, demonstrating that focusing on high-impact commitments is a broadly effective strategy for dLM post-training.
Pivot selection via information gain
The core innovation of Pivot-SD is the pivot-selection criterion. For each denoising step, the model computes the reduction in uncertainty (e.g., entropy) over all still-masked positions if a given token were committed. Tokens that produce the largest uncertainty drop are flagged as pivots. This metric can be computed efficiently because the diffusion process already maintains posterior distributions over masked positions; the pivot mask is simply the set of tokens whose commitment yields the greatest expected information gain.Dual supervision signals
Pivots derived from trajectories that reach a correct or high-reward final state are trained with standard cross-entropy, reinforcing the token choices that lead to successful outputs. Pivots from trajectories that fail or underperform are subjected to targeted unlikelihood, which explicitly discourages the model from repeating those low-impact or harmful commitments. The remainder of each failed trajectory—tokens that were not selected as pivots—is left unchanged, avoiding unnecessary updates that could disrupt the diffusion dynamics.Empirical results
In experiments across mathematical reasoning and code generation benchmarks, Pivot-SD consistently outperforms full-sequence supervised fine-tuning and diffusion RL baselines that operate under comparable compute budgets. The improvements are achieved with only 200 training questions and four rollouts per question, confirming the method's data and compute efficiency. The selective distillation of pivots not only accelerates convergence but also yields stronger final performance than approaches that treat every token equally.Limitations and analysis
The information-gain pivot selection relies on accurate uncertainty estimates from the diffusion model; if the model's posterior approximations are poorly calibrated, pivot quality may degrade. The method currently focuses on offline self-distillation from the same model family; extending the framework to cross-model distillation or online interaction remains open. Additionally, the 200-question budget is small relative to the scale of typical language model training, so results may vary on tasks with broader or more diverse reasoning patterns.Practical takeaways for dLM practitioners
Developers working with masked diffusion language models can implement Pivot-SD by computing uncertainty reduction per token during denoising, selecting the top pivots, and applying cross-entropy or unlikelihood losses accordingly. The approach requires minimal overhead—only a single pass to identify pivots and a standard distillation loop—and delivers measurable gains with a fraction of the data needed for full-sequence training. This makes it especially attractive for domains where labeled reasoning data are scarce but diffusion models are already available.Read the paper on arXiv