Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models

1 INTRODUCTION

• Confidence-based decoding uses local predictability but does not explicitly account for how revealing a token changes the difficulty of later predictions.Committing more positions can reduce rounds, but overly aggressive commitments risk irreversible errors; conservative schedules preserve quality at a computational cost.Source: §1 · pp. 1–2
• A less predictable token can encode a consequential semantic, syntactic, or reasoning decision and reduce uncertainty across several remaining positions.This distinction matters particularly in structured tasks such as mathematical reasoning and code synthesis, where intermediate decisions induce downstream dependencies.Source: §1 · p. 2
• 10.12× and 5.12× speedups over baselines on GSM8K for LLaDA-8B-Instruct and Dream-7B-Instruct, respectively, while maintaining strong or superior task performance.The evaluation spans mathematical reasoning and code generation tasks.Source: §1 · p. 2
• dSpark constructs unmasking schedules that prioritize commitments for their downstream help to remaining predictions, then distills those schedules into the model for parallel inference without inference-time schedule search.A KL-divergence regularizer helps keep the distilled model close to the original distribution and preserve generation quality.Source: §1 · p. 2

Diffusion Language Models

• Diffusion language-model work developed stronger modeling objectives and capable backbones, including LLaDA and Dream, while iterative generation still leaves room to improve inference efficiency.The cited work also explores blockwise generation and converting autoregressive models to diffusion models.Source: §2 · p. 3

Efficient Diffusion Language Decoding

• Efficient decoding methods target cache reuse, block-based generation, fewer-step distillation, earlier token readiness, multi-block decoding, or lookahead-based commitment selection.The section identifies dParallel and d3LLM as closely related approaches.Source: §2 · p. 3

3.1 MASKED DIFFUSION LANGUAGE MODELS

• Masked diffusion language models learn to recover tokens from partially masked sequences, using a forward process that independently replaces each token with MASK with probability t.The noise level ranges from clean input to complete masking, and a bidirectional denoiser is trained to reconstruct masked tokens.Source: §3.1 · p. 3

3.2 PARALLEL DECODING AND UNMASKING SCHEDULES

• At each decoding round, a forward pass predicts tokens at eligible masked response positions, and the decoder selects a nonempty subset to reveal.The sequence begins fully masked, and selection can use confidence thresholds or predictive-entropy budgets.Source: §3.2 · p. 3
• An unmasking schedule specifies which positions are revealed in each round; Figure 1 measures how individual reveals affect other masked positions and shows that larger uncertainty area is associated with more forward passes.The figure reports mean surprisal drop and the net number of other positions made ready.Source: §3.2 · p. 3, Figure 1 · p. 4

4.1 MOTIVATION

• Revealing a less predictable token can reduce uncertainty at other positions more than revealing a confident token, so efficient schedules should account for downstream usefulness.An accumulated uncertainty area captures how early uncertainty falls; optimized schedules finish verified responses in fewer rounds, with the greatest gains on responses needing the most rounds.Source: §4.1 · p. 4

4.2 UNCERTAINTY-GUIDED SCHEDULE CONSTRUCTION

• For each plausible candidate reveal, dSpark measures uncertainty reduction at the other masked positions and promotes the candidate with the largest positive score.If no candidate has a positive score, promotion is skipped; the resulting trajectory changes reveal order while keeping the correct response fixed.Source: §4.2 · pp. 4–5

4.3 SCHEDULE DISTILLATION

• The student trains on promoted-trajectory states using cross-entropy and confidence sharpening, which encourages earlier commitments where its predictions are already correct.The sharpening term strengthens confidence in correct predictions.Source: §4.3 · pp. 5–6
• Because sharpening can also affect unsuccessful paths, dSpark uses failed trajectories for selective KL regularization toward the frozen base model’s distributions.The anchoring limits changes where reliable token-level supervision is unavailable, while correct constructed trajectories remain free of the KL penalty.Source: §4.3 · p. 6

5.1 EXPERIMENTAL SETUP

• Experiments use LLaDA-Instruct and Dream-Instruct, LoRA fine-tuning, and multi-block decoding, evaluated on GSM8K, MATH, HumanEval, and MBPP.Math uses accuracy and code uses pass@1; results are averaged over three evaluation runs and compared with vanilla models, Fast-dLLM, dParallel, d3LLM, and D2F.Source: §5.1 · p. 7

5.2 MAIN RESULTS

• 10.12× the base model’s decoding parallelism on LLaDA GSM8K preserves accuracy; dSpark also has the highest AUP across all four LLaDA tasks.It improves AUP by 33.3% over d3LLM on GSM8K, matches d3LLM pass@1 with higher TPF on HumanEval, and improves both on MBPP.Source: §5.2 · p. 7
• 5.12× the base model’s TPF on Dream GSM8K comes with a 1.5% accuracy improvement.Against d3LLM, TPF rises 30.4% on MBPP and 36.4% on HumanEval, with higher pass@1 on both; dSpark has the highest TPF and AUP on MATH and MBPP.Source: §5.2 · p. 7
• dSpark achieves the highest TPF and AUP across all four benchmarks on both backbones, with consistent gains in the accuracy–parallelism trade-off.The relative contributions of accuracy retention and increased parallelism vary across tasks.Source: §5.2 · p. 7, Table 1 · p. 8
• dSpark’s higher TPF comes with retained task performance: it leads both TPF and accuracy on GSM8K for LLaDA and Dream.Figure 4 compares accuracy against TPF, with curves from TPF = 1 to each method’s operating point and a dashed Pareto frontier; dSpark also avoids pass@1 loss versus d3LLM on LLaDA HumanEval and improves both measures on Dream.Source: §5.2 · p. 7, Figure 4 · p. 9

5.3 ABLATION STUDY

• The full recipe reaches 10.1 TPF on GSM8K and 6.7 TPF on HumanEval, with accuracies of 78.9% and 39.6%, respectively, and the highest AUP on both.Schedule alone raises TPF but lowers accuracy; KL regularization alone improves accuracy but reduces TPF.Source: §5.3 · pp. 8–9, Table 2 · p. 9
• At a 4:1 correct-to-incorrect trajectory ratio, LLaDA GSM8K reaches 10.12 TPF and 78.94% accuracy; a 1:1 ratio lowers TPF to 7.90 for only a small accuracy gain.Using disjoint correct trajectories for KL regularization instead of incorrect trajectories at 4:1 reduces accuracy from 78.94% to 77.10% and TPF by 30.5%.Source: Table 3 · p. 9, §5.3 · p. 9

• dSpark constructs unmasking schedules using the downstream utility of token commitments and distills them into the model for parallel inference.The conclusion contrasts token utility with confidence alone: less predictable tokens can provide more useful context for resolving the remaining sequence.Source: §6 · p. 10
• On GSM8K, dSpark reports 10.12× and 5.12× speedups for LLaDA-8B-Instruct and Dream-7B-Instruct, respectively, while maintaining strong task performance.Source: §6 · p. 10
• The conclusion identifies unmasking schedules as a learnable component that can narrow the gap between theoretical parallelism and practical dLLM efficiency.Source: §6 · p. 10

• The work concerns dLLM inference efficiency and involves neither human subjects nor sensitive personal data.Source: ETHICS STATEMENT · p. 10

• Generative AI was used only for language polishing and clarity; authors developed and verified the technical content, experiments, analysis, and conclusions.Source: AI USE STATEMENT · p. 10

A.1 WALL-CLOCK THROUGHPUT AND KV CACHING

• With the static CUDA-graph cache, dSpark reaches 313 tokens/s at 0-shot, 260 at 5-shot, and 300 at 512 tokens, 1.3–2.8× its uncached throughput.Accuracy remains within 1.3 points of the uncached result; prefix caching raises 5-shot throughput from 92 to 155 tokens/s, while caches barely change short-prompt throughput.Source: §A.1 · p. 14, Table 4 · p. 14
• 249 tokens/s on GSM8K is 16.6× the base LLaDA throughput and 13.2% above d3LLM, with 78.94% accuracy versus 78.82% for LLaDA.The comparison is for batch size 1; the reported throughput corresponds to a wall-clock speedup while preserving backbone task performance.Source: §A.1 · p. 14, Table 5 · p. 14

A.2 GENERATION LENGTH

• At generation lengths of 128, 256, and 512 tokens, dSpark has the highest TPF throughout and the highest accuracy in five of six settings.At L = 512 on MATH, it reaches 35.64% accuracy and 6.79 TPF, versus 30.46% and 6.28 TPF for d3LLM.Source: §A.2 · p. 15, Table 6 · p. 15

A.3 GENERAL KNOWLEDGE

• On general-knowledge tasks, dSpark matches the base model’s 63.0% MMLU accuracy and scores 47.2% on ARC-C, a 0.4-point drop from base.The authors interpret these results as largely preserving the backbone’s general-knowledge performance.Source: §A.3 · p. 15, Table 7 · p. 15

B TRAINING DETAILS

• Both backbones are fine-tuned with LoRA on attention and feed-forward projections while pretrained weights remain frozen; their optimization settings match, but training corpora differ.Adapters are merged into the backbone before evaluation.Source: B TRAINING DETAILS · p. 15

C EVALUATION DETAILS

• Decoding uses a 0.90–0.95 confidence threshold, blocks of 32 tokens, and at most 256 new tokens; a new block enters the window when the active block is 50% decoded.The next block becomes fully active once its predecessor is 50% decoded, and the main results use no KV cache.Source: C EVALUATION DETAILS · p. 15
• Figure 5 plots accuracy by confidence across 10 equal-width bins, with a diagonal for perfect calibration and a hatched confidence–accuracy gap; dSpark stays closer to calibration than the compared baselines.The bottom panel shows the share of predictions in each bin.Source: Figure 5 · p. 16, Table 8 · p. 16
• TPF is generated non-padding tokens divided by forward passes, while AUP integrates accuracy over parallelism using TPF-sorted operating points.The AUP calculation keeps points within five accuracy points of the best method; each method contributes its one-token-per-forward point and reported operating point.Source: C EVALUATION DETAILS · p. 16

D.1 CANDIDATE FILTER

• The filter admits promoted reference tokens among the model’s top 10 predictions with probability at least 0.02, rather than requiring confidence.These bounds exclude tokens the model has effectively ruled out and keep search small because each candidate is scored with its own forward pass; the authors did not tune the bounds end-to-end.Source: §D.1 · pp. 16–17
• Revealing an uncertain token makes 5.1× more other positions ready and lowers their surprisal 4.7× more than revealing a confident token; promoted low-probability tokens leave accuracy unchanged within sampling noise while saving 6.5 forward passes for the least likely tokens.The token-utility comparison motivates avoiding a stricter confidence filter; the finishing comparison uses 240 promoted-token states, 80 per probability range.Source: §D.1 · p. 17, Table 10 · p. 17

D.2 TOKEN PROMOTION STRATEGY.

• With KL regularization, surprisal-reduction promotion reaches 10.1 TPF and 78.9% accuracy, exceeding highest-surprisal promotion at 8.9 TPF and 77.8% accuracy.Without KL regularization, the proposed criterion reaches 11.4 TPF versus 7.6 for the natural schedule and 8.3 for random promotion, but highest-surprisal promotion lowers accuracy to 73.3%.Source: §D.2 · p. 17, Table 11 · p. 18

D.3 KL WEIGHT

• Setting KL weight β to 1 raises GSM8K accuracy from 74.19% to 78.94% and MATH-500 accuracy from 32.81% to 34.40%, while retaining 10.12 and 4.96 TPF, respectively.Raising β to 2 adds only 0.16 GSM8K accuracy while reducing TPF from 10.12 to 8.1; on MATH-500, accuracy is unchanged and TPF falls.Source: §D.3 · p. 18

E QUALITATIVE EXAMPLE

• For the raspberry-count prompt (6 clusters of 20 plus 67 individual raspberries; reference answer 187), all four shown generations are marked correct, with forward-pass counts of 199, 23, 25, and 17.Source: E QUALITATIVE EXAMPLE · p. 19
• For the salary-and-award prompt, three shown generations are marked wrong at 45, 46, and 256 forward passes, while one is marked correct at 28.The prompt’s reference answer is 114,200.Source: E QUALITATIVE EXAMPLE · p. 20