Edit-Based Refinement for Parallel Masked Diffusion Language Models

Conference: ICML 2026
arXiv: 2605.09603
Code: https://github.com/renhouxing/ME-DLM
Area: Diffusion Language Models / Parallel Decoding / LLaDA / Text Generation
Keywords: Masked Diffusion, edit-based refinement, edit distance supervision, parallel decoding

TL;DR

ME-DLM introduces a lightweight "decode-then-edit" stage to masked diffusion language models (e.g., LLaDA): the first stage performs standard unmasking to generate a draft, and the second stage applies parallel corrections using three types of token-level edits (replace/delete/insert). The supervision signal is derived from the shortest edit script (edit distance). With only 1/8 diffusion steps, it surpasses LLaDA-Instruct by +11.6 on HumanEval and +33.6 on GSM8K.

Background & Motivation

Background: Masked Diffusion Language Models (MDLM) such as LLaDA and Dream have achieved parity with autoregressive LLMs at the billion-parameter scale, with the main advantage being parallel decoding—filling multiple mask tokens in one step, saving time compared to autoregressive models.

Limitations of Prior Work: When the number of tokens predicted in parallel per step increases from 1 to 4, 8, or even 16, the generation quality drops sharply. The paper provides an intuitive example: with a training set containing only "2+2=4", "2+3=5", "3+2=5", the model, sampling tokens independently by marginal probability, can generate "2+2=5"—each token is locally most probable, but the combination violates arithmetic.

Key Challenge: The MDLM training objective is token-level cross-entropy \(\mathcal{L} \propto \mathbb{E}[-\log p_\theta(x_{0,i}|x_t)]\), modeling only the marginal distribution at each position. However, during parallel decoding, the model simultaneously takes argmax over the set \(\mathcal{S}\), implicitly assuming conditional independence across positions: \(p_\theta(x_{0,\mathcal{S}}|x_t)\approx\prod_{i\in\mathcal{S}}p_\theta(x_{0,i}|x_t)\). Marginal optimality ≠ joint optimality, which is the root cause of multi-token parallel decoding failures.

Goal: To address this "lack of joint consistency" without altering the LLaDA training paradigm or increasing the total number of diffusion steps.

Key Insight: The authors observe that the draft produced by parallel decoding is already close to correct, with only sporadic structural errors (extra/missing/wrong tokens). By retaining the parallel unmasking for the draft in the first stage and adding a lightweight edit refinement stage for local corrections, one can enjoy both the speed of parallelism and joint consistency.

Core Idea: Decompose the diffusion process into two stages: "mask diffusion (draft generation) + edit diffusion (local refinement)". The edit stage uses three token-level edit actions (replace/delete/insert), supervised by the shortest edit script (edit distance) from draft to target.

Method

Overall Architecture

A two-stage diffusion process, sharing the same set of parameters (LLaDA-8B) throughout, with three progressive training stages:

• Mask diffusion stage: Starting from a fully masked sequence, unmask tokens iteratively according to schedule \(\{t_K>\dots>t_0=0\}\), filling multiple mask tokens in parallel at each step to obtain a complete draft \(x^{(0)}\). This is identical to LLaDA.
• Edit diffusion stage: Starting from \(x^{(0)}\), at each step, the model predicts for each token position a pair of actions \((c_i,n_i)\)—\(c_i\in\mathcal{V}\cup\{\text{[DEL]}\}\) indicates replace/delete/keep at the current position, \(n_i\in\mathcal{V}\) indicates what to insert after the current position (if \(n_i=c_{i+1}\), no insertion). A deterministic operator \(A\) applies these actions to the sequence: \(x^{(t+1)}=A(x^{(t)},\{(c_i,n_i)\})\). All actions are predicted in parallel, but the application couples the entire sequence, introducing dependencies among tokens at the application layer.
• Termination: The process stops when the model predicts "no edit" (i.e., \(c_i=x_i\) and no insertion) for all positions or when the maximum number of steps is reached.

Key Designs

1. Token-level (c, n) edit actions + deterministic application operator:

   • Function: A pair of token-level outputs \((c_i,n_i)\) simultaneously expresses replace, delete, and insert operations, enabling parallel prediction but sequential application of edits.
   • Mechanism: The transition \(p_\theta(x^{(t+1)}|x^{(t)})\equiv\prod_{i=1}^{L_t}p_\theta(c_i,n_i|x^{(t)})\) is independent at the prediction layer, but the application operator \(A\) sweeps left-to-right: if \(c_i=\text{[DEL]}\), delete \(x_i^{(t)}\); otherwise, replace \(x_i^{(t)}\) with \(c_i\); then, if \(n_i\neq c_{i+1}\), insert \(n_i\) after position \(i\). For boundary cases, repeated insertions are represented canonically (e.g., "a a a b" as "insert a between a and b"), and prompt/generated boundaries are valid insertion points.
   • Design Motivation: The parallelism advantage of MDLM must be preserved, so the prediction must be factorized; however, joint consistency requires coupling between positions. The authors place "coupling" in the deterministic application step, circumventing the challenge of explicit joint modeling. This is the most ingenious engineering insight of the paper.

2. Edit distance supervision + canonical mapping:

   • Function: Provides deterministic token-level edit supervision \((c_i^\star,n_i^\star)\) during training, guiding the model to learn "minimal correction" rather than "major rewriting".
   • Mechanism: During training, the model first generates an intermediate state \(x^{(m)}\) (mask diffusion for n steps + edit diffusion for m steps), then computes the shortest edit script (replace/delete/insert sequence) from \(x^{(m)}\) to ground-truth \(x^\star\) using the classic edit distance algorithm, and maps the script to \((c_i^\star,n_i^\star)\) for each position via a canonical rule. If multiple tokens need to be inserted at the same position, only the first is supervised in the current step, with the rest deferred to subsequent steps.
   • Design Motivation: Edit distance is the deterministic, optimal (though non-differentiable) solution—given \((x^{(m)},x^\star)\), the shortest edit script (under the canonical mapping) is unique, providing unambiguous training signals. Learning only minimal edits encourages conservative corrections, so the model naturally outputs no edits and terminates once the sequence stabilizes, aligning with the convergence semantics of diffusion. This is much simpler than RL/RLHF for training edit policies.

3. Three-stage curriculum training + mask/edit step allocation at inference:

   • Function: Enables the model to gradually transition from "only mask diffusion" to "capable of edit refinement", avoiding error accumulation in early training.
   • Mechanism: (i) Stage 1 trains on Nemotron-Pretraining-SFT to predict the current and next token, laying the foundation for \((c_i,n_i)\) next-token prediction; (ii) Stage 2 fine-tunes on R1-Distilled data with standard masked diffusion, yielding a strong baseline; (iii) Stage 3 interleaves mask and edit training on the same data, gradually increasing \(m\) from 0. At inference, by default, 1/4 of the total budget is allocated to edit, with an edit step cap of 32; for example, with 1/8 budget (64 steps), 48 are mask and 16 are edit.
   • Design Motivation: Directly training edit leads to poor initial drafts and excessive editing burden; curriculum allows the model to first master "drafting" before gradually learning "refinement", a common trick in diffusion training. Allocating more steps to mask at inference is justified because better drafts require fewer edits; Table 3 confirms that with a 1/1 budget, only 6-9 edit steps are needed for convergence.

Loss & Training

• Three-stage progressive fine-tuning of the same LLaDA-8B parameters; Stage 1: lr=5e-5, batch=2048; Stage 2: lr=5e-5, batch=128; Stage 3: lr=1e-5, batch=128; total training time: ~213 hours on 64×H800 GPUs.
• Inference: mask diffusion → edit diffusion, edit step cap 32, early stop if all edits are empty.

Key Experimental Results

Main Results

On six math and code benchmarks, average gains under different budgets (Budget = total steps × tokens per step / sequence length):

Budget	LLaDA-Instruct	ME-DLM Stage-2	ME-DLM Stage-3	Stage-3 vs Stage-2
1/1	45.3	55.7	60.0	+4.3
1/2	42.5	50.7	55.4	+4.7
1/4	32.3	37.7	46.4	+8.7
1/8	20.9	19.3	32.6	+13.3

On specific datasets, at 1/8 budget (8 tokens per step, 64 steps total):

Dataset	LLaDA-Instruct	ME-DLM Stage-3	Gain
HumanEval	12.2	25.0	+12.8
HumanEval+	9.8	22.6	+12.8
MBPP	17.5	26.7	+9.2
GSM8K	50.3	83.8 (84.8 @ 1/1)	Highly significant
MATH-500	20.2	34.4	+14.2

(The "+11.6 HumanEval / +33.6 GSM8K" in the abstract refers to ME-DLM Stage-3 vs LLaDA-Instruct, using representative numbers from different budgets.)

Ablation Study

Step allocation experiments (1/8 budget = 64 total steps):

m/e (mask/edit)	HumanEval	GSM8K	Notes
64/0 (only mask)	Sharp drop	Sharp drop	Validates parallel decoding failure
32/32	Moderate	Moderate	Balanced but insufficient mask steps
48/16 (default)	Best	Best	More mask steps yield better drafts, 16 edit steps suffice
0/64 (only edit)	Sometimes OK at low budget	-	Flexible but less stable

Edit convergence steps:

Budget	Edit step cap	HumanEval actual	MATH-500 actual
1/1	32	6.2	7.4
1/2	32	21.6	17.8
1/4	32	27.6	24.1
1/8	16	15.2	14.7

Key Findings

• Smaller budgets yield greater edit gains: At 1/1 budget, Stage-3 vs Stage-2 is only +4.3, but at 1/8 it jumps to +13.3—indicating that edit refinement is designed to remedy aggressive parallel decoding.
• Increasing mask steps reduces required edit steps (Table 3): At 1/1, only 6-9 edit steps are needed for convergence; at 1/4, 26-27 steps are needed—confirming the intuition that better drafts require fewer edits, and that edit diffusion is a true convergence process rather than open-ended rewriting.
• The +33.6 on GSM8K is the most striking result: Jumping from 50.3 to 83.8 shows that mathematical reasoning is particularly sensitive to "joint consistency" (one wrong token ruins the whole problem), making edit refinement almost essential for such tasks.
• Improvements on code tasks are less pronounced than on math: Possibly because code tokens are strongly constrained (syntax), so parallel decoding errors are less severe and code can still compile; in math, one wrong number invalidates the entire answer.

Highlights & Insights

• "Factorized prediction + deterministic application coupling" is a clever decoupling design: It retains parallelism in prediction while enforcing joint consistency in deterministic application, circumventing the fundamental conflict between parallelism and joint modeling. This trick can be transferred to any token-level parallel generation framework (including non-mask diffusion, speculative decoding, etc.).
• Edit distance as a supervision signal is an underrated tool: While RLHF/DPO are popular in the LLM era, when the goal is "minimal correction rather than arbitrary rewriting", edit distance offers determinism, minimality, and computability, making it more robust than RL-trained reward models.
• "Self-generated trajectory for self-correction" training paradigm: Stage 3 uses the model's own draft as edit supervision, aligning the training and inference distributions and fundamentally avoiding exposure bias.
• Broad applicability: Diffusion video generation, diffusion speech, and parallel non-autoregressive MT all face the "marginal optimality ≠ joint optimality" issue; the ME-DLM approach is transferable.

Limitations & Future Work

• Requires three-stage progressive training, which is more complex than direct fine-tuning of LLaDA, and Stage 1 alone takes 150 hours.
• Edit stage requires self-rollout to generate training data, making each training step much more expensive than standard SFT.
• Diminishing returns as budget approaches 1/1: When parallelism is not aggressive, edit refinement offers little cost-effectiveness.
• Canonical mapping ambiguity: The same edit distance can correspond to multiple shortest scripts; the authors fix one, but it may not be optimal.
• Only validated on code + math: It remains untested whether "minimal editing" benefits open-ended generation or dialogue scenarios.
• Edit operations are token-level only: Cannot perform span-level rewrites (e.g., reordering entire paragraphs).
• Future directions: (i) Make edit steps dynamically adaptive rather than fixed; (ii) Learn a lightweight edit policy to replace edit distance supervision; (iii) Combine edit refinement with best-of-N sampling for test-time scaling.

Related Work & Insights

• vs Soft Mask / EvoToken: Both soften the mask representation, while ME-DLM leaves the mask unchanged and adds an edit stage post-decoding—a complementary approach. On GSM8K at 1/4 budget, ME-DLM outperforms Soft Mask by +14.3 and EvoToken by +4.3.
• vs LRD / adaptive stopping: LRD dynamically adjusts steps based on convergence, but budgets are not directly comparable; ME-DLM compares under fixed budgets for clearer conclusions.
• vs Speculative decoding / Medusa: Those are draft-verify approaches for autoregressive models; ME-DLM is draft-edit for diffusion, conceptually similar but technically distinct.
• vs CDLM / remasking methods: These attempt to correct erroneous tokens, but remasking can only replace, not insert/delete, and cannot handle "missing token" errors; ME-DLM supports all three operations.
• vs Levenshtein Transformer / EditNAR: Edit-based non-autoregressive MT approaches; ME-DLM transfers these ideas to diffusion LMs and, combined with LLaDA as a modern backbone, achieves large-scale model results.

Rating

• Novelty: ⭐⭐⭐⭐ Uses edit refinement to address joint consistency in parallel diffusion; the technical combination is simple and effective
• Experimental Thoroughness: ⭐⭐⭐⭐ Six math + code benchmarks, four budgets, step allocation and convergence ablations, but lacks open-ended generation evaluation
• Writing Quality: ⭐⭐⭐⭐⭐ The "2+2=5" failure example is highly persuasive; formulas and algorithm tables are very clear
• Value: ⭐⭐⭐⭐⭐ The +33.6 on GSM8K is a key contribution to the practical utility of diffusion LMs; code is open-sourced

Related Papers

• [ICLR 2026] Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding
• [ICLR 2026] Toward Safer Diffusion Language Models: Discovery and Mitigation of Priming Vulnerabilities
• [ICLR 2026] DreamOn: Diffusion Language Models For Code Infilling Beyond Fixed-size Canvas
• [ACL 2025] Segment-Level Diffusion: A Framework for Controllable Long-Form Generation with Diffusion Language Models
• [ACL 2025] DiffLM: Controllable Synthetic Data Generation via Diffusion Language Models