Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models
cs.LG, cs.CL
Submitted: 2026-10-02
Updated: 2026-10-02
Comments: EMNLP 2026 Main (Oral)
License: http://creativecommons.org/licenses/by/4.0/
The gist: Masked diffusion language models (dLMs) offer a promising parallel alternative to autoregressive models for complex reasoning.
Terminology
Abstract
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.
Sources
- Program Synthesis with Large Language Models
- Evaluating Large Language Models Trained on Code
- Reasoning in Diffusion Large Language Models is Concentrated in Dynamic Confusion Zones
- Training Verifiers to Solve Math Word Problems
- Rethinking Reasoning with MDLMs: Early Exits, Post-hoc Reasoning, and Beyond
- Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
- Offline Learning and Forgetting for Reasoning with Large Language Models
- Diffusion-State Policy Optimization for Masked Diffusion Language Models
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation
- Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning
- Neural Text Generation with Unlikelihood Training
- GIFT: Guided Importance-Aware Fine-Tuning for Diffusion Language Models
- Improving Sampling for Masked Diffusion Models via Information Gain
- Dream 7B: Diffusion Large Language Models
- ACECODER: Acing Coder RL via Automated Test-Case Synthesis
- Simple Policy Gradients for Reasoning with Diffusion Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks