Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models

Eliya Nachmani; Haim Permuter; Omer Luxembourg

arxiv: 2506.19037 · v4 · pith:BOEC332Ynew · submitted 2025-06-23 · 💻 cs.CL · cs.AI· cs.IT· cs.LG· cs.NE· math.IT

Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models

Omer Luxembourg , Haim Permuter , Eliya Nachmani This is my paper

Pith reviewed 2026-05-19 07:32 UTC · model grok-4.3

classification 💻 cs.CL cs.AIcs.ITcs.LGcs.NEmath.IT

keywords masked diffusion language modelsdilated unmaskingparallel decodingnon-autoregressive generationinference speedupentropy boundlanguage model sampling

0 comments

The pith

Dilated Unmasking Scheduler partitions positions into non-adjacent groups to minimize joint entropy gain and enable fast parallel decoding in masked diffusion models.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Masked diffusion language models promise non-autoregressive generation but lose quality when multiple tokens are unmasked together because existing methods ignore interactions between positions. The paper introduces the Dilated Unmasking Scheduler that divides the sequence into non-adjacent dilated groups and at each step selects the group whose unmasking minimizes an upper bound on the increase in joint entropy. This inference-only approach turns the quality-speed tradeoff into a deterministic function of block size and recovers most of the performance that naive parallel unmasking loses. On benchmarks covering math, code, knowledge, and instruction following, it outperforms confidence-based planners while delivering up to 5.8 times wall-clock speedup over token-by-token decoding without any change to the underlying denoiser.

Core claim

By partitioning sequence positions into non-adjacent dilated groups and unmasking them in parallel according to the group that minimizes an upper bound on joint entropy gain at each denoising step, the Dilated Unmasking Scheduler allows masked diffusion language models to generate text faster than sequential decoding while keeping most of the quality that would be lost under standard parallel unmasking strategies.

What carries the argument

The Dilated Unmasking Scheduler, which partitions token positions into non-adjacent dilated groups and chooses at each step the group minimizing an upper bound on joint entropy gain.

If this is right

DUS recovers most of the performance lost under traditional parallel unmasking strategies.
It yields up to 5.8× wall-clock speedup over token-by-token MDLM decoding without modifying the denoiser.
The speedup becomes deterministic and set by the chosen block size B.
DUS outperforms confidence-based planners across math, code, general-knowledge, and instruction-following benchmarks.
When applied as a drop-in post-filter, dilated spacing also improves adaptive samplers.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The dilated-group idea could be tested in other non-autoregressive sampling schemes to reduce cross-position interference.
Variable block sizes chosen dynamically per sequence might tighten the quality-speed curve further.
Because the method requires no model changes, it can be stacked with distillation or quantization for cumulative gains.
The entropy-bound heuristic may lose effectiveness on very long contexts where distant dependencies dominate.

Load-bearing premise

Partitioning into non-adjacent dilated groups and minimizing the upper bound on joint entropy gain produces schedules whose quality loss remains small enough to be offset by the measured speed gains.

What would settle it

On short sequences where exhaustive search over all possible parallel unmasking orders is feasible, an optimal scheduler would produce measurably higher accuracy than DUS at the same number of parallel steps.

read the original abstract

Masked diffusion language models (MDLMs) promise fast, non-autoregressive text generation, yet existing samplers, which pick tokens to unmask based on model confidence, ignore interactions when unmasking multiple positions in parallel and effectively reduce to slow, autoregressive behavior. We propose the Dilated Unmasking Scheduler (DUS), an inference-only, planner-model-free method that partitions sequence positions into non-adjacent dilated groups and unmasks them in parallel so as to minimize an upper bound on joint entropy gain at each denoising step. By explicitly trading off the number of network calls against generation quality, DUS recovers most of the performance lost under traditional parallel unmasking strategies. Across math (GSM8K, MATH500), code (HumanEval, MBPP), general-knowledge (BBH, MMLU-Pro), and instruction following (IFEval) benchmarks, DUS outperforms confidence-based planners and turns the diffusion-specific quality-speed trade-off into a deterministic, predictable speedup set by the block size $B$, yielding up to $5.8\times$ wall-clock speedup over token-by-token MDLM decoding without modifying the underlying denoiser. Applied as a drop-in post-filter, dilated spacing also improves adaptive samplers. Code is available at https://github.com/omerlux/DUS.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

DUS gives a clean inference-only way to dial in speed versus quality for masked diffusion LMs via dilated blocks and an entropy upper bound, but the evidence tying the bound to real gains is still light.

read the letter

The key point is that this paper introduces a Dilated Unmasking Scheduler that splits positions into non-adjacent dilated groups and orders the parallel unmasking steps by minimizing an explicit upper bound on joint entropy gain. The result is a deterministic knob (block size B) that trades network calls for quality, and they report recovering most of the performance lost to naive parallel unmasking while hitting up to 5.8x wall-clock speedup over token-by-token decoding on GSM8K, MATH, HumanEval, and the other listed benchmarks. It also works as a drop-in improvement for existing adaptive samplers.

Referee Report

3 major / 2 minor

Summary. The paper introduces the Dilated Unmasking Scheduler (DUS) for masked diffusion language models (MDLMs). It partitions token positions into non-adjacent dilated groups and selects the parallel-unmasking order at each denoising step by minimizing an explicit upper bound on joint entropy gain. The method is inference-only and planner-model-free. Empirical results on GSM8K, MATH500, HumanEval, MBPP, BBH, MMLU-Pro and IFEval show that DUS recovers most of the quality lost by standard parallel unmasking, yields up to 5.8× wall-clock speedup over token-by-token MDLM decoding, and can be used as a drop-in post-filter for adaptive samplers. Speedup is presented as a deterministic function of block size B.

Significance. If the central empirical claim holds, the work supplies a practical, training-free way to control the quality-speed trade-off in MDLMs and turns an otherwise heuristic scheduling problem into a predictable function of a single hyper-parameter B. The public code release and the fact that the scheduler is derived from an explicit (if upper) bound are positive features.

major comments (3)

[§3] §3 (entropy-bound derivation): the manuscript derives an upper bound on joint entropy gain but provides no analysis of the tightness of this bound relative to the true joint entropy under the model's attention. Without a gap analysis or a proof that the selected order remains near-optimal when higher-order dependencies are present, it is unclear whether the minimization step actually drives the reported recovery of performance or whether any fixed dilated ordering would suffice.
[Experiments] Experimental section (benchmarks and ablations): the abstract and results claim consistent outperformance and recovery of most lost quality, yet no error bars, no multiple-run statistics, and no ablation that isolates the entropy-minimization step from the mere use of dilated groups are reported. This makes it difficult to assess whether the central construction is load-bearing for the observed gains on GSM8K/MATH/HumanEval etc.
[§4] §4 (speedup claim): the reported wall-clock speedup is stated to be a direct function of block size B. However, the manuscript does not quantify how much of the speedup is attributable to the specific ordering chosen by the bound versus the dilation pattern alone, leaving the causal link between the proposed scheduler and the 5.8× figure incompletely supported.

minor comments (2)

[§3] Notation for the dilated groups and the precise definition of the upper-bound objective could be stated more formally (e.g., as an explicit optimization problem) to aid reproducibility.
[Figures] Figure captions and axis labels in the speed-quality trade-off plots should explicitly indicate the block size B used for each curve.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive and detailed feedback. We address each major comment point by point below, indicating where revisions will be made to the manuscript.

read point-by-point responses

Referee: [§3] §3 (entropy-bound derivation): the manuscript derives an upper bound on joint entropy gain but provides no analysis of the tightness of this bound relative to the true joint entropy under the model's attention. Without a gap analysis or a proof that the selected order remains near-optimal when higher-order dependencies are present, it is unclear whether the minimization step actually drives the reported recovery of performance or whether any fixed dilated ordering would suffice.

Authors: We thank the referee for this observation. The upper bound is derived via subadditivity of entropy combined with the reduced cross-attention between non-adjacent dilated positions in the transformer. While we acknowledge the absence of a formal gap analysis or optimality proof under higher-order dependencies, the bound remains useful for guiding scheduling because it is computable in closed form without additional model calls. Empirically, DUS outperforms both random and fixed dilated orderings on the reported benchmarks, indicating that the minimization step contributes to quality recovery. In the revision we will add a short discussion of the bound's derivation assumptions and an ablation comparing entropy-minimizing order versus fixed dilated order within the same groups. revision: partial
Referee: Experimental section (benchmarks and ablations): the abstract and results claim consistent outperformance and recovery of most lost quality, yet no error bars, no multiple-run statistics, and no ablation that isolates the entropy-minimization step from the mere use of dilated groups are reported. This makes it difficult to assess whether the central construction is load-bearing for the observed gains on GSM8K/MATH/HumanEval etc.

Authors: We agree that the current experimental presentation would be strengthened by statistical reporting and targeted ablations. We will rerun all main experiments across multiple random seeds, reporting means and standard deviations. We will also add an explicit ablation that fixes the dilated grouping but replaces the entropy-minimization scheduler with either a random or sequential ordering inside each group, thereby isolating the contribution of the bound minimization. These additions will appear in the revised experimental section and supplementary material. revision: yes
Referee: [§4] §4 (speedup claim): the reported wall-clock speedup is stated to be a direct function of block size B. However, the manuscript does not quantify how much of the speedup is attributable to the specific ordering chosen by the bound versus the dilation pattern alone, leaving the causal link between the proposed scheduler and the 5.8× figure incompletely supported.

Authors: The wall-clock speedup is governed by the number of parallel unmasking steps, which is strictly determined by the block size B and the dilation pattern; the entropy-minimization ordering affects only which tokens are chosen within each parallel step and therefore has negligible impact on measured runtime. To make this distinction explicit, we will add wall-clock timing results for dilated groups using random intra-group ordering and compare them directly to DUS timings. This will clarify that the reported speedups (including the 5.8× figure) stem from the dilation schedule while the bound minimization is responsible for quality recovery. We will revise the relevant paragraph in §4 and the caption of the speedup figure accordingly. revision: partial

Circularity Check

0 steps flagged

No circularity: derivation from explicit upper bound on joint entropy is independent of evaluation data

full rationale

The paper defines DUS by partitioning into non-adjacent dilated groups and selecting order via minimization of an explicit upper bound on joint entropy gain (described in abstract and §3). This construction uses only the model's attention structure and the bound itself; no fitted parameters from target benchmarks enter the scheduler definition. Speedup is stated as a deterministic function of block size B, and quality recovery is measured on held-out sets (GSM8K, MATH500, HumanEval, etc.) after the schedule is fixed. No equations reduce the performance claim to a quantity defined by the same data, and no self-citation chain is invoked to justify the bound or the dilation choice. The derivation therefore remains self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

1 free parameters · 0 axioms · 0 invented entities

The method introduces one tunable integer (block size B) that directly sets the parallelism level and therefore the speedup. No new physical or mathematical entities are postulated. The entropy upper bound is treated as a standard information-theoretic quantity rather than an ad-hoc invention.

free parameters (1)

block size B
Integer that determines the number of parallel groups and therefore the exact speedup factor; chosen by the user to trade quality against speed.

pith-pipeline@v0.9.0 · 5782 in / 1418 out tokens · 22059 ms · 2026-05-19T07:32:28.628377+00:00 · methodology

discussion (0)

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

IndisputableMonolith/Foundation/ArithmeticFromLogic.lean embed_strictMono_of_one_lt unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

DUS partitions sequence positions into non-adjacent dilated groups and unmasks them in parallel so as to minimize an upper bound on joint entropy gain at each denoising step... R = ⌈loga B⌉, st = ⌊B/a^t⌋... Pt = {k ∉ Ut−1 | (k−1) mod st = 0}
IndisputableMonolith/Foundation/AlexanderDuality.lean alexander_duality_circle_linking unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

Under a fast-mixing first-order Markov chain... I(Xim;Xin) ≤ −1/2 log(1−ρ^{2d}) = O(ρ^{2d})... dilated pattern maximizes distances

What do these tags mean?

matches: The paper's claim is directly supported by a theorem in the formal canon.
supports: The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends: The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses: The paper appears to rely on the theorem as machinery.
contradicts: The paper's claim conflicts with a theorem or certificate in the canon.
unclear: Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

Drifting Objectives for Refining Discrete Diffusion Language Models
cs.CL 2026-05 unverdicted novelty 7.0

TokenDrift refines discrete diffusion language models by applying anti-symmetric drifting to soft-token features during training, yielding large reductions in generation perplexity at low NFEs.
Saber: An Efficient Sampling with Adaptive Acceleration and Backtracking Enhanced Remasking for Diffusion Language Model
cs.AI 2025-10 unverdicted novelty 6.0

Saber improves both speed and accuracy of diffusion language models on code generation by dynamically adjusting unmasking steps and reverting low-confidence tokens via backtracking.