PaperlayerView exact revision on arXiv

Source-linked AI summary

Discrete Flow Matching

Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, Yaron Lipman

arXiv:2407.15595v2cs.LGcs.AI

TL;DR

Flow and diffusion models excel on continuous signals, but their performance on discrete sequential data still trails autoregressive models. Discrete Flow Matching generalizes discrete flows across probability paths, and the authors report improved generative perplexity and coding-task results among non-autoregressive models.

  • Problem

    Flow and diffusion models still perform less well on discrete sequential data than autoregressive models.

    Source: §1 · p. 1

  • Method

    Discrete Flow Matching defines a general framework over probability paths and a unified closed-form formula for generating probability velocities.

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

  • Results

    The method improves generative perplexity over current non-autoregressive methods and solves coding tasks at rates not previously achieved by such models, as far as the authors are aware.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The framework provides a large design space of discrete non-autoregressive generative models for further exploration.

    Source: §5 · p. 10

    A simplification of the generating-velocity formula applies only to conditional paths with independent coupling and an i.i.d. source.

    Source: §2.4 · p. 7

Abstract

from arXiv · show

Despite Flow Matching and diffusion models having emerged as powerful generative paradigms for continuous variables such as images and videos, their application to high-dimensional discrete data, such as language, is still limited. In this work, we present Discrete Flow Matching, a novel discrete flow paradigm designed specifically for generating discrete data. Discrete Flow Matching offers several key contributions:(i) it works with a general family of probability paths interpolating between source and target distributions; (ii) it allows for a generic formula for sampling from these probability paths using learned posteriors such as the probability denoiser ($x$-prediction) and noise-prediction ($ε$-prediction); (iii) practically, focusing on specific probability paths defined with different schedulers improves generative perplexity compared to previous discrete diffusion and flow models; and (iv) by scaling Discrete Flow Matching models up to 1.7B parameters, we reach 6.7% Pass@1 and 13.4% Pass@10 on HumanEval and 6.7% Pass@1 and 20.6% Pass@10 on 1-shot MBPP coding benchmarks. Our approach is capable of generating high-quality discrete data in a non-autoregressive fashion, significantly closing the gap between autoregressive models and discrete flow models.

1 Introduction

Discrete Flow Matching adapts flow-based generation to discrete data with broader probability paths, posterior-based velocity formulations, and scheduler-controlled sampling. The paper reports coding-benchmark results and conditional-text perplexity as evidence of its generative performance.

1 Introduction

  • Discrete diffusion and flow approaches had adapted continuous methods through embeddings or processes over discrete state spaces, but still lagged autoregressive models on discrete sequences.The cited approaches include Dieleman et al. (2022) and Stark et al. (2024) for continuous embeddings, and Austin et al. (2021a) and Campbell et al. (2022) for discrete state spaces.Source: §1 · p. 1
  • A 1.7B-parameter model achieved 6.7% Pass@1 and 13.4% Pass@10 on HumanEval, and 6.7% Pass@1 and 20.6% Pass@10 on 1-shot MBPP.The model was trained on the same data mix as Llama-2 and CodeLlama.Source: §1 · p. 1
  • Conditional text generation reached a generative perplexity of 9.7, versus 22.3 for a 1.7B autoregressive model and 8.3 for Llama-2 7B.Perplexity was measured by the Llama-3 8B model.Source: §1 · p. 1
  • The framework allows arbitrary source-target couplings and time-dependent schedulers, and expresses generating probability velocity directly through learned posteriors and schedulers.It also provides a unified theory and algorithm for corrector sampling and iterations; the paper reports that tuning path and corrector schedulers substantially improves generation quality.Source: §1 · p. 1
  • The authors frame Discrete Flow Matching as a step toward bridging the performance gap between discrete diffusion and autoregressive models, while noting further improvements may be possible.They point to the method’s design space as an avenue for further enhancements.Source: §1 · p. 2

2 Discrete Flow Matching

Discrete Flow Matching models sequences as discrete token states and constructs source-to-target probability paths whose sampling dynamics are learned through discrete-time probability velocities. The framework supports flexible couplings and schedulers, and its denoiser- and noise-prediction velocity forms match those of continuous Flow Matching.

2.1 Setup and notations

  • Discrete sequences are arrays of N tokens drawn from a vocabulary of size d, giving state space D = [d]^N.Source: §2.1 · p. 2

2.2 Source and target distributions

  • The method trains on source-target pairs under a coupling, focusing on all-mask sources and using either independent U-coupling or partially masked C-coupling, which improves conditional sampling.Source: §2.2 · pp. 2–3

2.3 Probability paths

  • A predefined path interpolates source and target distributions using arbitrary source-target couplings, with conditional paths required to equal the source and target at the endpoints.Source: §2.3 · p. 3
  • Conditional paths can be convex mixtures of component probabilities with time-dependent schedulers, including convex-interpolant and uniform-noise instantiations.Source: §2.3 · p. 3

2.4 Generating Probability Velocities

  • Discrete sampling uses a continuous-time Markov chain that independently updates each token, with probability velocities generating the path through the discrete continuity equation.Source: §2.4 · pp. 4–5
  • With the same schedulers, denoiser and noise-prediction parameterizations yield generating velocity fields of the same form as in continuous Flow Matching.Source: Table 1 · p. 6, §2.4 · p. 7
  • For perfectly trained posteriors, combining forward- and backward-time velocities supports valid corrector sampling that either generates the path or converges to its distribution at fixed time.The simplified backward-time velocity under the specified path, independent coupling, and i.i.d. source does not require estimating p0|t.Source: §2.4 · p. 7
  • Training minimizes posterior prediction loss, whose minimizer is the corresponding posterior used to construct the generating velocity.Source: §2.4 · p. 7

3 Related work

The section situates Discrete Flow Matching against discrete flows, masked modeling, and discrete diffusion, emphasizing its generalized CTMC formulation and its connection to MaskGit.

Discrete Flows

Masked modeling

  • With masked source p(x)=δm(x), Discrete Flow Matching yields a time-independent probability denoiser, connecting it to MaskGit (Ghazvininejad et al., 2019; Chang et al., 2022).Unlike MaskGit’s confidence-prioritized unmasking, its sampling independently unmasks tokens according to a rule justified asymptotically; corrector sampling can re-mask tokens while preserving pt and improves generation quality.Source: §3 · p. 8

Discrete diffusion

4 Experiments

Experiments evaluate language modeling, code generation, and fully discrete CIFAR10 image generation. The reported results cover generative perplexity, coding success, and image quality, with scheduler choice identified as important for performance.

4 Experiments

  • Experiments use masked-source probability paths, unconditional and conditional couplings, a probability denoiser, and scheduled corrector steps; the authors identify scheduler choice as pivotal to performance.The setup uses the generating velocity for sampling and evaluates scheduler choices in an appendix ablation.Source: §4 · p. 9

4.1 Language modeling

  • On 150M-parameter OpenWebText models, the method outperforms all non-autoregressive baselines in generative perplexity across all tested NFE values.The comparison uses basic sampling settings without temperature, top-probability, or corrector steps, and tests paths defined by equations 9 and 10.Source: §4.1 · p. 9
  • For conditional text generation, the 1.7B-parameter models narrow the generative-perplexity gap with autoregressive models while matching Llama-3 8B in entropy; C-coupling yields slightly better perplexity than U-coupling.Conditions are prefixes from the first 1,000 OpenWeb samples.Source: §4.1 · p. 9

4.2 Image generation

4.1 Language modeling

  • The 1.7B-parameter U-coupling model achieves non-trivial results on HumanEval and 1-shot MBPP, demonstrating non-autoregressive coding-task capability.The evaluation reports pass@{1, 10, 25}; oracle length conditions on the ground-truth solution length.Source: §4.1 · pp. 9–10

5 Conclusions and future work · Appendix · A Related works, continuation

Discrete Flow Matching expands the design space of discrete non-autoregressive generation and supports language models with improved generative perplexity and coding-task performance. The paper identifies sampling efficiency and probability-path exploration as future directions, while the continuation reviews related continuous, autoregressive, and masked-generation approaches.

5 Conclusions and future work

  • Discrete Flow Matching improves generative perplexity over current non-autoregressive methods and solves coding tasks at rates the authors say had not previously been achieved by such models.The authors report fewer network evaluations than autoregressive models, but lower sampling efficiency than the continuous counterpart.Source: §5 · p. 10
  • The authors identify improving sampling efficiency and exploring probability paths beyond those studied as future work.Discrete Flow Matching still requires more sampling effort than its continuous counterpart despite reducing network evaluations relative to autoregressive models.Source: §5 · p. 10
  • The authors see potential for discrete non-autoregressive models to close or surpass the gap with autoregressive models and state that their paradigm adds no significant societal risks beyond existing large language models.Source: §5 · p. 10

A Related works, continuation

B Further implementation details

The appendix details safe adaptive sampling, prefix conditioning, an NFE bound for mask modeling, and how trained posteriors can be used with changed schedulers.

Safe sampling

  • An arbitrary sampling step can yield negative probabilities and clamping error, while an adaptive step guarantees a valid PMF with o(h) sampling error.The tradeoff is potentially more function evaluations; the paper also specifies an adaptive step for its probability-denoiser parameterization.Source: B Further implementation details · p. 14

Conditioning

  • Training samples a random-length prefix mask, and conditional sampling restores conditioned tokens after each update.The prefix length N0 is sampled uniformly from 0 to N during training.Source: B Further implementation details · p. 15

NFE bound

  • For mask modeling without a corrector, Algorithm 1 requires at most N forward-pass evaluations because unchanged states need no recomputation.The probability denoiser is time-independent in this setting, so the forward pass need not be repeated when no mask token has been unmasked.Source: B Further implementation details · p. 15

Post training scheduler change

  • A trained posterior for a conditional path can be used with a different scheduler via the posterior scheduler-change formula, while mask-modeling posteriors remain unchanged.For mask modeling, the posterior is time-independent, so changing the scheduler does not affect it.Source: B Further implementation details · p. 15

C Code infilling

The section evaluates code infilling by masking spans and predicting them from unmasked tokens, a setup described as most similar to training. Performance improves with less input masking, and oracle prediction length boosts pass@1 on fully masked sequences.

C Code infilling

  • Code infilling predicts masked spans from unmasked tokens, and this evaluation setup is described as most similar to the training process.Prompts contain various masked spans.Source: C Code infilling · p. 15
  • On HumanEval, 1.5B-parameter models’ pass@1 and compiles@1 improve as input masking decreases across rates from 0.0 to 1.0.The evaluation provides oracle lengths for masked spans; Figure 4 shows the two metrics as a function of masking rate.Source: C Code infilling · p. 15, Figure 4 · p. 15
  • For fully masked sequences, providing the oracle prediction length significantly improves pass@1 scores, increasing the reported value from 6.7 to 11.6.Source: C Code infilling · p. 16

D Ablations

The square path scheduler achieved the highest text performance and was used for language models. For corrector sampling, the authors chose a = b = 0.25 to balance perplexity and entropy.

Path scheduler choice

  • The a = 0, b = 2 square scheduler achieved the highest text performance and was used for language models.The authors describe it as a cubic polynomial scheduler equivalent to a square function.Source: Figure 6 · p. 16, D Ablations · p. 16
  • The path-scheduler training comparison used 150M-parameter models for all a, b ∈ {0, 1, 2, 3}, generating 1,000 samples per model at temperatures τ ∈ {0.8, 0.9, 1}.Samples were generated with Algorithm 1 and the scheduler each model was trained on, then evaluated for generative perplexity with GPT-2.Source: D Ablations · p. 16

Corrector scheduler

  • On 1.7B-parameter models, smaller corrector values of a and b lowered perplexity but reduced entropy, so a = b = 0.25 was chosen as a balance.The corrector ablation used the optimal path schedulers from the preceding section.Source: D Ablations · pp. 16–17, Figure 7 · p. 17

E Theory and proofs … E.5 Backward-time generating velocity for i.i.d. source p(x0) and simple paths

The theory appendix derives discrete-flow divergence and velocity identities, including conditional-to-marginal generation, forward and backward conditional-path velocities, and a simpler reverse-time velocity for i.i.d. sources.

E.1 Computation of the discrete divergence

  • Discrete divergence at a state is computed by adding outgoing flux and subtracting incoming flux.The derivation uses that the relevant indicator is nonzero only when the states match or differ at the token being considered.Source: §E.1 · pp. 17–18

E.2 Conditional velocities lead to marginal velocities

  • Theorem 2 shows that conditional probability velocities generate the marginal path when averaged as a posterior-weighted convex combination.The proof differentiates the marginal path and applies Bayes’ rule to obtain the conditional-velocity average.Source: §E.2 · p. 18

E.3 Probability velocities generating conditional probability paths

  • Theorem 3 provides a generating probability velocity for the conditional paths, and its proof verifies the local conditions and continuity equation.The paths are those defined in equations 7 and 8.Source: §E.3 · pp. 19–20

E.4 Backward-time generating probability velocity.

  • For backward sampling, if the continuity equation holds and −ut meets the required conditions, stepping from Xt ∼ pt approximates Xt−h ∼ pt−h + o(h).The backward-time generator is −ut for the conditional paths defined in equations 7 and 8.Source: §E.4 · pp. 20–21
  • The backward-time generating velocity is obtained by replacing the forward construction’s argmin with argmax.The backward sampling formula is justified under the continuity-equation and velocity conditions.Source: §E.4 · p. 20

E.5 Backward-time generating velocity for i.i.d. source p(x0) and simple paths

  • Under the i.i.d.-source and independent-coupling setup, the simple-path condition makes the paths independent of the i-th source token.This condition supports the simplified backward-time generating velocity.Source: §E.5 · p. 21
  • For an i.i.d. source and independent coupling, a backward-time generating velocity can be simplified without learning the posterior p0|t.The simplified velocity is equivalent to the general one because their difference is divergence-free.Source: §E.5 · p. 21

E.6 Corrector steps

For perfectly trained posteriors, Theorem 4 characterizes when the velocity generates the probability path and when fixed-time corrector steps converge to its distribution.

E.6 Corrector steps

  • With αt, βt > 0 for t∈(0,1), ū_t is a probability velocity: αt−βt=1 generates pt, while αt−βt=0 yields fixed-time convergence for sufficiently small steps.The theorem assumes perfectly trained posteriors. The generating claim follows because the velocity satisfies the continuity equation.Source: §E.6 · p. 22
  • For αt−βt=0, the corrector transition has pt as its stationary distribution and converges when it is irreducible and aperiodic.For sufficiently small h, the transition is a convex combination of PMFs; strict positivity of p1|t or p0|t ensures irreducibility, and a positive self-transition ensures aperiodicity. Iterations change one token at a time, with parallel sampling via equation 12 as an approximation.Source: §E.6 · p. 22

E.7 Training … F Inference time

The appendices characterize training and posterior properties, give continuous-flow parameterizations and a scheduler-change result, and report inference latency relative to autoregressive decoding.

E.7 Training

  • The minimizer of the training loss matches the target posterior, because minimizing cross-entropy makes the predicted distribution equal to that posterior.The proof states this equality for each i ∈ [N].Source: §2.4 · p. 7, §E.7 · p. 23

E.8 Time-independent posterior for masked modeling

  • For paths defined by equations 7 and 9 with source p(x) = δm(x), the posterior and probability denoiser are time-independent.The proof also notes time independence for dummy source distributions and convex paths.Source: §E.8 · p. 23

E.9 Continuous Flow Matching

  • Continuous Flow Matching’s denoiser parameterization obtains the marginal generating velocity field by marginalizing the conditional field with posterior pt(x1|x).The appendix supplies the conditional denoiser field and identifies its marginalization as the continuous-flow denoiser parameterization.Source: §E.9 · p. 24
  • Continuous Flow Matching’s noise-prediction parameterization obtains the marginal generating velocity field by marginalizing the conditional field with posterior pt(x0|x).The appendix provides the conditional field for pt(x|x0) and identifies its marginalization as the noise-prediction parameterization.Source: §E.9 · p. 24

E.10 Scheduler change formula

  • For a conditional probability path as in equation 9, Proposition 8 gives a scheduler-change relation for any two schedulers and their corresponding posteriors, with t′ = κ^-1κ′.The proposition also applies to the marginal probability path defined by equation 7.Source: §E.10 · pp. 24–25

F Inference time

  • At 256 NFEs, the proposed method is ∼2.5x faster than the autoregressive model on HumanEval (19.97 vs. 50.94 seconds per example); at 512 NFEs, latency is roughly equal.Measurements used one A100 GPU with 80 GB RAM and batch size 1; this compares latency, not throughput, which the passage says is significantly better for autoregressive decoding due to KV caching.Source: F Inference time · p. 25

G Experimental setup … I Textual generations

This section describes the text and image experimental setups, reports a scheduler ablation and code-generation examples, and presents code infilling alongside conditional and unconditional textual generations from the proposed method and baselines.

G.1 Text

  • Text experiments use 150M- and 1.7B-parameter DiT models, with OpenWebText for small models and Llama-2 or CodeLlama mixes for large models.Evaluation reports within-sequence token entropy and generative perplexity averaged over 1,000 generated samples; sampling uses a double-precision categorical sampler.Source: §G.1 · p. 26

G.2 Image

  • Cubic scheduling generally gives the best CIFAR10 FID, while FID varies substantially across scheduler choices.The study trains with Linear, Quadratic, Cubic, and Cosine schedulers and evaluates each model with all four; corrector sampling is disabled and temperature is 1.Source: Figure 8 · p. 27, §G.2 · p. 27

H.1 Success cases

  • Success-case examples show generated code for balance checking and character-shift encoding and decoding.The examples also include filtering integers and finding common elements between lists.Source: §H.1 · p. 28

H.2 Failure cases

  • Failure-case examples include attempted solutions for case inversion and rolling maximum tasks.The displayed case-inversion code returns the input when it is None, and the rolling-maximum snippet initializes its maximum to zero.Source: §H.2 · p. 30

H.3 Infilling

  • Infilling examples include code snippets for breadth-first search and binary search.The passage also includes a bubble-sort implementation fragment.Source: §H.3 · pp. 33–34
  • Infilling snippets show recursive binary_search calls and a Dijkstra single-source shortest-path implementation using an adjacency-matrix graph.Source: §H.3 · pp. 34–35

I Textual generations

  • Textual-generation examples include conditional and unconditional outputs alongside several baselines, with conditional prompts marked in gray.Source: I Textual generations · p. 36

I.1 Conditional generation · I.2 Unconditional generation

The passages cover a North Korea sanctions proposal, a Celtic squad error, and artist profiles, followed by loosely phrased advice on goals, planning, and timing and a reflection on a film project.

I.1 Conditional generation

  • The United States asked the UN Security Council to impose an oil embargo on North Korea and freeze Kim Jong-un’s assets.The passage says this was in response to revelations that North Korea had restarted nuclear work.Source: §I.1 · p. 36
  • A Celtic administrative error left Efe Ambrose out of the squad, although he was available for the Champions League first leg and faced a potential three-match ban.The account says the club wrongly believed he would be suspended for that match.Source: §I.1 · p. 36
  • Carl Jara is described as a Cleveland sand sculptor who learned mold-based techniques and later carved sand sculptures for clients and special events.The passage also says he won a Cleveland art competition and had worked as an artist for over 25 years.Source: §I.1 · p. 37

I.2 Unconditional generation

  • The advice recommends setting goals around what you want to accomplish and taking responsibility for your work and self-care.It says management cannot plan what happens to you and that you should make sure the goal is one you want to achieve.Source: §I.2 · p. 38
  • The advice frames action as a sequence of planned steps, urging readers to distinguish each action from the next step and adjust when motivation wanes.It also presents mistakes in day-to-day planning as opportunities to correct course.Source: §I.2 · p. 38
  • The passage says timing is part of taking action and urges readers to act promptly when they know what to do.It uses changes to team performance and customer service as an example.Source: §I.2 · p. 38
  • A speaker describes a film project as a complex, tragic story about conflict, individual needs, and human resilience.The speaker says the project brings together collaborators and actors and refers to a 35-year period of slavery.Source: §I.2 · p. 38

Related papers

Loading 2407.15595v2…