PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Seo Hyun Kim, Sunwoo Hong, Younwoo Choi, Chen-Hao Chao, Se-Young Yun, Rahul G. Krishnan

arXiv:2610.03665v1cs.LGcs.CL

TL;DR

Masked diffusion language models lack post-training methods that select the commitments shaping their responses, rather than training on final text or whole denoising steps. Pivot-SD uses information gain to select pivots and distills successful and failed trajectories; with 200 questions, it outperforms full-sequence SFT and diffusion RL on math and code benchmarks.

  • Problem

    Most dLM post-training recipes train on final text or whole denoising steps rather than selecting commitments that reduce uncertainty over remaining masked positions.

    Source: §1 · p. 1, §3.1 · p. 4

  • Method

    Pivot-SD selects high-information-gain commitments and replays their partially masked states, applying cross-entropy to successful trajectories and unlikelihood to failed ones.

    Source: §1 · p. 2

  • Results

    With 200 training questions, Pivot-SD gives LLaDA-8B-Instruct the highest mean accuracy on four math and code benchmarks, outperforming full-sequence SFT and diffusion RL.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Pivot-SD penalizes pivots from failed trajectories while leaving valid steps around them intact.

    Source: §5 · p. 8

    The objective has hyperparameters set per configuration, and some adjustment across domains and backbones is expected.

    Source: §6 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Masked diffusion models record when tokens are committed, but most post-training recipes do not use that record to select training tokens. Pivot-SD trains only high-impact commitments and reports stronger benchmark results than full-sequence SFT and diffusion RL baselines.

  • Most dLM post-training trains on final text or whole denoising steps, so influential commitments can be weighted like late fills and correct intermediate steps penalized with errors.Unlike autoregressive models, dLMs commit tokens in confidence-dependent order, recording which token was written where and at what step.Source: §1 · p. 1
  • Pivot-SD identifies pivots as commitments at denoising steps with the largest normalized reduction in entropy over remaining masked positions.These steps show marked entropy drops, while uncertainty changes little at other steps.Source: §1 · pp. 1–2
  • The offline method samples trajectories once, stores the partially masked states at pivot commitments, and replays those states for training.Correct-trajectory pivots receive cross-entropy; incorrect-trajectory pivots receive unlikelihood.Source: §1 · p. 2
  • With 200 training questions, Pivot-SD gives LLaDA-8B-Instruct the highest mean accuracy on four math and code benchmarks, outperforming full-sequence SFT and diffusion RL.The reported comparison includes RL runs trained on ten times as many questions for five times as many steps; Pivot-SD uses one-sixth to one-eighth of their wall-clock time.Source: §1 · p. 2
  • Pivot-SD computes loss on only 10 tokens per trajectory, under 4% of the 256-token response budget.The introduction also reports about 40% fewer FLOPs than budget-matched RL.Source: §1 · p. 2

2 Related Work

Prior work addresses diffusion-language-model reasoning, process supervision, and token weighting through different supervision units and signals. Pivot-SD instead targets selected commitments within denoising trajectories using verifier outcomes.

Diffusion language models

  • Diffusion generation passes through partially masked states rather than a fixed left-to-right order, and Pivot-SD takes supervision from that state sequence.The background traces masked diffusion from discrete-state approaches to LLaDA.Source: §2 · p. 2

Reasoning with masked diffusion language models

  • Recent masked-diffusion reasoning methods use online RL, whereas Pivot-SD samples trajectories once from a frozen model and localizes supervision using their outcomes.This offline approach can use both successful and failed trajectories.Source: §2 · p. 2

Process supervision and reasoning credit assignment

  • Process-supervision methods score textual reasoning steps in autoregressive solutions, while Pivot-SD supervises denoising commitments that can reshape many remaining positions.The commitment unit reflects that masked diffusion tokens are generated out of order.Source: §2 · pp. 2–3

Token importance and weighted fine-tuning

  • Token-selection methods score final-sequence positions, but Pivot-SD scores each commitment by its effect on positions still masked at that step.The final sequence does not record this step-specific effect; reported ablations compare pivot-only, selected-step uniform, entropy-selected, and random-pivot supervision.Source: §2 · p. 3

Positioning of our work

  • Pivot-SD supervises selected denoising commitments and uses verifier outcomes for positive and negative targets, without process reward models or step-level labels.Process supervision scores textual steps (Lightman et al., 2024; Wang et al., 2024), GIFT reweights final-sequence positions (Xu et al., 2025), and step-aware diffusion RL assigns rewards over denoising intervals (Chen et al., 2025; Xie et al., 2026; Oba et al., 2026).Source: §2 · p. 3

3 Methodology

Pivot-SD turns sampled denoising trajectories into offline, outcome-conditioned supervision for selected commitments rather than treating all masked tokens equally. It selects pivots by normalized entropy reduction, then applies cross-entropy or unlikelihood only at those commitments.

3.1 Preliminaries and Problem Setup

  • Standard masked SFT weights every masked position equally, whereas Pivot-SD selects commitments and sets each update's direction according to trajectory success.The setup defines a prompt x and target response y; the method addresses the difference between trivial fills and commitments that determine the rest of the response.Source: §3.1 · pp. 3–4

3.2.1 Trajectory Generation and Candidate Pivots

  • Each sampler commitment becomes a candidate pivot recorded by its denoising step, position, and sampled token, so supervision represents a state-conditioned commitment event.The same token at the same position can matter differently depending on how many neighboring positions remain masked.Source: §3.2.1 · p. 4

3.2.2 Information-Gain Pivot Selection

  • Pivot-SD ranks steps by average entropy reduction over positions that remain masked, excludes steps with changing candidate sets, and supervises the top K commitments.The score is comparable across steps; scoring reuses sampler forward passes, and the experiments use K = 10, with little change for 5 ≤ K ≤ 20.Source: §3.2.2 · p. 5

3.2.3 Training objective

  • Each selected pivot receives cross-entropy for a successful trajectory or targeted unlikelihood for a failed one, leaving other positions in failed trajectories untouched.The trajectory outcome determines the update direction; unlikelihood is applied only at selected pivots.Source: §3.2.3 · p. 5

3.3 Algorithmic Implementation

  • The algorithm samples trajectories from the frozen base model, assigns outcomes, selects top-K steps by information gain, and optimizes an offline dataset of selected pivots.The pipeline is summarized as trajectory sampling followed by fine-tuning.Source: §3.3 · p. 5, §3.2.3 · p. 5

3.4 Relation to SFT and Online RL

  • Pivot-SD combines offline supervised losses with outcome labels from successful and failed trajectories, while applying updates only to selected commitments.Because the pivot dataset is sampled once from the frozen base model, hyperparameters can be tuned offline on that dataset; training cannot stall on zero-reward batches observed for online RL at this budget.Source: §3.4 · pp. 5–6

4 Experiments

Under a small training budget, Pivot-SD leads the main 256-token benchmark comparison, and ablations, compute accounting, backbone transfer, and out-of-domain evaluation probe its behavior and scope.

4 Experiments

  • The experiments use 200 training questions per dataset and four trajectories per question; budget-matched RL uses the same questions and rollouts, while extended RL uses 2,000 questions and 5,000 steps.Evaluation covers MATH500, GSM8K, HumanEval+, and MBPP+ at 256 tokens, with 512-token evaluation measuring length transfer.Source: §4 · pp. 6–7

4.1 Main Results

  • At 256 tokens, Pivot-SD leads all four benchmarks, gaining 2.0 points on MATH over SFT-GT, 1.9 on GSM8K over 5,000-step diffu-GRPO, and 4.7 on HumanEval+ over SFT-SD.On MBPP+, its 0.8-point margin over SFT-GT is within one standard deviation; at 512 tokens it leads on MATH, GSM8K, and MBPP+, but scores 41.46 on HumanEval+ versus 42.07 for both the base model and budget-matched diffu-GRPO.Source: §4.1 · p. 7

4.2 What Makes Pivot-SD Work?

  • With positive and negative updates, pivot-local loss beats applying the loss to every masked token in the selected state by 2.5 to 5.1 points across all four benchmarks.With positive updates alone, neither target is consistently better; the authors say the localization gain appears with the negative branch, consistent with state-wide unlikelihood pushing down correct tokens in failed states.Source: §4.2 · pp. 7–8
  • Pivot-SD beats both random-selection controls on all four benchmarks; with unlikelihood, random steps and random tokens cost 7.5 and 2.6 MATH points, respectively.Random-Step Pivot scores 29.93 on MATH versus the untrained base model's 31.40; in the positive-only setting, entropy pivot stays within 2.3 points of Positive Pivot on every benchmark.Source: §4.2 · p. 8

4.3 Compute Efficiency

  • Pivot-SD samples 800 trajectories once in 2.5 hours and fine-tunes on cached pivots in 0.3 hours, totaling 2.8 hours; its FLOPs are about 1.65× lower than budget-matched RL.RL rollouts during training take 3.6–5.4 hours at 1,000 steps and 17.3–22.9 hours at 5,000 steps.Source: §4.3 · p. 8
  • Pivot-SD has the highest average accuracy while using about 8× less time than 5,000-step diffu-GRPO and 1.3 to 1.9× less time than budget-matched RL.Figure 3 plots average accuracy against total wall-clock training time on two GPUs.Source: Figure 3 · p. 7, §4.3 · p. 8

4.4 Backbone Generalization

  • On Dreamv0-Instruct-7B, Pivot-SD improves over the base model and both SFT baselines on all four benchmarks.The paper reports this as transfer of the gains to a second masked diffusion model.Source: §4.4 · p. 8

4.5 Out-of-Domain Transfer

  • With the same 200-question budget, Pivot-SD has the best out-of-domain overall average, 48.61, and the best average rank across training domains.It ranks first among GSM8K-trained models, second among AceCode-trained models, and third among MATH-trained models; Overall is the micro-average across source-target pairs.Source: §4.5 · p. 8

5 Conclusion

The conclusion presents Pivot-SD as an offline self-distillation method that trains only on high-impact commitments in masked diffusion language models. With 200 training prompts, it reports better performance than full-sequence SFT and online RL baselines at a fraction of their compute.

5 Conclusion

  • Pivot-SD uses information gain to select commitments where uncertainty over remaining positions drops most, then reinforces successful pivots with cross-entropy and penalizes failed ones with unlikelihood.Targeting failed pivots leaves valid steps around them intact.Source: §5 · p. 8
  • With only 200 training prompts, Pivot-SD outperforms full-sequence SFT and online RL baselines while using a fraction of their compute.Source: §5 · p. 8

6 Limitations

Pivot-SD’s design has open hyperparameter and scope constraints, while its credit assignment focuses on selected realized commitments and information-gain pivots.

6 Limitations

  • Some hyperparameters are set per configuration and may need adjustment across domains and backbones; an adaptive entropy-collapse rule remains open.Source: §6 · p. 8
  • Pivot-SD operates within a denoising block, so long-context reasoning requires extending step-conditioned credit assignment across block boundaries.Source: §6 · p. 9

A Credit-Assignment Landscape

  • Pivot-SD uniquely combines training on selected realized commitments with localized negative targets from failed trajectories among the compared methods.Source: A Credit-Assignment Landscape · p. 11

B Illustration of Information-Gain Pivot Selection

C Experimental Details and Configurations

The experiments use a shared response-length budget and decoding setup, with specified hardware, optimization settings, and evaluation lengths.

C Experimental Details and Configurations

D Multi-Seed Evaluation Protocol

The multi-seed protocol re-randomizes the offline training pipeline, but extended-budget RL and length-transfer results are single runs; the paper does not claim statistical significance.

D Multi-Seed Evaluation Protocol

  • Controlled-budget methods use three runs, but extended-budget RL baselines and 512-token length-transfer results are single runs.Source: D Multi-Seed Evaluation Protocol · p. 11
  • Each run re-samples the 200-question set, the trajectory pool, and the training seed, so reported variance covers the offline pipeline that produces Dpivot.Source: D Multi-Seed Evaluation Protocol · p. 11
  • The authors report means and standard deviations but do not claim statistical significance; code-benchmark standard deviations are rough estimates because the benchmarks are small and use three runs.Source: D Multi-Seed Evaluation Protocol · p. 12
  • On MATH, GSM8K, and HumanEval+, Pivot-SD’s one-standard-deviation intervals do not overlap those of the strongest competitor; MBPP+ is treated as a tie.The MBPP+ comparison is 47.12 ± 1.20 for Pivot-SD versus 46.32 ± 0.45 for SFT-GT.Source: D Multi-Seed Evaluation Protocol · p. 12

E Training Budget Sweep

Across the budget sweep, Pivot-SD remains competitive with larger-budget baselines, including at the largest training budget, although several exceptions occur.

E Training Budget Sweep

  • The sweep varies q across 50, 200, 400, and 2000 and retrains every method under the same protocol.Source: E Training Budget Sweep · p. 12
  • At q = 2000, Pivot-SD is best on all four benchmarks, even as online RL baselines train for 5,000 steps.Source: E Training Budget Sweep · pp. 12–13
  • With 200 questions, Pivot-SD matches or exceeds every baseline at every budget on GSM8K and HumanEval+, including baselines trained on ten times the data.Source: E Training Budget Sweep · p. 13
  • Baselines lead in four cases, usually by at most 1.7 points; the larger exception is wd1++ at q = 400 on MATH with 39.2.wd1++ reaches 35.4 at q = 2000 and 33.27 in the q = 200 multi-seed runs.Source: E Training Budget Sweep · p. 13

F Pivot Budget K

The pivot budget controls how much of each trajectory receives supervision; K = 10 lies in a strong range, while larger budgets reduce selectivity and random-step selection underperforms.

F Pivot Budget K

  • K = 10 is best or tied for best on all four benchmarks at K = 5, 10, and 20.The budget was fixed before experiments; Table 10 reports the sweep on a shared trajectory pool.Source: F Pivot Budget K · p. 13
  • At K = 40, MATH falls to 22.8 and both code benchmarks drop, while GSM8K is unchanged.Larger budgets admit steps that barely change remaining entropy, extending unlikelihood to positions that did not shape the trajectory.Source: F Pivot Budget K · p. 13
  • For K = 5–20, Random-Step reaches 17.6–26.6 on MATH, versus 37.8–38.6 for Pivot-SD, so sparsity alone does not explain the gain.Random-Step supervises the same number of positions as Pivot-SD.Source: F Pivot Budget K · p. 13

G Compute and FLOPs Accounting

Pivot-SD uses less compute than budget-matched RL and performs strongly on out-of-domain evaluations. Its compute accounting includes offline trajectory generation and the FLOPs, model calls, and supervised-token counts.

G Compute and FLOPs Accounting

  • Pivot-SD samples 800 offline trajectories with 256 sampler steps, for 204,800 generation NFE; its information-gain scores reuse those passes.Budget-matched RL uses 358,080 online NFE after rollout, policy, and scoring passes; supervised-token counts include only positions receiving a loss.Source: G Compute and FLOPs Accounting · pp. 13–14
  • 1,875 PFLOPs for Pivot-SD versus 3,093 PFLOPs for each budget-matched RL run, about 1.65× fewer total FLOPs.Most of the difference comes from generation: 1,678 versus 2,819 PFLOPs.Source: G Compute and FLOPs Accounting · p. 14

H Out-of-Domain Generalization

  • Pivot-SD has the best average among GSM8K-trained models and the best overall micro-average; its averages are within 0.7 points of best for MATH- and AceCode-trained models.Table 12 evaluates models trained on one domain on benchmarks from other domains.Source: H Out-of-Domain Generalization · p. 14

I Reward Sparsity in wd1++ and the Bootstrap Reward

The wd1++ bootstrap experiment examines sparse rewards on code, alongside comparisons with Pivot-SD. The paper also describes artifact use, risks from offline self-distillation and verifier-based labeling, and broader impacts.

I Reward Sparsity in wd1++ and the Bootstrap Reward

  • On the 200-prompt AceCode subset, wd1++ stopped learning because zero group reward yielded zero advantage and its likelihood and unlikelihood gradients canceled.The checkpoints at steps 500 and 1000 were identical.Source: I Reward Sparsity in wd1++ and the… · p. 14
  • With bootstrap shaping, wd1++ resumes code learning and gains 3.97 points on MBPP+ at 256 tokens by step 1000.The bootstrap reward gives completions with zero unit-test pass rate up to +0.05 for basic syntactic requirements; Table 13 reports the shaped-run results.Source: I Reward Sparsity in wd1++ and the… · pp. 14–15, Table 13 · p. 15
  • Even with reward shaping, wd1++ does not surpass Pivot-SD: on HumanEval+, Pivot-SD leads by 6.9 points, 40.43 versus 33.54.On MBPP+, the shaped run reaches 47.09 versus Pivot-SD at 47.12 ± 1.20; the wd1++ run uses the same AceCode subset, and Pivot-SD uses no reward shaping.Source: I Reward Sparsity in wd1++ and the… · p. 15, Table 12 · p. 15

J Artifact Licenses and Intended Use

  • The benchmarks and models are used under their stated licenses, and the generated trajectories and evaluation scripts are to be released under MIT.The paper says the datasets contain no personally identifying information or offensive content.Source: J Artifact Licenses and Intended Use · pp. 15–16

K Potential Risks

  • Outcome filtering does not remove biases or hallucinations in the frozen base model’s trajectories, and training on them can amplify these issues..Source: K Potential Risks · p. 16
  • Trajectory labeling relies on exact-match checks for mathematics and unit-test execution for code; open-ended tasks without clear verifiers may invite reward hacking or superficial proxies..Source: K Potential Risks · p. 16

L Broader Impacts

  • The authors report about one-eighth the wall-clock time of a 5,000-step diffu-GRPO run, using 200 prompts and one offline sampling pass, as making post-training research feasible on modest academic hardware..Source: L Broader Impacts · p. 16
  • The authors suggest localized unlikelihood may aid alignment by penalizing selected commitments rather than valid fragments across a full sequence; they also note dual-use risks.The approach relies on an external verifier and a fixed offline dataset.Source: L Broader Impacts · p. 16

Related papers

Loading 2610.03665v1…