PaperlayerView exact revision on arXiv

Source-linked AI summary

LED: Improving Masked Diffusion Language Models by Lookahead Editing

Yuyan Zhou, Zeyuan Yang, Yuxuan Song, Zhicheng Cai, Hongli Yu, Weiying Ma, Hao Zhou, James Kwok

arXiv:2610.23122v1cs.LGcs.CL

TL;DR

The paper asks whether masked diffusion models can gain iterative revision without losing their simpler transition schedule. It proposes Lookahead Editing, which refines future predictions in a rewritable proposal before committing tokens, and reports average gains over standard MDM SFT across eight benchmarks.

  • Problem

    Masked diffusion models are easier to optimize but have limited iterative revision and computational expressivity, while uniform diffusion models support revision and greater expressivity at higher training difficulty.

    Source: §1 · p. 2, §3 · p. 3

  • Method

    Lookahead Editing adds a rewritable proposal state that iteratively refines future predictions before tokens are committed through masked diffusion.

    Source: §1 · p. 2, §6 · p. 10

  • Results

    LED SFT gains +11.4 and +7.7 points over standard MDM SFT across eight benchmarks for Qwen2.5-7B and Qwen3-8B, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    LED combines iterative revision with the simple mask-based transition schedule of MDMs.

    Source: §6 · p. 10

    The analysis of learning-rate and excess-loss differences assumes unbiased gradient noise and a model expressive enough to represent the data distribution.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

Diffusion large language models (dLLMs) have emerged as a promising alternative to autoregressive language models, enabling non-autoregressive text generation through parallel decoding. Existing dLLMs can be broadly categorized into two families: masked diffusion models (MDMs) and uniform diffusion models (UDMs) In this work, we theoretically show that MDMs are easier to optimize due to their mask-based transition schedule but have lower computational expressivity, whereas UDMs enable iterative state revision and greater expressivity but are substantially harder to train. This trade-off raises a natural question: can we obtain the expressivity of iterative revision without giving up the simple, easy-to-optimize transition schedule of MDMs? Motivated by this analysis, we propose Lookahead Editing (\ours{}), which augments masked diffusion decoding with a rewritable proposal state. At each decoding step, the model performs several lookahead updates to iteratively refine future predictions in this proposal state before committing tokens to the masked sequence. This design introduces a rewritable proposal that uses additional compute to improve expressivity while keeping the simple mask-based transition schedule of MDMs. We fine-tune Qwen2.5-7B and Qwen3-8B to obtain LED-\texttt{7B} and LED-\texttt{8B}, respectively. Across eight benchmarks, LED SFT yields average gains of $+11.4$ and $+7.7$ points over standard MDM SFT. Moreover, using SFT alone, LED-\texttt{8B} outperforms LLaDA2.1-16B Q-mode, which is trained with continued pre-training, supervised fine-tuning, and reinforcement learning.

1 Introduction

Diffusion language models offer parallel denoising as an alternative to autoregressive decoding. The paper analyzes a training–expressivity trade-off and proposes Lookahead Editing to combine iterative revision with masked diffusion, reporting gains over standard MDM fine-tuning.

1 Introduction

  • Autoregressive models generate tokens sequentially in a fixed left-to-right order, limiting their use of hardware parallelism; diffusion models instead enable parallel denoising over discrete states.The passage describes autoregressive decoding as largely memory-bound.Source: §1 · p. 1
  • MDMs have simpler optimization because their denoiser-independent transition schedule eliminates the schedule-matching loss, whereas UDMs incur an additional model-dependent schedule-matching loss.The analysis connects this difference to a smaller excess loss for MDMs.Source: §1 · p. 2
  • MDMs remain within TC0, while UDMs’ rewritable state supports NC1 computations such as evaluating nested Boolean expressions.The paper characterizes this as a computational-expressivity limitation of MDMs despite arbitrary-order prediction and bidirectional attention.Source: §1 · p. 2
  • Lookahead Editing (LED) adds a rewritable proposal state for iterative future-token refinement while retaining MDM’s mask-based transition schedule.At inference, proposal updates alternate with commit stages; committed tokens remain fixed, while reliable proposal tokens are preserved and erroneous ones revised.Source: §1 · p. 2
  • +11.4 and +7.7 points are LED SFT’s average gains over standard MDM SFT across eight benchmarks for Qwen2.5-7B and Qwen3-8B, respectively.Using SFT alone, LED-8B also outperforms LLaDA2.1-16B Q-Mode, which uses continued pre-training, supervised fine-tuning, and reinforcement learning.Source: §1 · p. 2

2 Preliminary

The preliminaries define discrete diffusion as corruption and denoising over vocabulary states, then introduce circuit complexity as a framework for characterizing parallel computation.

2.1 Diffusion Language Models

  • A diffusion language model defines a Markov trajectory of vocabulary states from clean X0 to terminal noisy X1, with forward corruption and learned reverse denoising.MDMs use an absorbing transition toward a mask token, while UDMs transition toward a uniform vocabulary distribution.Source: §2.1 · pp. 2–3
  • The discussion focuses on block diffusion language models, which partition sequences into blocks for discrete diffusion.This scope is stated to apply unless otherwise noted.Source: §2.1 · p. 3

2.2 Circuit Complexity

  • Circuit complexity characterizes parallel computation by circuit size, depth, fan-in, and gate types; NCk, ACk, and TCk form nested classes with increasingly permissive gates.The classes satisfy NCk ⊆ ACk ⊆ TCk ⊆ NCk+1.Source: §2.2 · p. 3

3 Optimization-Expressivity Trade-off

The section frames a trade-off: masking simplifies optimization but restricts revision, while rewritable diffusion supports greater computational expressivity.

3.1 Masking Eliminates Schedule Learning.

  • For MDMs, qθ(S)=p(S) for all θ, so the schedule-matching loss Lwhen is zero; UDM schedule matching depends on its denoiser.Source: §3.1 · p. 4
  • Under unbiased gradient noise and a sufficiently expressive model, UDMs require a smaller learning rate and incur larger expected loss than MDMs under the same where loss.Source: §3.1 · p. 4

3.2 Rewriting improves Expressivity.

  • Even with bidirectional attention and O(log n) chain-of-thought, constant-block MDMs remain in TC0, no stronger than autoregressive decoding.Source: §3.2 · p. 4
  • With polynomially many revisions over constant workspace, UDMs can achieve NC1.Source: §3.2 · p. 5

4 Proposed Method: Lookahead Editing

LED adds a rewritable proposal state to masked diffusion, refining future predictions before commitment while retaining mask-based decoding. Training uses pretrained MDM proposals to teach selective editing.

4 Proposed Method: Lookahead Editing

  • LED introduces a rewritable proposal state to masked diffusion, targeting greater expressivity while preserving MDMs’ simple transition schedule.Source: §4 · p. 5

4.1 Training

  • Training adds a pretrained-MDM proposal branch alongside clean and masked branches; cross-entropy supervision teaches LED to retain correct proposals and revise errors.Source: §4.1 · p. 5, Figure 2 · p. 6

4.2 Inference

  • At each decoding step, LED makes K lookahead updates to the proposal, then uses the refined proposal to commit high-confidence tokens until the block is decoded.Source: Figure 3 · p. 6, §4.2 · pp. 6–7

5 Experiments

LED improves over standard MDM fine-tuning across eight benchmarks and two model scales, while also comparing favorably with more extensively post-trained dLLMs. Lookahead steps provide speed–accuracy operating points, and performance depends on the proposal threshold and lookahead depth.

5.1 Experiment Setup

  • The evaluation covers eight zero-shot benchmarks in math, science, and code, using curated open-source post-training datasets and 7B–16B diffusion-model baselines.The benchmark suite includes MATH, AIME 2024/2025, GPQA, HumanEval, HumanEval+, MBPP, and MBPP+; reported metrics are accuracy or Pass@1.Source: §5.1 · p. 7

5.2 Comparison with Existing dLLMs

  • LED-8B scores 68.5 on average, exceeding SFT-only Fast-dLLM by 26.8 points and XDLM by 44.4 points.It also exceeds CPT-trained Nemotron-8B by 9.0 points and CPT+SFT+RL-trained LLaDA2.1-16B by 0.6 points on average.Source: §5.2 · p. 8
  • +11.4 points for the 7B model and +7.7 points for the 8B model are LED’s average gains over the corresponding standard MDM-SFT baselines.LED raises the average score from 45.4 to 56.8 for 7B and from 60.8 to 68.5 for 8B; gains span math, science, and code.Source: §5.2 · p. 8

5.3 Speed-Accuracy Trade-off

  • At 36.7% AIME 2025 accuracy, LED-8B reaches 31 tokens/s versus LLaDA2.1-16B’s 20 tokens/s; at 33.3%, it reaches 44 versus 25 tokens/s.LED-8B also matches Nemotron-8B’s 33.3% accuracy at a higher throughput of 44 versus 40 tokens/s.Source: §5.3 · pp. 8–9, Figure 4 · p. 8
  • On HumanEval+, LED-8B reaches 81.7% at 43 tokens/s, exceeding LLaDA2.1-16B’s 81.1% at 32 tokens/s and Nemotron-8B’s 79.3% at 40 tokens/s.Varying LED’s lookahead steps provides different operating points; the figure reports a stronger Pareto frontier across both benchmarks.Source: Figure 4 · p. 8, §5.3 · p. 9

5.4 Hyper-parameter Analysis

  • With K = 2, raising τ from 0.6 to 0.7 increases AIME 2025 accuracy from 33.3% to 43.3% and HumanEval+ from 79.9% to 84.1%.Among the configurations shown, τ = 0.7 with K = 2 performs best on both benchmarks.Source: §5.4 · p. 9, Figure 5 · p. 9
  • On AIME 2025, accuracy remains 33.3% at K = 0 and K = 1, rises to 43.3% at K = 2, and can drop at K = 3 or K = 4.With τ = 0.9, K = 4 still reaches 83.5% on HumanEval+; longer lookaheads may accumulate inaccurate proposals, while filtering can preserve their usefulness.Source: §5.4 · pp. 9–10

6 Conclusion

The conclusion presents LED as a way to combine masked diffusion’s simple transition structure with iterative refinement. It reports gains over standard MDM fine-tuning and favorable speed–accuracy comparisons with existing dLLM baselines.

6 Conclusion

  • LED adds a rewritable proposal state to masked diffusion, refining future predictions before committing tokens while preserving its simple transition structure.The conclusion reports consistent improvements over standard MDM fine-tuning across eight reasoning and code-generation benchmarks, alongside a stronger speed–accuracy Pareto front than existing dLLM baselines.Source: §6 · p. 10

A Proof of theorem 3.1

The proof compares transition schedules under the forward and reverse processes for masked diffusion. It identifies the transition time of a clean token as having the same distribution in both processes.

A Proof of theorem 3.1

  • The proof defines S(X) as the transition time at which a clean token reaches the absorbing mask state.It uses the event S(X) ≤ t exactly when X_t = m to express the schedule’s cumulative distribution.Source: A Proof of theorem 3.1 · p. 13
  • Under the forward process, the probability that the transition time is at most t is 1 − α_t.This follows from the probability that the token is masked by time t.Source: A Proof of theorem 3.1 · p. 13
  • The reverse-process transition-time distribution has the same density as the forward schedule, so their induced schedule distributions coincide.The proof equates the densities and concludes p(S) = q_θ(S).Source: A Proof of theorem 3.1 · p. 13

B Proof of theorem 3.2

The proof develops the UDM transition schedule and its reverse-process counterpart, then combines schedule terms into a schedule-matching loss. The reverse model’s rate is obtained by marginalizing over its prediction of the clean state.

B Proof of theorem 3.2

  • The proof defines a UDM trajectory’s transition schedule as the ordered times at which its state changes.It derives the forward transition rate from the uniform-diffusion kernel and uses it to obtain the schedule density.Source: B Proof of theorem 3.2 · p. 14
  • Because the reverse model does not observe X_0, its reverse rate marginalizes over the denoiser’s clean-state prediction.The proof defines the reverse generative process in increasing time and derives the marginal schedule intensity from the current state.Source: B Proof of theorem 3.2 · pp. 15–16
  • Combining the forward and reverse schedule terms yields the schedule-matching loss.The derivation converts a sum over observed transition times into an integral under the forward schedule distribution.Source: B Proof of theorem 3.2 · p. 17

C Proof of theorem 3.3

The proof establishes that the UDM Hessian dominates the MDM Hessian and that local gradient-noise covariance is proportional to the Hessian. It then derives a step-size condition for convergence and the limiting second moment.

C Proof of theorem 3.3

  • The UDM Hessian dominates the MDM Hessian because their difference is positive semidefinite.The proof writes HU = HM + Hwhen and establishes Hwhen ⪰ 0.Source: C Proof of theorem 3.3 · p. 18
  • For negative-log-likelihood losses, the local gradient-noise covariance is Σj = 1/B Hj for j ∈ {M, U}.The proof relates the two models’ Hessians through HU = HM + Hwhen.Source: C Proof of theorem 3.3 · p. 19
  • The condition 0 < η < 2 λmax(HU) implies 0 < η < 2 λmax(Hj) for both j ∈ {M, U}.For every positive eigenvalue, this gives |1 − ηλj,i| < 1, so ρ(Aj) < 1 and Ak_j → 0 as k → ∞.Source: C Proof of theorem 3.3 · pp. 19–20
  • The limiting second-moment diagonal entry is p∞,j,i = η B(2 − ηλj,i).This expression is given for the i-th diagonal entry in the eigenbasis.Source: C Proof of theorem 3.3 · p. 20

D Proof of theorem 3.5

The proof shows that an MDM decoding trajectory with logarithmically many steps can be evaluated by a polynomial-size, constant-depth threshold circuit. It does so by simulating one forward pass and enumerating candidate trajectories in parallel.

D Proof of theorem 3.5

  • One forward pass of the MDM can be computed by a polynomial-size, constant-depth threshold circuit.The block-causal attention simulation replaces the causal prefix endpoint i with the block endpoint r(i).Source: D Proof of theorem 3.5 · pp. 21–22
  • For T ≤ c log n, the number of candidate MDM trajectories is polynomial in n.Each trajectory is specified by T block-local decoding states from a constant-size set.Source: D Proof of theorem 3.5 · p. 22
  • For each candidate trajectory, all consistency checks Eτ,t can be evaluated in parallel.The state before step t depends on the input and earlier candidate actions, without requiring preceding model evaluations to run sequentially.Source: D Proof of theorem 3.5 · p. 22
  • The circuit Cn computes exactly the output of the original T-step MDM, and Cn,j ∈ TC0 for every output bit.The proof uses the unique trajectory consistent with the deterministic decoding rule.Source: D Proof of theorem 3.5 · p. 22

E Proof of theorem 3.6

The proof constructs a UDM that simulates a branching-program computation by updating a rewritable workspace at each decoding step. It establishes that every function in NC1 has such a simulation with the stated resource bounds.

E Proof of theorem 3.6

  • The UDM simulates one branching-program instruction per decoding step by maintaining σt in a rewritable workspace.The state is updated as σt = σt−1 ◦ ρt(x), starting from the identity element.Source: E Proof of theorem 3.6 · p. 22
  • At each step, attention retrieves the input bit at the instruction’s address, and the feed-forward layer computes the next permutation state.The workspace token is rewritten at the same location between decoding steps.Source: E Proof of theorem 3.6 · pp. 23–24
  • After T(n) decoding steps, the output is determined by whether the final permutation state equals α.The proof uses the branching program’s property that fn(x) = 1 iff σT(n) = α.Source: E Proof of theorem 3.6 · p. 25
  • The construction uses O(log n) hidden dimension, O(log n) bits of fixed-point precision, and O(1) workspace tokens.The proof also states that the branching-program length is polynomial in n.Source: E Proof of theorem 3.6 · p. 25
  • Every f ∈ NC1 belongs to TUDM[poly(n), poly(n), O(log n), 0, O(1)].The proof concludes this after constructing a UDM simulation for an arbitrary f ∈ NC1.Source: E Proof of theorem 3.6 · pp. 22, 25

Related papers

Loading 2610.23122v1…