PaperlayerView exact revision on arXiv

Source-linked AI summary

Birth, Death and Diff: Insertion-Based Sequence Generation with Corrections

Durga Prasad Maram, Dhruvesh Patel, Tahira Naseem, Benjamin Rozonoyer, Gaurav Pandey, Ramón Fernandez Astudillo, Andrew McCallum

arXiv:2610.34310v1cs.LGcs.CL

TL;DR

Insertion-based masked diffusion models support flexible-length parallel generation, but irreversible decisions and uncertainty across alignments make aggressive decoding difficult. δMDM learns to revise its own proposals through deletion and substitution. It Pareto-dominates FlexMDM on speed–quality trade-offs, with reported gains on OpenWebText and TinyGSM.

  • Problem

    Insertion-based masked diffusion models face irreversible generation errors and diffuse token predictions from multiple alignments, both of which make decoding more conservative.

    Source: §1 · p. 1, §2.3 · p. 4, §2.2 · p. 4

  • Method

    δMDM learns correction operations from its own aggressive proposals and interleaves insertion and unmasking with deletion and substitution.

    Source: §1 · p. 2, §3.3 · p. 7

  • Results

    δMDM Pareto-dominates FlexMDM on speed–quality trade-offs, halves its OpenWebText generative perplexity at 1024 steps, and matches its TinyGSM accuracy with half as many NFEs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across the reported tasks, δMDM combines better inference-time scaling with a stable from-scratch training recipe based on discrete diffusion.

    Source: §7 · p. 10

    Training uses a single aggressive proposal step, so errors may compound across inference rounds and produce states unseen during training.

    Source: §6 · p. 10

Abstract

from arXiv · show

Insertion-based masked diffusion models enable flexible-length, parallel generation by interleaving mask insertion with token unmasking. However, because their generation decisions are irreversible, aggressive parallel decoding can introduce errors that subsequent steps cannot correct, necessitating conservative decoding with more forward passes (NFEs). Moreover, they exhibit higher entropy than standard MDMs because they marginalize over multiple alignments during training and therefore require even more conservative decoding in the absence of a correction mechanism. To address this limitation, we introduce ${\delta}$MDM, which makes intermediate generations revisable by equipping an insertion-based model with substitution and deletion operations. During training, we expose the model to errors caused by aggressive unmasking and insertion, teaching it to correct its own mistakes. At inference time, we interleave insertion and unmasking with correction steps that can revise previously unmasked tokens and remove unnecessary masks. This mechanism makes aggressive decoding practical, improving both generation speed and quality. ${\delta}$MDM Pareto-dominates FlexMDM, a vanilla insertion-based MDM, on the speed--quality trade-off. On high-diversity unconditional text generation on OpenWebText, it halves the generative perplexity relative to FlexMDM at the highest inference budget. It also achieves lower generative perplexity at all inference budgets compared to MDM trained from scratch with variable-length padded sequences. On low-diversity conditional text generation on TinyGSM, it matches FlexMDM accuracy with half as many NFEs.

1 INTRODUCTION

Insertion-based masked diffusion models can place tokens flexibly, but irreversible decisions and uncertainty across possible alignments make aggressive decoding difficult. δMDM adds self-correction, and evaluations report a better speed–quality trade-off than FlexMDM across tasks.

1 INTRODUCTION

  • Insertion-based models cannot revise unmasked tokens or remove unnecessary masks, while averaging over alignments spreads probability across possible target tokens.This multiple-alignments problem raises predictive entropy and contributes to conservative decoding.Source: §1 · pp. 1–2
  • δMDM interleaves insertion and unmasking with deletion and substitution, learning corrections from its own aggressive proposals without an auxiliary model.Time-blindness and the ability to correct mistakes allow more aggressive drafting while improving generation speed and quality.Source: §1 · p. 2
  • δMDM Pareto-dominates FlexMDM on speed–quality trade-offs, halves its OpenWebText generative perplexity at 1024 steps, and matches its TinyGSM accuracy with half as many NFEs.It also outperforms both FlexMDM and MDM at 1024 steps using 128 sampling steps.Source: §1 · p. 2
  • The paper also presents ablations of its design choices and qualitative analyses of δMDM generation behavior.The introduction identifies these as additional contributions beyond the evaluation across synthetic planning, text generation, and reasoning tasks.Source: §1 · pp. 2–3

2 PRELIMINARIES

The section lays out how insertion-based MDMs represent partial sequences and train and generate them, then specifies δMDM’s correction operators and procedures. It describes the multiple-alignments ambiguity and the irreversible-generation limitation that motivate correction.

2.1 INSERTION-BASED MASKED DIFFUSION MODELS

  • Insertion-based MDMs represent partial sequences with tokens and their ordered positions in the clean target, while dropped positions are tracked as gaps.The rightmost gap is immediately left of [E]; [D] is a bookkeeping marker, not part of the model vocabulary.Source: §2.1 · p. 3

2.2 TRAINING

  • The multiple-alignments problem makes both token and gap targets ambiguous because a partial sequence can align to the same target in several ways.The model observes x_t but not its alignment; a mask may correspond to different target tokens under different alignments.Source: §2.2 · p. 4

2.3 GENERATION

  • The standard sampler uses tau-leaping to superpose independent Poisson unmasking and insertion decisions at each time step, but cannot undo either decision later.Once a token is unmasked it cannot be changed, and an inserted mask cannot be removed, which the paper calls the irreversible generation problem.Source: §2.3 · p. 4

3.1 THE Expand AND Remove OPERATORS

  • Expand inserts masks in the gaps and shifts later cells, while Remove deletes selected cells and shifts those to their right leftward.For alignment bookkeeping, new cells match missing ground-truth positions from left to right; surplus insertions receive −1.Source: §3.1 · p. 5

3.2 TRAINING

  • δMDM trains one time-blind network on aggressive proposals, then uses substitution and deletion losses to correct wrong fills and remove surplus masks.Proposal sampling greedily fills masked positions and also permits substitutions at already-unmasked positions; insertion counts are sampled from a random fraction of predicted gap lengths.Source: §3.2 · pp. 5–6

3.3 GENERATION

  • Generation alternates proposal and correction steps so the model can substitute previously unmasked tokens, delete unnecessary new masks, and unmask remaining masks.Unless otherwise stated, each proposal step is followed by a correction step; both count toward the forward-pass budget.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

The experiments evaluate δMDM on unconditional OpenWebText generation, conditional GSM8K reasoning, and correction behavior. δMDM improves the speed–quality trade-off over FlexMDM, while analyses show how insertion scale and correction timing affect outcomes.

4.1 UNCONDITIONAL GENERATION ON OPENWEBTEXT

  • At T=1024, δMDM’s generative perplexity is 20.35, versus 40.37 for FlexMDM and 38.16 for MDM; at T=128, it beats both baselines at T=1024.The OWT comparison reports a 47% reduction relative to MDM at T=1024 and an 8× reduction in sampling steps at T=128.Source: Table 1 · p. 8, §4.1 · p. 8

4 EXPERIMENTS

  • The evaluation uses GPT-2-tokenized OWT sequences padded to 1024 tokens and compares models across forward-pass budgets.δMDM is trained for 400k steps, while MDM and FlexMDM are trained for 800k; each δMDM proposal and correction counts toward inference budget T.Source: §4 · p. 7, §4.1 · p. 8, Table 1 · p. 8

4.1 UNCONDITIONAL GENERATION ON OPENWEBTEXT

  • On OWT, δMDM has lower generative perplexity than FlexMDM at all budgets except T=32, and its token entropy remains stable across budgets.At T=1024, δMDM’s token entropy is generally below FlexMDM’s; compared with DILM-S, it has higher token entropy while attaining lower generative perplexity at reported budgets.Source: §4.1 · p. 8

4.2 CONDITIONAL GENERATION ON TINYGSM

  • On GSM8K, greedy decoding fills each selected mask with the highest-probability token; the test set contains 1.32K problems with exact-answer pass@1 evaluation.The training data consists of 11.8M grade-school math problems converted to Python solutions.Source: §4.2 · pp. 8–9
  • On GSM8K, δMDM outperforms FlexMDM at every budget and matches its accuracy with half the steps at low budgets.MDM is strongest at T≤250, but its accuracy plateaus around 40–41%; δMDM overtakes it from T=500 onward.Source: Table 2 · p. 9, §4.2 · p. 9

4.3 ANALYSIS

  • Substitutions peak early in generation, while deletions dominate near the end on both OWT and GSM8K.The authors relate early substitutions to low-context predictions and late deletions to excess masks; they hypothesize that lower entropy after the correction pass permits more unmasking per step.Source: §4.3 · pp. 9–10
  • For OWT, Gen. PPL is lowest at moderate insertion scales, about λ=1.1–1.25 for T≥256, and rises sharply beyond that range.On GSM8K, with correction from the start, accuracy rises to λ=2 at T=250 and λ=1.75 at T∈{500,1000}; late correction peaks at λ=1.5 and trails by roughly 10 points or more at the largest scales.Source: §4.3 · p. 9, Figure 5 · p. 9, Figure 6 · p. 9

5 RELATED WORK

Prior correction work primarily targets fixed-length MDMs, while approaches for variable-length settings address narrower operations or tasks. δMDM’s context is a literature that includes remasking, correction heads, and length-changing operations.

5 RELATED WORK

  • Most prior MDM correction approaches focus on fixed-length settings, while variable-length self-correction has received less attention.Wu et al. (2026) add adaptive mask insertion and removal for code infilling but do not consider substitutions.Source: §5 · p. 10
  • Correction methods include confidence-based remasking, added correction heads, and alternating unmasking and correction; Schiff et al. (2026) limit corrections to substitutions.Correction-head supervision has used replacement tokens, model-sampled tokens, or human-injected synthetic errors.Source: §5 · p. 10
  • Other approaches enable token flips through mixed noise or modify the denoising architecture, while Wu et al. (2026) use adaptive insertion and deletion for code infilling.The passage notes that Zhao et al. (2025) require architectural changes and retraining.Source: §5 · p. 10

6 LIMITATIONS

The corrector is trained on one aggressive proposal, while inference chains proposals built on the model’s earlier outputs, creating exposure bias when errors compound into unseen states.

6 LIMITATIONS

  • Inference-time errors can compound across proposal-correction rounds, producing states absent from single-step corrector training and resulting in exposure bias.The authors suggest unrolling several rounds or fully on-policy training as possible remedies, but say these would be considerably more complex.Source: §6 · p. 10
  • The current framework does not allow deletion of visible tokens because doing so would cause ambiguous supervision.Source: §6 · p. 10
  • More expressive corrections through visible-token deletion are therefore outside the current framework.Source: §6 · p. 10
  • Further scaling and adaptation studies are left for future work.Source: §6 · p. 10

7 CONCLUSION AND OUTLOOK

The conclusion presents δMDM as a variable-length masked diffusion model with insertion, deletion, and substitution, and reports speed–quality gains over FlexMDM. It also frames self-correction as a possible route to adapting pretrained diffusion models.

7 CONCLUSION AND OUTLOOK

  • δMDM Pareto-dominates FlexMDM on speed–quality and scales better with inference budget across unconditional and conditional text generation tasks.Its gains are associated with time-blind operation and training on corruptions produced by its own proposals.Source: §7 · p. 10
  • The authors suggest adapting pretrained diffusion models with self-correction objectives, in the spirit of fixed-point iteration, as a possible route to faster, higher-quality generation.Source: §7 · p. 10
  • The authors report using generative AI for parts of model implementation and analysis code, figure and table formatting, and polishing some manuscript text.The two main authors reviewed AI-generated code and content, and all authors take responsibility for the final work.Source: AI USE STATEMENT · p. 11
  • Anonymized source code is available, and trained checkpoints and logged training-loss curves are slated for public release upon acceptance.The paper points to Section 3 for algorithm details and Section 4 for datasets, preprocessing, and architecture.Source: REPRODUCIBILITY STATEMENT · p. 11

A SUMMARY OF NOTATION

B EXPERIMENTAL DETAILS

The experimental setup uses shared architecture and optimization settings, with dataset-specific sequence lengths and insertion scales. Proposal and correction steps both count toward the inference budget.

B EXPERIMENTAL DETAILS

  • δMDM is trained for 400k steps, while MDM and FlexMDM are trained for 800k steps.ILM and DILM-S use implementations and 800k-step checkpoints provided by Patel et al. (2026a).Source: B EXPERIMENTAL DETAILS · p. 16
  • OpenWebText sequences are padded to length 1024 and use insertion scale λ = 0.90; TinyGSM sequences use length 512 and λ = 1.75.Correction follows every proposal from offset 0, and both step types count toward budget T.Source: B EXPERIMENTAL DETAILS · p. 16
  • All trained models use the architecture and optimization settings listed in Table 4.The table’s parameter count covers the transformer body and excludes token embeddings.Source: Table 4 · p. 16

C ADDITIONAL TASKS

The star-graph task evaluates exact-match path generation across difficulty splits, while Figure 7 compares δMDM and FlexMDM over sampling budgets. The supplied table caption identifies the T=500 result scope.

C.1 SYNTHETIC PLANNING: STAR GRAPHS

  • The star-graph task asks the model to generate the unique path between given nodes and reports exact match at sampling budget T=500.It uses easy, medium, and hard splits with increasing difficulty from variable arm lengths and non-junction start nodes.Source: §C.1 · p. 16, Table 5 · p. 16
  • The input consists of a flattened edge list followed by the start node, target node, and EOS token.The insertion scale factor is 2, selected as the best value from a sweep over I.Source: §C.1 · p. 16
  • Figure 7 plots star-hard exact match against sampling budget for δMDM and FlexMDM at T ∈ {32, 64, 128, 256, 512}.Source: Figure 7 · p. 17

D.1 MULTIPLE ALIGNMENT PROBLEM

On OWT, FlexMDM’s predicted-token entropy is higher than MDM’s for most of generation, which the paper attributes to supervision spread across multiple candidate positions. A second figure examines δMDM entropy before and after correction at T=512.

D.1 MULTIPLE ALIGNMENT PROBLEM

  • FlexMDM has much higher predicted-token entropy than MDM for most of generation because multiple candidate positions share supervision for each mask.Figure 9 examines the same quantity for δMDM before and after correction at T=512.Source: §D.1 · p. 17
  • Figure 8 plots predicted-token entropy at masked positions over time for FlexMDM and MDM on OWT.The accompanying text attributes FlexMDM’s higher entropy to the multiple-alignments problem.Source: Figure 8 · p. 17, §D.1 · p. 17
  • Figure 9 compares δMDM’s masked-position token entropy before and after correction at T=512.Source: Figure 9 · p. 17, §D.1 · p. 17

D.2 TIME-BLINDNESS

This appendix reports time-blind FlexMDM results, examines entropy and correction behavior across settings, and describes insertion-scale, sequence-length, and related-work comparisons.

D.2 TIME-BLINDNESS

D.3 ENTROPY VS. GENERATION PERPLEXITY ACROSS SAMPLING BUDGETS

  • The appendix plots entropy versus generation perplexity for remaining budgets T∈{32, 64, 128, 256, 1024}, complementing the main-text T=512 plot.Source: §D.3 · p. 18, Figure 10 · p. 18

D.4 HOW ARE CORRECTIONS DISTRIBUTED?

  • Substitutions peak early in generation, whereas deletions increase toward the end, as sequences grow and isolated late substitutions become difficult.Correction intensity is the percentage of eligible positions corrected at a time step.Source: §D.4 · p. 18, Figure 11 · p. 19, Figure 12 · p. 19

D.5 GSM8K ACCURACY VS. INSERTION SCALE

  • Figure 13 plots GSM8K pass@1 against insertion scale at sampling budgets T=250, 500, and 1000.Source: Figure 13 · p. 20

D.6 LENGTH OF GENERATED SEQUENCES ON OWT

E EXTENDED RELATED WORK

  • Prior correction work includes remasking heads and alternating correction; variable-length methods include mask expansion and deletion, while DreamOn lacks substitutions.The cited discussion says DreamOn therefore leaves errors from carry-over unmasking persistent.Source: E EXTENDED RELATED WORK · pp. 20–21

F RANDOMLY SAMPLED GENERATIONS

The section presents randomly sampled OpenWebText generations and GSM8K examples, including a trajectory that highlights substitution steps.

F.1 OPENWEBTEXT

  • OpenWebText samples include sports reporting, Google and Android topics, and other news-like passages.Source: §F.1 · p. 21

F.2 GSM8K

  • GSM8K samples include arithmetic word problems about test scores, TV sales, and staged dehumidifier use.Source: §F.2 · pp. 23–24
  • A GSM8K trajectory with T=250 marks substitutions in orange across three intermediate steps and the final completion.Source: §F.2 · p. 25

Related papers

Loading 2610.34310v1…