Discrete Flow Matching

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

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.15595v2Jul 22, 2024cs.LGcs.AI

TL;DR

Discrete diffusion and flow models remain limited on high-dimensional discrete sequential data compared with autoregressive models. Discrete Flow Matching introduces a general discrete flow framework with unified probability-velocity formulas and reports improved non-autoregressive language-modeling results while narrowing the gap to autoregressive models.

Problem
Diffusion and flow models perform less well on discrete sequential data than on continuous signals and autoregressive models, motivating discrete generative approaches.
Method
Discrete Flow Matching generalizes discrete and continuous flow matching through broader probability paths, unified generating and corrector velocities, and selectable noise and corrector schedulers.
Results
Discrete Flow Matching outperforms all evaluated non-autoregressive baselines in generative perplexity across function-evaluation budgets and narrows the conditional-generation gap with autoregressive models.
Takeaways & Limitations
The framework provides a large design space for non-autoregressive discrete generation, including large-scale language models that improve perplexity and solve coding tasks at reported rates.
Takeaways & Limitations
The reported simplification applies under specific path assumptions, including a monotonically non-decreasing noise schedule and conditional paths with independent coupling and an i.i.d. source.

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 targets the weaker performance of diffusion and flow models on discrete sequences by extending flow matching to a broad discrete framework. The paper presents this approach as a potential bridge toward autoregressive-level performance while supporting non-autoregressive generation.

Motivation: Diffusion and flow models remain weaker on discrete sequential data than on continuous signals and autoregressive models.Prior discrete approaches either embed data into continuous space or operate directly over discrete state spaces.
Approach: Discrete Flow Matching generalizes discrete flow modeling with broader source-target probability paths, arbitrary couplings, and time-dependent schedulers.It also provides unified generating-velocity and corrector-sampling formulations.
Approach: The method is designed for high-quality discrete generation without autoregressive decoding.The introduction frames it as a discrete non-autoregressive generative approach.
Examples: Figure 1 illustrates code generation with gray code conditions and yellow model generations across left-to-right prompting and complex infilling.The middle and right subfigures show the infilling setup.

2 Discrete Flow Matching

Discrete Flow Matching defines a continuous-time Markov-chain framework for transforming discrete source samples into target samples along tractable probability paths. Its discrete probability velocities support sampling formulas analogous to continuous Flow Matching, with theoretical guarantees for generation and fixed-time convergence.

Setup and notations: Discrete sequences are modeled as N tokens drawn from a vocabulary of size d, giving the state space D = [d]^N.The paper represents distributions over these sequences with probability mass functions.
Source and target distributions: Source samples X0 ∼ p are transformed into target samples X1 ∼ q using training pairs drawn from a coupling π.Independent U-coupling and partially masked conditional C-coupling are considered, with the latter reported to improve conditional sampling.
Probability paths: The framework uses arbitrary source-target couplings and conditional probability paths whose endpoint distributions are concentrated on x0 and x1.Schedulers specify convex combinations of conditional probabilities and may vary by sequence location or be shared across tokens.
Generating Probability Velocities: Discrete Flow Matching models probability change through a continuous-time Markov chain whose tokens update independently according to learned probability velocities.Euler sampling uses a user-defined time step h, and the velocity may be directly modeled or parameterized through denoising or noise prediction.
Generating Probability Velocities: The discrete continuity equation measures outgoing and incoming probability flux through a discrete divergence over neighboring sequences differing by at most one token.This extends the continuous divergence intuition to the discrete state space.
Generating Probability Velocities: The generating velocity fields have the same form as continuous Flow Matching under denoiser or noise-prediction parameterizations, while a unified corrector velocity supports additional sampling procedures.For perfectly trained posteriors, the theory guarantees path generation when α_t − β_t = 1 and convergence to p_t under repeated fixed-time sampling when α_t − β_t = 0.

3 Related work

The paper builds on discrete-flow and discrete-diffusion work by generalizing probability paths, couplings, velocity formulas, and corrector sampling. It also connects masked Discrete Flow Matching to MaskGit while identifying a sampling difference that affects performance.

Discrete flows: Discrete Flows provides the closest foundation, while Discrete Flow Matching generalizes it to arbitrary source-target couplings and broader probability paths.The framework also unifies generating probability velocities in a closed-form formulation.
Discrete flows: Unified probability-velocity formulas recover the continuous Flow Matching form while extending it to discrete state spaces.The paper distinguishes probability velocity from a CTMC rate matrix because it is defined over pairs of discrete states.
Discrete flows: Noise and corrector schedulers provide a design space in which particular choices boost results, with κ_t = t reproducing Campbell et al. (2024).The general corrector velocity unifies corrector iterations and stochastic sampling.
Masked modeling: In masked modeling, the probability denoiser is time-independent, paralleling MaskGit’s unmasking probabilities.This follows from the time-independent posterior for the masked paths considered.
Masked modeling: Discrete Flow Matching samples token updates independently using its probability velocity, whereas MaskGit prioritizes tokens by confidence.The paper reports that confidence prioritization helps at very low NFE but introduces sampling bias at higher NFE.
Discrete diffusion: Discrete diffusion methods model categorical data through corruption processes, while later work extends continuous-time diffusion and probability-ratio learning to discrete spaces.These approaches include D3PM, Argmax flows, continuous-time discrete diffusion, and discrete score matching.

4 Experiments

The experiments evaluate Discrete Flow Matching across language, code, and fully discrete image generation, using masked sources and multiple couplings and paths. Results show improved generative perplexity, narrowed conditional-text gaps with autoregressive models, non-trivial coding performance, and strong CIFAR10 image results.

Experimental scope: Experiments cover language modeling, code generation, and fully discrete CIFAR10 image generation.The study uses small 150M-parameter and large 1.7B-parameter settings across unconditional, conditional, and coding tasks.
Experimental setup: The experimental setup uses masked sources, unconditional and conditional couplings, several probability paths, a probability denoiser, and velocity-based sampling.Scheduler and corrector choices are also varied in the experiments.
Language modeling: Our method outperforms all baselines in generative perplexity for every tested number of function evaluations in the 150M-parameter unconditional comparison.All methods use OpenWebText and basic sampling settings without temperature or corrector steps.
Language modeling: Conditional 1.7B-parameter models narrow the generative-perplexity gap with autoregressive models while maintaining entropy comparable to Llama-3 8B.The C-coupling model produces slightly better conditional-task perplexity than the U-coupling model.
Image generation: 3.63 FID at 1024 NFE is achieved on fully discrete CIFAR10, outperforming MaskGIT and Campbell et al. (2024).Inception score shows a similar trend; MaskGIT performs better at low NFE but deteriorates at higher NFE.

5 Conclusions and future work

The paper introduces Discrete Flow Matching as a general framework for discrete non-autoregressive generation. It reports improved generative perplexity and coding performance, while identifying sampling efficiency and probability-path exploration as future directions.

Discrete Flow Matching generalizes continuous flow matching and discrete flows into a framework for discrete non-autoregressive generation.
The framework provides a large design space for discrete non-autoregressive generative models.
The trained language models achieve improved generative perplexity compared with current non-autoregressive methods.
The models solve coding tasks at rates the authors report as previously unattained by non-autoregressive models.
Discrete Flow Matching reduces network evaluations relative to autoregressive models but remains less sampling-efficient than its continuous counterpart.
The authors identify exploring broader probability paths and further closing the gap with autoregressive models as future directions.

B Further implementation details

The implementation details address safe sampling, conditional generation, scheduler changes, and network-evaluation efficiency. Adaptive steps preserve valid probability distributions, while masked modeling can bound evaluations by the sequence length.

Safe sampling: Adaptive step sizes prevent invalid transition probabilities during sampling while retaining an o(h) sampling error.The fix may require more function evaluations.
Conditioning: Training conditions on prefixes of randomly sampled length N0 < N and adjusts the input according to the mask.The prefix length is sampled uniformly from 0 to N.
Conditioning: Conditional sampling replaces updated tokens with the specified conditioned tokens after each step.The update uses the condition mask to combine the current sample and condition.
NFE bound: For mask modeling, sampling without corrector steps requires at most N network evaluations.The forward pass need not be recomputed when no token is unmasked.
Post training scheduler change: A trained posterior can be used with a different sampling scheduler through the scheduler-change formula.For mask modeling, the posterior is time-independent and therefore unaffected by scheduler changes.

C Code infilling

The paper evaluates Discrete Flow Matching for code infilling with masked spans and reports that less input masking improves HumanEval performance. Oracle span lengths substantially help when the sequence is fully masked.

Code infilling: Code infilling predicts masked spans from unmasked context and most closely matches the training process.The evaluation uses prompts containing multiple masked spans.
Evaluation: HumanEval code-infilling experiments report Pass@1 and compiles@1 across masking ratios from 0.0 to 1.0.The model receives the oracle length for each masked span.
Results: Pass@1 and compiles@1 improve as the input masking level decreases.
Figure 4: Figure 4 plots Pass@1 and compiles@1 for 1.5B-parameter models as a function of input masking ratio on HumanEval.
Results: 6.7 vs. 11.6 Pass@1 shows that oracle prediction length substantially improves performance for fully masked sequences.

D Ablations

Ablations examine path, corrector, and temperature scheduler choices. The experiments select a square-function path scheduler and a corrector setting that balances perplexity against entropy.

Scheduler ablations: The ablations compare path schedulers, corrector schedulers, and temperature scheduling as distinct sampling-design choices.
Path scheduler: Path-scheduler experiments train 150M-parameter models across a,b ∈ {0,1,2,3} and evaluate samples at temperatures 0.8, 0.9, and 1.
Path scheduler: The cubic polynomial scheduler with a = 0 and b = 2 achieves the highest text-modality performance and is used for language models.This setting is equivalent to a square function.
Corrector scheduler: Corrector-scheduler ablations evaluate 1.7B-parameter models using parameter values a,b ∈ {0,0.25,0.5} and α ∈ {10,15,20}.
Corrector scheduler: Smaller corrector-scheduler a and b values lower perplexity but somewhat reduce entropy.
Corrector scheduler: The selected corrector setting a = b = 0.25 balances perplexity and entropy.

E Theory and proofs

The paper develops discrete probability velocities through conditional and marginal probability paths, including forward and backward-time formulations. For i.i.d. source distributions, it derives a simpler backward sampling velocity that avoids learning the posterior p0|t.

Conditional and marginal velocities: Theorem 2 shows that a conditional probability velocity generates the corresponding marginal probability path after marginalization with Bayes-rule posteriors.The proof differentiates the marginal path and relates conditional and marginal velocities.
Conditional path velocities: Theorem 3 provides a generating probability velocity for the conditional paths defined by equations 7 and 8.The result is established by verifying the velocity conditions and the Continuity Equation.
Backward-time sampling: The framework also derives a backward-time generating probability velocity for reverse sampling along the same conditional paths.Reverse-time sampling changes the optimization direction and uses the negative forward velocity.
I.i.d. sources and simple paths: For i.i.d. source distributions, the backward sampling formula can be simplified because the alternative velocity generates the same probability path while avoiding p0|t.The equivalence follows from a divergence-free difference between the two velocity fields.

E.6 Corrector steps

The corrector formulation uses a family of probability velocities parameterized by αt and βt. Under specific parameter settings, the resulting dynamics either generate the intended path directly or converge through repeated fixed-time sampling.

Corrector velocity: Theorem 4 states that, for perfectly trained posteriors and αt, βt > 0, the corrector velocity is a valid probability velocity.The theorem applies for t ∈ (0, 1) and establishes the Continuity Equation conditions.
Path generation: When αt − βt = 1, the corrector velocity generates the probability path pt.This follows because the corresponding velocity satisfies the Continuity Equation.
Corrector sampling: When αt − βt = 0, repeated sampling at fixed t with sufficiently small h converges to a sample from pt.The proof uses a transition matrix whose stationary distribution is pt and establishes convergence through irreducibility and a-periodicity.

E.7 Training

The training section establishes posterior-learning objectives and connects discrete denoiser and noise-prediction parameterizations to generating velocity fields. It also gives scheduler transformations and reports lower latency for non-autoregressive decoding at 256 NFEs.

Training objective: The minimizer of the training objective matches the model prediction ẇt(xi|Xt; θ) to the corresponding true conditional prediction.The objective reduces to cross-entropy minimization over tokenwise conditional distributions.
Denoiser parameterization: With source p(x) = δm(x), the posterior and probability denoiser are time-independent.This provides a theoretical basis for using a denoiser without explicit time dependence in this setting.
Velocity parameterizations: The generating velocity field can be parameterized using either probability denoising (x-prediction) or noise prediction (ε-prediction).The section supplies formulas for both parameterizations and their posterior-based marginalizations.
Scheduler transformations: Proposition 8 relates posteriors under any two schedulers through the inverse scheduler transformation.This allows scheduler changes to be expressed through transformed time coordinates.
Inference time: At 256 NFEs, the proposed method averages 19.97 seconds per HumanEval example versus 50.94 seconds for the autoregressive model.At 512 NFEs, both methods have roughly the same latency; the comparison concerns latency rather than throughput.

G Experimental setup

The experiments use text and code datasets, small and large DiT models, and task-specific evaluation metrics. CIFAR10 experiments compare schedulers and show substantial FID variation, with the Cubic scheduler generally performing best.

Data: The study uses OpenWebText, the Llama-2 data mix, and the CodeLlama data split across small and large model experiments.Small models use OpenWebText, while large models use the Llama-2 and CodeLlama mixes.
Models: The models include 150M-parameter small models and 1.7B-parameter large models using DiT transformer architectures.The large architecture uses 48 layers, 24 attention heads, and hidden dimension 1536.
Evaluation metrics: Generative perplexity is computed as the average likelihood of generated text under a second, usually stronger, model over 1000 samples.The experiments also report token entropy as a sequence-level diversity measure.
CIFAR10 scheduler ablation: CIFAR10 FID varies substantially across schedulers, with the Cubic scheduler generally performing best in both training and evaluation.Each scheduler is used for both training and evaluation to assess cross-scheduler effects.
CIFAR10 scheduler ablation: Figure 8 compares CIFAR10 FID across Linear, Quadratic, Cubic, and Cosine schedulers for training and evaluation.The experiment uses no corrector sampling and sets temperature to 1.
Coding evaluation: The coding evaluation includes HumanEval and 1-shot MBPP examples among the generated benchmark tasks.The surrounding evaluation material describes comparisons of the proposed schedulers and sampling algorithm with prior methods.

I.1 Conditional generation

The passages describe administrative, biographical, strategic, and cinematic material, without establishing a coherent account of conditional generation.

An administrative account describes a defender’s availability and a club’s mistaken suspension understanding.
A separate sentence states that the defender is available for the Maribor first leg despite his club’s belief that he should be suspended.
A biographical passage introduces Carl Jara as a professional sand sculptor from Cleveland, Ohio.
A strategic passage outlines planning goals and subsequent steps for executing actions at work.
A film-related passage emphasizes humour, humanity, tragedy, and an intense narrative about a fractured world.