PaperlayerView exact revision on arXiv

Source-linked AI summary

Edit Flows: Flow Matching with Edit Operations

Marton Havasi, Brian Karrer, Itai Gat, Ricky T. Q. Chen

arXiv:2506.09018v3cs.LGcs.AI

TL;DR

Non-autoregressive sequence models often rely on fixed token positions and cannot insert or remove tokens. Edit Flows model sequences as a CTMC over edit operations, and show improvements over mask models across benchmarks while surpassing autoregressive models in image captioning.

  • Problem

    Many non-autoregressive models use rigid, fixed-position representations and iteratively unmask or replace tokens, but cannot add or remove tokens.

    Source: §1 · p. 1

  • Method

    Edit Flows model generation as a CTMC over full sequences, predicting position-relative insertion, deletion, and substitution operations to support variable-length outputs.

    Source: §1 · pp. 1–2

  • Results

    Edit Flows consistently improved over mask construction across large-scale benchmarks, surpassed autoregressive models in image captioning, and achieved a 138% relative improvement over the mask model in code generation.

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

  • Takeaways & Limitations

    Edit operations provide a non-autoregressive approach to variable-length generation, with reported gains over mask models across image captioning, text, and code benchmarks.

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

    Training conditions on a randomly selected subset of the input, which may not match evaluation scenarios and may diminish the learning signal.

    Source: §6 · p. 11

Abstract

from arXiv · show

Autoregressive generative models naturally generate variable-length sequences, while non-autoregressive models struggle, often imposing rigid, token-wise structures. We propose Edit Flows, a non-autoregressive model that overcomes these limitations by defining a discrete flow over sequences through edit operations$\unicode{x2013}$insertions, deletions, and substitutions. By modeling these operations within a Continuous-time Markov Chain over the sequence space, Edit Flows enable flexible, position-relative generation that aligns more closely with the structure of sequence data. Our training method leverages an expanded state space with auxiliary variables, making the learning process efficient and tractable. Empirical results show that Edit Flows outperforms both autoregressive and mask models on image captioning and significantly outperforms the mask construction in text and code generation.

1 Introduction

Edit Flows model sequence generation as a CTMC over full sequences, using insertions, deletions, and substitutions to support position-relative, variable-length generation. Across large-scale benchmarks, they outperform all baselines on image-to-text generation and improve over the mask model in code generation.

1 Introduction

  • Unlike rigid non-autoregressive models that only unmask or replace tokens, Edit Flows can add or remove tokens and adapt output length to context.The cited prior models use fixed token positions and cannot add or remove tokens.Source: §1 · p. 1
  • Edit Flows represent generation as a CTMC over full sequences, estimating rates for insertions, deletions, and substitutions conditioned on the current sequence.This supports position-relative, variable-length generation without masking or padding tokens during training or inference.Source: §1 · p. 1
  • A Flow Matching training procedure adds auxiliary variables for a possible edit chain, yielding a tractable objective while keeping those variables hidden from the model.The model learns to infer the auxiliary variables.Source: §1 · p. 1
  • Edit Flows outperformed all baselines, including the autoregressive model, on image-to-text generation and achieved a 138% relative improvement over the mask model on code generation.The benchmarks include image-to-text at 280M parameters and code and open-ended text at 1.3B parameters.Source: §1 · pp. 1–2
  • Edit Flows define a sequence-level probability path for CTMC-based modeling across varying lengths, unlike prior work focused on token-level transitions.The reported evaluations span image captioning, open-ended text, and code generation.Source: §1 · p. 2

2 Preliminaries

The preliminaries introduce CTMCs as sequence-state processes governed by transition rates, then describe Discrete Flow Matching as transport between source and target distributions. Existing constructions commonly use token-wise paths, while mask construction simplifies transitions but has limits for variable-length generation.

2.1 Continuous-time Markov Chains

  • A CTMC rate specifies infinitesimal transitions between states, while the forward equation describes probability change as flow into a state minus flow out.Source: §2.1 · pp. 2–3

2.2 Discrete Flow Matching

  • Discrete Flow Matching couples source and target samples and uses conditional paths to train a CTMC that transports the source distribution to the data distribution.Source: §2.2 · p. 3
  • Common token-wise paths factorize changes by position and permit parallel sampling, but generalizing to broader sequence transitions can make conditional CTMC training intractable.Source: §2.2 · p. 3
  • Figure 3 forms aligned sequences, applies scheduler-selected edits to produce z_t, removes ε tokens to obtain x_t, and uses model output in a Monte Carlo loss estimate.Source: Figure 3 · p. 4, §2.2 · p. 4
  • Mask construction learns transitions between a special mask token and ordinary tokens, but lacks inherent variable-length support and padding can encourage overprediction of padding tokens.Source: §2.2 · p. 4

3 Edit Flows

Edit Flows define a CTMC over sequences whose transitions are insertions, deletions, or substitutions, allowing variable-length generation. Auxiliary alignments make training tractable, while localized edits and additional sampling procedures extend the framework.

3.1 Edit Flows: a continuous-time Markov chain using edit operations

  • Edit Flows parameterize a sequence-space CTMC whose transitions change the current sequence by one insertion, deletion, or substitution.The state space contains possible sequences up to a maximum length, enabling variable-length generation and encompassing existing constructions as special cases.Source: §3.1 · p. 4

3.2 Training Edit Flows

  • An auxiliary alignment process supplies a tractable training signal by selecting an edit path when multiple operation sets connect the same sequences.The alignment is sampled during training without being exposed to the model; the resulting training recipe permits Bregman-divergence objectives.Source: §3.2 · p. 6
  • The auxiliary-process Flow Matching result recovers a CTMC directly in the sequence space by marginalizing a time-dependent process in an augmented space.The result also applies beyond Edit Flows, including training with localized propagation rates.Source: §3.2 · p. 6

3.3 Algorithms and advanced techniques for Edit Flows

  • Sampling advances the learned CTMC by independently proposing insertions, deletions, and substitutions, then applying the edits simultaneously.The procedure uses a first-order approximation and assigns each operation probability hλt,i(Xt).Source: §3.3 · p. 7
  • Combining forward and reverse rates adds inference-time computation for self-correction without modifying the sample distribution.The method takes a forward step followed by a reverse-time step.Source: §3.3 · p. 7
  • Localized edit operations encourage locally consistent subsequences and yield a 48% increase in Pass@1 on code generation.The construction propagates edit occurrences to nearby operations rather than factorizing them independently.Source: §3.3 · pp. 7–8

4 Related work

Related work contrasts Edit Flows with token-wise discrete generation, heuristic edit-based variable-length models, and iterative correction methods. It also situates relative and localized operations as a shift from absolute-position constructions.

Discrete diffusion and flows for language modeling

  • Existing discrete diffusion reverses corruption, while discrete flows transport between distributions; prior token-wise constructions had not reliably surpassed the mask construction.Edit Flows instead use the CTMC construction to support variable-length generation through position-relative edits.Source: §4 · p. 8

Non-autoregressive variable length generation

  • Many variable-length approaches use a separate length predictor or heuristic edits; Edit Flows model sequence edits as a flow rather than an imitation-learned expert policy.Prior insertion-only or edit-based approaches include sequential prediction, while extending direct insertion derivations to deletion and substitution is described as challenging.Source: §4 · p. 8

Relative positions for language modeling

  • Token-wise non-autoregressive methods account for exact token positions, whereas Edit Flows use relative and localized operations in their method, sampling, and architecture.The paper suggests this may help explain performance but does not establish that explanation.Source: §4 · p. 8

Iterative editing models

  • Unlike iterative correction systems that separate generation from refinement or resample whole sequences, Edit Flows generate directly through token edits with one model.The cited contrasts include constrained samplers and methods using external objectives or backpropagation at test time.Source: §4 · pp. 8–9

5 Experiments

The experiments compare Edit Flows with autoregressive and Mask DFM baselines across image captioning, text, and code generation; performance varies by task.

5 Experiments

  • The baselines are an autoregressive model and Mask DFM, the most relevant and best-performing non-autoregressive text-generation framework in the comparison.Image captioning uses MS COCO and Image Captioning 3M; code uses the CodeLlama datamix, and text uses DCLM-baseline 1.0.Source: §5 · pp. 9–10
  • The evaluated variants are default Edit Flow, Uniform X0 + Edit Flow, and Localized Edit Flow, which encourages localized edits during generation.The default begins from the empty sequence and uses insertions and deletions; Uniform X0 begins from training-token samples.Source: §5 · p. 9
  • With the same compute and memory as Mask DFM, Edit Flows process 3× more training data per iteration using variable-length batches without special padding tokens.All models and baselines are trained with the same compute budget.Source: §5 · p. 9

Image captioning

  • For short captions, non-autoregressive models can outperform autoregressive models, and Edit Flow improves substantially over other models; Localization is on par with default Edit Flow.Training on Image Captioning 3M lets the models match the performance of cited prior works using large pretrained models.Source: §5 · p. 10

Text benchmarks

  • On text benchmarks, Edit Flow is significantly better than Mask DFM but slightly behind the autoregressive model.The evaluation uses zero-shot benchmarks with Llama3 1.3B models trained on DCLM-baseline 1.0.Source: §5 · p. 10, Table 2 · p. 10

Code benchmarks

  • On MBPP, Localized Edit Flow improves Pass@1 by 48% over non-localized Edit Flow and 138% over Mask DFM.The code benchmark uses Llama3 1.3B models trained on the CodeLlama datamix.Source: §5 · p. 10, Table 3 · p. 10

6 Limitations

The authors identify two empirical limitations: a pretraining setup favoring autoregressive models and text benchmarks that assess likelihood rather than generated-text quality.

6 Limitations

  • The pretraining configuration favors autoregressive models because they see all conditioning contexts, whereas Edit Flow training uses a randomly selected subset that may mismatch evaluation contexts.Such a mismatch may diminish the resulting learning signal.Source: §6 · p. 11
  • Text benchmarks evaluate likelihood rather than generated-text quality; because non-autoregressive models lack closed-form likelihood, the study uses a noisy ELBO estimate.These benchmarks do not directly assess the quality of text generated by the models.Source: §6 · p. 11

7 Conclusion

The conclusion presents Edit Flows as a position-relative, variable-length approach with consistent gains over mask-based construction, while initial results remain behind autoregressive models in text and code. Qualitative examples show stochastic caption generation and iterative code correction.

7 Conclusion

  • Edit Flows use position-relative edits and support variable-length generation, with consistent improvement over the mask construction across large-scale benchmarks.Initial results surpass autoregressive models in image captioning but fall slightly behind in text benchmarks and code generation.Source: §7 · p. 11

A Qualitative Examples

  • The examples include stochastic captions and code correction in which the model makes 117 edits over 300 steps to reach a correct implementation from an incorrect one.Intermediate states can contain extra tokens that are later deleted.Source: Figure 5 · p. 15, Figure 6 · p. 16, Figure 7 · p. 17

B Theorems and proofs

The appendix establishes how rates on an augmented sequence-and-auxiliary space generate the desired joint process, then gives a construction for a time-independent conditional and its deterministic special case.

B Theorems and proofs

  • Theorem 3.1 states that an augmented-space rate generating p_t(x,z) supports the flow-matching construction.The appendix then uses a lemma to simplify construction of such a rate.Source: B Theorems and proofs · pp. 17–18
  • The proof also shows that marginalizing the augmented process yields a sequence-level rate and distribution satisfying the forward equation.The argument checks rate conditions for the marginalized rate.Source: B Theorems and proofs · pp. 17–18
  • When p_t(x|z)=p(x|z) is time-independent, a rate over Z generating p_t(z) yields a rate over X×Z generating p_t(x,z)=p(x|z)p_t(z).The proof checks rate conditions and the Kolmogorov forward equation.Source: B Theorems and proofs · pp. 18–19
  • For a deterministic conditional p_t(x|z)=δ_f(z)(x), the appendix specializes the joint-flow construction to x=f(z).The proof applies the preceding lemma and evaluates the rate where x=f(z).Source: B Theorems and proofs · p. 19

B.1 A Bregman divergence as the training loss for Edit Flows

The section describes a Bregman-divergence training loss and advanced Edit Flow techniques, including reverse-time rates, classifier-free guidance, and localized propagation paths.

B.1 A Bregman divergence as the training loss for Edit Flows

  • The training loss is expressed using the Bregman divergence associated with the chosen ϕ.The supplied passages introduce the corresponding divergence and then state that the training loss can be written using it.Source: §B.1 · p. 20

C Advanced techniques for Edit Flows

  • Sampling uses a first-order approximation: it independently determines insertion, deletion, and substitution probabilities, then performs the edits simultaneously.For insertion or substitution, the new token is sampled from the corresponding edit distribution.Source: C Advanced techniques for Edit Flows · p. 20
  • Naïve rate CFG performed best across benchmarks, fixed rate CFG was close, and weighted rate CFG was consistently worse.When combined with reverse rates, CFG was applied to both forward and reverse rates.Source: C Advanced techniques for Edit Flows · p. 21
  • Combining forward and reverse rates can provide corrector steps that preserve the path marginal while correcting marginal-distribution errors.The section also states that (1 + α) times the forward rate plus α times the reverse rate generates the same path in forward time for α ≥ 0.Source: C Advanced techniques for Edit Flows · p. 21

C.1 Localized propagation paths

  • The localized path addresses long-sequence noise by letting true mask values encourage neighboring mask values to become true.The construction uses independent CTMC processes whose diagonal entries switch independently and whose off-diagonal entries propagate to neighbors.Source: §C.1 · p. 22
  • Localized-path mask variables can be sampled in two steps, and the computations within each step can be fully parallelized.The steps sample independent switch times and then the number of neighbors reached by propagation.Source: §C.1 · pp. 23–24

D Training data analysis

Figure 9 reports that 54% of training sequences are shorter than 1024 tokens; 57% of those are self-contained.

D Training data analysis

  • 54% of training sequences are shorter than 1024 tokens, and 57% of those start with <BOS> and contain fewer than 1024 tokens total.The figure defines these as self-contained sequences.Source: Figure 9 · p. 25

E Further experimental details

The experimental details describe training compute, model conditioning and architecture, and how sampling settings were selected for evaluation.

E Further experimental details

  • Edit models ingested 6T tokens during 500,000 training steps, compared with 2T tokens for Autoregressive and Mask DFM models.All models used batch size 4096 distributed across 16 × 8 H100 GPUs; Edit models do not use compute for missing sequence tokens.Source: E Further experimental details · p. 25
  • Image conditioning used frozen CLIP embeddings appended as prompts and mapped to the sequence-model dimension with a one-layer MLP projector.For text and code, the conditioning portion was randomly chosen; 10% of sequences dropped conditioning for unconditional prediction and CFG scaling.Source: E Further experimental details · pp. 25–26
  • For code-benchmark pass@1 and pass@10 evaluation, sampling parameters—temperature, top_p, sampling steps, CFG, and the divergence-free component—were tuned separately for each model to maximize performance.Source: E Further experimental details · p. 26

F Model preference for minimal edits

On a toy dataset, the learned coupling favors sequence pairs requiring fewer edits rather than preserving the uniform training coupling.

F Model preference for minimal edits

  • The toy dataset used only insertions and deletions, with both endpoint distributions uniform over length-four strings containing A and B.The training coupling itself was uniform, and the path deleted every source character and inserted every target character.Source: F Model preference for minimal edits · pp. 26–27
  • AAAA was 20× more likely to generate AAAA than BBBB, corresponding to no edits versus four insertions and four deletions.More generally, the highest learned coupling values corresponded to pairings requiring few edits, and the lowest to pairings requiring many.Source: F Model preference for minimal edits · p. 27

Related papers

Loading 2506.09018v3…