Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models

1. Introduction

• DUS partitions each length-B block into dilated groups, achieving O(log B) denoising iterations instead of O(B) under a Markov assumption.The schedule uses a fixed dilation pattern and requires no retraining or additional planner modules.Source: §1 · pp. 1–2
• Figure 1 compares DUS and self-confidence speed-accuracy curves and shows a B=32 GSM8K trace where DUS remains coherent while self-confidence truncates.Source: Figure 1 · p. 2
• As a model-agnostic inference-time drop-in, DUS requires no architecture or training changes.Source: §1 · p. 2
• Under fast-mixing assumptions, DUS approaches the joint-entropy bound at each iteration relative to baselines with the same step budget.Source: §1 · p. 2
• Across math, code, general-knowledge, and instruction-following benchmarks, DUS reduces denoiser calls by up to an order of magnitude and consistently improves quality over confidence-based planners.Experiments use LLaDA-8B, Dream-7B, and DiffuCoder-7B.Source: §1 · p. 2
• Dilated spacing also improves EB/CB adaptive samplers at matched or modestly higher NFE without changing their scoring or selection criteria.Source: §1 · p. 2

2. Related Work

• Self-planners rank masked positions by confidence, entropy, or probability margin, while EB/CB heuristics do not explicitly model dependencies among jointly unmasked tokens.Source: §2 · pp. 2–3
• Path Planning uses a pretrained BERT to score candidate token sets, finding better scores in various experiments at the cost of auxiliary model calls.Source: §2 · p. 2
• Semi-AR diffusion reveals contiguous blocks sequentially but still incurs O(B) denoiser calls per block; this work focuses exclusively on inference.Source: §2 · pp. 2–3
• At 7–8B parameters, LLaDA and Dream demonstrate practical parallel decoding, while DiffuCoder reduces left-to-right bias and supports fully parallel decoding without semi-AR blocks.Source: §2 · p. 3
• KV-cache methods reuse stable transformer tensors across denoising steps and report up to 3× speedups with minimal quality loss, at the cost of extra memory.FreeCache also uses an auxiliary AR guide model.Source: §2 · p. 3
• Jump Your Steps selects non-uniform noise levels to reduce denoiser calls, but its score-matching target models are too small for meaningful benchmark comparison here.Many discrete masked diffusion models also use cross-entropy objectives, making direct log-likelihood scheduling less straightforward.Source: §2 · p. 3

3.5. DUS as Predefined Planner

• DUS deterministically reveals preselected, non-adjacent positions in parallel, reducing denoiser evaluations per block from O(B) to O(log B).For block size B and base a > 1, it completes a block in R = ⌈log_a B⌉ evaluations, versus O(B) for token-by-token unmasking.Source: §3 · p. 3, §3.5 · p. 5

3.1. Notation

• The planner selects which masked indices to reveal, while revealed values are sampled from the fixed denoiser’s token-wise distributions; at inference, the planner is the remaining degree of freedom.Theoretical analysis assumes a stationary, ergodic, fast-mixing variable-length Markov chain.Source: §3.1 · p. 3

3.2. MDLM Inference in Practice

• Because inference-time masking is an edge case of training, denoisers may be poorly calibrated there, motivating multiple unmasking-denoising iterations within semi-autoregressive blocks.Typical large-scale diffusion LLM inference partitions sequences into blocks of length B and completes each block over several iterations.Source: §3.3 · p. 4

3.3. Detailed Formulation

• In block-wise sampling, the planner’s ideal objective is joint conditional entropy, but parallel sampling uses a tractable sum-of-marginals surrogate whose entropy gap should be minimized.Token-by-token unmasking achieves the tight bound, whereas full parallel sampling corresponds to revealing the whole block as one group.Source: §3.3 · pp. 4–5

3.4. Self-Planners Guided by Denoiser Confidence

• Denoiser-guided self-planners rank positions using per-token confidence or entropy scores, but do not explicitly control dependence among positions selected together.Selecting strongly coupled, often nearby positions can therefore yield a large entropy gap.Source: §3.4 · p. 5

3.5. DUS as Predefined Planner

• DUS spaces early reveals widely and fills remaining gaps at finer resolution, with an optional uncertainty-based skip heuristic that is not used in the default schedule.For B = 8 and a = 2, the schedule uses three iterations: {1, 5}, then {3, 7}, then {2, 4, 6, 8}.Source: §3.5 · p. 5

3.6. Theoretical Analysis of DUS

• Under fast-mixing VLMC assumptions, dilated spacing makes selected tokens less dependent and the joint-entropy gap small in coarse iterations, while later context may help correct local inconsistencies.The analysis does not claim conditional independence in later, finer iterations; strong long-range dependencies can still produce a larger gap.Source: §3.6 · p. 6

3.7. Experimental Setup

• The experimental setup evaluates five MDLM variants with block sizes B ∈ {8, 16, 32, 64}, comparing DUS with self-confidence and adaptive sampler baselines across benchmark families.The listed benchmark families include math and code (Table 1) and general knowledge (Table 2); the setup describes EB- and CB-Sampler as adaptive baselines.Source: §3.7 · pp. 6–7, Table 1 · p. 7, Table 2 · p. 7
• For block length B, DUS unmasks an average of log2 B tokens per denoiser iteration across a block.The setup reports relative NFE speedup as B/log2 B, the ratio of token-by-token total NFE to the experiment’s total NFE; planners are compared at matched B and NFE budgets.Source: §3.7 · p. 7
• DUS spacing is also applied as a post-filter to EB or CB selections for adaptive samplers.Source: §3.7 · p. 7
• A GSM8K problem with B = 32 is visualized for DUS and self-confidence in Figures 1b and 1c, respectively.Source: §3.7 · p. 7

4.1. Math and Coding Experiments

• At the same NFE budget, DUS improves performance by up to 27% over self-confidence; larger blocks enable up to 10× fewer iterations with some loss in final accuracy.Smaller blocks yield higher accuracy but require more denoiser rounds.Source: §4.1 · p. 8
• Across all four benchmarks, DUS matches or improves the self-confidence baseline, reaching up to 40% gains on math and up to 20% on code while using less NFE than token-by-token inference.Source: §4.1 · p. 8
• On code benchmarks, LLaDA and DiffuCoder base models can outperform their instruct counterparts under DUS, unlike with self-confidence, where instruct tuning is generally superior.Source: §4.1 · p. 8

4.2. General Knowledge Experiments

• At matched speedups, DUS consistently beats self-confidence on BBH by 1–5% and on MMLU-Pro by 5–9%.The gains are smaller than those reported for math and code.Source: §4.2 · p. 8

4.3. Planners Parallelism Performance

• DUS performs best overall; with random planners, dilated-incremental scheduling recovers much of fixed-k’s accuracy loss and can exceed fixed-k self-confidence.For self-confidence, fixed-k slightly outperforms dilated-incremental, indicating that its selection benefits from uniform parallelism.Source: §4.3 · p. 8
• Across three models, DUS reaches near-token-by-token accuracy with 3–4× fewer NFE, while EB and CB operate at higher NFE.The comparison is for IFEval scores versus NFE with early termination.Source: §4.3 · p. 8, Figure 2 · p. 8

4.4. Beyond DUS: Dilated Spacing as a Drop-in Filter for Adaptive Schedulers

• Dilated spacing improves aggressive EB and CB post-filter accuracy in all reported cases, with HumanEval gains up to +13.4% for EB and +12.2% for CB.The broader sweep found at least one initial gap that improved both accuracy and NFE for each model–dataset pair and sampler.Source: §4.4 · p. 9, Table 4 · p. 9

• Across diffusion LLM experiments, DUS improves downstream accuracy by up to 27% on challenging math and code benchmarks while reducing denoising iterations by an order of magnitude.Source: §5 · p. 9
• As a post-filter on EB- and CB-Sampler selections, dilated spacing improves accuracy at matched or modestly higher NFE.Source: §5 · p. 9
• DUS requires no changes to model architecture or training, and its schedule is orthogonal to KV caching, which reduces per-step cost.Source: §5 · p. 9

Impact Statement

• DUS reduces the calls needed for generation while maintaining or improving output quality, with potential to lower computational costs and energy consumption.Source: Impact Statement · p. 10

A.1. Use of Large Language Models

• The authors report using an LLM only for minor draft editing, not for ideation, experiments, analysis, data generation, or technical content.Source: §A.1 · p. 11

A.2. Masked Diffusion: Forward Process and Training Objective

• During training, the forward process independently masks each token with probability t; the sequence is fully observed at t=0 and fully masked at t=1.Source: §A.2 · p. 11
• The denoiser is trained to reconstruct original tokens with a reweighted cross-entropy loss over masked positions, which upper-bounds negative log-likelihood.The 1/t reweighting compensates for fewer masked tokens at low noise levels.Source: §A.2 · p. 12
• At inference, a planner selects masked positions, the fixed denoiser predicts their token distributions, and the sampler updates the state.Source: §A.2 · p. 12

A.3. Datasets and Settings

• Evaluations use NVIDIA Tesla V100 GPUs, report average NFEs for fair comparison, and use the Language Model Evaluation Harness.The code is publicly available for reproducibility.Source: §A.3 · p. 12
• The benchmark settings include full GSM8K, MBPP, and HumanEval datasets, sampled BBH and MMLU-Pro subsets, and generation lengths of 256 except 512 for coding tasks.IFEval uses generation length 1024 for planner-parallelism experiments and 256 for hybrid post-filter evaluation.Source: §A.3 · p. 12
• Evaluations use few-shot chain-of-thought prompting, with shot counts varying by benchmark; generation stops when EOS and all preceding tokens are unmasked.The usual unmasking parameter is k = log2 B unless stated otherwise.Source: §A.3 · p. 12

B.1. Block-Size Effect on Speedups

• Under a fixed 8× speedup budget, GSM8K accuracy rose from 13.33% at B=8 to 32.33% at B=32 for LLaDA-Base, and from 12.00% to 56.67% for LLaDA-Instruct.The study used 300 samples and total NFE=32; the tested block sizes B=8, 16, and 32 correspond to average NFEs per block of 1, 2, and 4.Source: §B.1 · p. 12, Table 6 · p. 13
• The authors attribute larger-block gains to wider initial token spacing, which reduces mutual information among unmasked tokens and helps later iterations correct errors; excessively large blocks may leave too little nearby context.The account says block size can improve quality without raising the compute budget, while noting the context constraint.Source: §B.1 · pp. 12–13

B.2. Skip Threshold Ablation

• On HumanEval with LLaDA-Base, c=0.3 raises pass@1 from 9.8% to 12.8% and lowers NFE from 43 to 31.Skip improves the best score in three of four cells, but the optimal threshold varies by task; c=0 remains a competitive baseline.Source: Table 7 · p. 13, §B.2 · p. 13

B.3. Single-Block DUS (B=G)

• At ×28 speedup, single-block DUS reaches 42.3% on GSM8K and 17.2% on MATH500 for LLaDA-Instruct, trading accuracy for further speedup.The skip threshold improves the best-row accuracy in three of six model-dataset settings; c=0 is best for LLaDA-Instruct on GSM8K and MATH500.Source: Table 8 · p. 13, §B.3 · p. 13

B.4. Parameter Sensitivity

• The default settings (a=2, base skip=1) win at every dataset and block-size pair; raising a costs up to 77% relative accuracy on MATH500 and 56% on GSM8K.Raising base skip from 1 to 2 saves approximately 25% NFE but costs 32%–69% relative accuracy across datasets.Source: Table 9 · p. 13, §B.4 · pp. 13–14

B.5. Speed-Quality Trade-off Analysis

• Across all datasets and speedup factors, DUS consistently scores higher than the self-confidence planner.The comparison covers GSM8K, HumanEval, and MBPP at speedups from 2.7× to 10.7×.Source: §B.5 · p. 14
• Larger blocks enable up to 10× fewer iterations but lower end accuracy, while smaller blocks yield higher accuracy at the cost of more denoiser rounds.Source: §B.5 · p. 14
• GSM8K scores decline steadily as inference speed increases, whereas code-benchmark trends are less smooth at higher speedups.The passage suggests this may reflect the code benchmarks’ already low performance at higher speedups.Source: §B.5 · p. 14

B.6. Runtime Breakdown: Sampler Overhead vs. Denoiser

• The model forward pass takes 247.34 ms per step, dominating sampler component costs.All components remain small compared with the forward pass.Source: §B.6 · p. 14
• Entropy computation takes approximately 1.26–2.01 ms per step and grows with block size, versus approximately 1.16–1.26 ms for confidence computation.Source: §B.6 · p. 14
• CB selection overhead increases with block size from approximately 0.15 to 0.72 ms, while EB selection remains relatively small.DUS uses only a lightweight schedule lookup.Source: §B.6 · p. 14

B.7. Token Distance Analysis Across Planners

• DUS achieves 2–3× wider average spacing than every alternative and isolates essentially all co-decoded tokens; self-confidence planners cluster most aggressively.For self-confidence, 49%–58% of co-decoded tokens are adjacent to another co-decoded token.Source: §B.7 · p. 14, Table 11 · p. 15

B.8. Wall-Clock Throughput

• For LLaDA-8B-Base on GSM8K, DUS reaches 59.51% accuracy over token-by-token decoding at B=16; at B=32, it is 5.8× faster with 49.36% accuracy.The wall-clock comparison uses a single NVIDIA RTX 6000 Ada GPU, generation length 256, and 105 samples after 5 warmup samples.Source: §B.8 · p. 15, Table 12 · p. 16

B.10. Dilated Spacing Post-Filter Details

• The dilated-spacing post-filter generally requires equal or slightly more NFE than the base adaptive sampler, usually with higher accuracy, especially for aggressive base settings.Performance is not strictly monotone in the initial gap, with occasional regressions on Dream-Instruct CB.Source: §B.10 · p. 16

C.1. Surrogate-to-Joint Entropy Gap Decomposition

• Lemma 3.3 considers the total loss from parallel sampling within disjoint groups of a block under the optimal-estimator assumption.Source: §C.1 · p. 16

B.9. Adaptive Scheduler Comparison

• DUS occupies the lowest-NFE end of the adaptive-sampler comparison, providing the most aggressive speedup at some accuracy cost; EB and CB reach higher accuracy at higher NFE.The ordering between EB-Sampler and CB-Sampler varies by task.Source: §B.9 · p. 15, Figure 5 · p. 17

C.2. Lemma 3.5 - Restate and proof

• Under a stationary, ergodic, finite-order fast-mixing VLMC, sufficiently spaced token indices make joint conditional entropy ε-close to summed marginal entropies.The lemma states that a distance threshold D_ε suffices for the approximation.Source: §C.2 · pp. 17–18

C.3. MI Decay Under Fast Mixing

• For fast-mixing finite-alphabet VLMCs, mutual information between a token and any finite, sufficiently distant future block decays geometrically with spacing.The assumption specifies an irreducible, aperiodic context chain; the lemma gives constants C and ρ for the decay bound.Source: §C.3 · pp. 18–19

D. Qualitative Generation Trajectories

• Trajectory examples compare DUS and fixed-block scheduling across four configurations; token colors mark first-unmasking step, and prompt lengths are reported.The configurations use two block sizes for each strategy, with fixed scheduling using self-confidence.Source: D. Qualitative Generation Trajectories · p. 19