PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Matching Guide and Code

Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky T. Q. Chen, David Lopez-Paz, Heli Ben-Hamu, Itai Gat

arXiv:2412.06264v1cs.LG

TL;DR

Flow Matching addresses the challenge of building generative models that transform simple source distributions into complex target distributions across diverse domains. The paper reviews FM’s mathematical foundations, design choices, and extensions, and provides code for applying it. FM supports extensions to discrete state spaces and manifolds, while its relationships to diffusion sampling are characterized under specified conditions.

  • Problem

    Generative modeling requires transforming samples from a known source distribution into samples from an unknown data distribution across varied state spaces and applications.

  • Method

    The paper presents a self-contained Flow Matching review centered on probability-path design, velocity-field regression, extensions, and accompanying PyTorch examples.

  • Results

    Under Gaussian-source, independent-coupling, fixed-scheduler, and score-parameterization conditions, diffusion Probability Flow ODE sampling is equivalent to FM sampling, while stochastic SDE sampling is equivalent to GM sampling.

  • Takeaways & Limitations

    FM provides a framework applicable to continuous, discrete, and manifold-valued generative tasks, with examples intended to support adoption and further development.

  • Takeaways & Limitations

    Geodesic conditional flows can be difficult on general manifolds without closed-form exponential and logarithmic maps and may concentrate probability at boundary points.

Abstract

from arXiv · show

Flow Matching (FM) is a recent framework for generative modeling that has achieved state-of-the-art performance across various domains, including image, video, audio, speech, and biological structures. This guide offers a comprehensive and self-contained review of FM, covering its mathematical foundations, design choices, and extensions. By also providing a PyTorch package featuring relevant examples (e.g., image and text generation), this work aims to serve as a resource for both novice and experienced researchers interested in understanding, applying and further developing FM.

1 Introduction

Flow Matching is a simple generative-modeling framework that has advanced image, video, speech, audio, protein, and robotics generation. This manuscript provides a self-contained reference covering FM’s foundations, design choices, and extensions, alongside code intended to support adoption and further development.

  • Flow Matching has pushed state-of-the-art performance across image, video, speech, audio, protein, and robotics generation.
  • The manuscript aims to provide a comprehensive, self-contained reference explaining Flow Matching’s design choices and community-developed extensions.
  • The accompanying codebase is intended to help newcomers quickly adopt and build upon Flow Matching for their own applications.
  • Flow Matching learns a velocity field whose ODE-defined flow transforms samples from a source distribution p into a target distribution q.
  • FM principles extend beyond R^d to discrete state spaces and Riemannian manifolds, supporting discrete generative tasks and applications such as protein folding.
  • Unlike diffusion’s forward-noising and reversal formulation, diffusion models can be viewed as a limiting case of bridge matching.

2 Quick tour and key concepts

Flow Matching constructs a probability path from a source distribution to a data distribution, learns its velocity field by regression, and samples by solving the resulting ODE. Conditional path construction makes the objective tractable, while a standalone PyTorch implementation demonstrates the complete workflow.

  • Core recipe: FM builds a path (p_t) from source p_0 = p to target p_1 = q, then trains a neural velocity field to describe its instantaneous sample velocities.
  • Core recipe: After training, sampling draws X_0 ∼ p and solves the velocity-field ODE to obtain a novel target sample X_1 ∼ q.
  • Core recipe: An ODE is determined by a time-dependent vector field that defines a time-dependent flow ψ_t.
  • Path design: The FM design uses a source distribution p := p_0 = N(0, I) and aggregates conditional paths associated with individual target examples.
  • Training: Conditioning converts the intractable marginal velocity-field objective into a tractable conditional Flow Matching loss.
  • Path design: The conditional path uses X_t|1 = t x_1 + (1 − t)X_0, with X_t|1 distributed as N(t x_1, (1 − t)^2 I).
  • Training: The conditional and marginal objectives provide the same gradients for learning the neural velocity field.
  • Implementation: The standalone implementation includes a neural network, midpoint ODE solver, linear interpolation, regression loss, training loop, and trajectory visualization.

3 Flow models

Flow models use time-dependent diffeomorphic maps, defined by velocity fields through ODEs, to transform a source distribution into a target distribution. Their continuity-equation formulation also permits exact likelihood computation, while simulation-based training requires precise ODE solutions.

  • Flow models: Flow models transform samples from a source distribution p into target samples from q using a deterministic, time-continuous bijection ψt.The flow is a time-dependent diffeomorphism of R^d.
  • Probability paths: The flow reshapes the source density into a time-indexed probability path, with later states depending deterministically on the current state.This deterministic dependence also gives the flow model its Markov property.
  • Velocity fields and ODEs: A velocity field ut defines the flow ψt through the ODE dψt(x)/dt = ut(ψt(x)).Under regularity conditions, the ODE has a unique local flow solution.
  • Existence conditions: Local Lipschitz regularity guarantees only local existence; reaching t = 1 for all x requires stronger conditions such as global Lipschitzness or integrability.The manuscript uses integrability later to guarantee existence almost everywhere until t = 1.
  • Likelihood computation: Flows permit tractable exact likelihoods through the continuity equation's Instantaneous Change of Variables identity.Computing the divergence becomes increasingly challenging as dimensionality d grows, motivating unbiased estimators.
  • Training flow models: Training flow models by likelihood requires precise ODE simulations and therefore differs from Flow Matching's simulation-free training framework.Only errorless ODE solutions yield unbiased gradients for the simulation-based objective.

4 Flow Matching

Flow Matching trains a neural velocity field along a designed probability path from a known source to a data distribution. Conditioning makes path and velocity construction tractable, and Conditional Flow Matching preserves the learning objective while enabling scalable training.

  • 4 Flow Matching: Flow Matching learns a neural velocity field that transports a known source distribution p to a target distribution q along a probability path.Sampling from the learned model consists of solving an ODE with the learned velocity field.
  • 4 Flow Matching: The framework designs a path from p0 = p to p1 = q, trains uθ to match the ground-truth velocity, and samples by ODE integration.The regression loss uses a vector dissimilarity such as the squared ℓ2-norm.
  • Conditional construction: Conditional Flow Matching simplifies training by conditioning the objective on a single randomly selected target example.The marginal path is obtained by aggregating the conditional paths.
  • Conditional construction: The marginal velocity is an average of conditional velocity fields, equivalently interpretable as a posterior-weighted average over target examples.The weights represent posterior probabilities of targets given the current sample.
  • Theoretical guarantee: Under Assumption 1, conditionally integrable velocity fields generating conditional paths produce the corresponding marginal probability path.This result is stated as the Marginalization Trick for all t ∈ [0, 1).
  • Theoretical guarantee: Flow Matching and Conditional Flow Matching have identical gradients, and the Conditional Flow Matching minimizer is the marginal velocity field.Thus the conditional objective is equivalent for learning purposes.
  • Design choices: Sampling t according to ω can outperform weighting the loss by ω(t) in large-scale image-generation tasks.The passage reports this as an empirical performance observation.

4.6 Solving conditional generation with conditional flows

Conditional Flow Matching reduces the design of conditional probability paths and velocities to constructing a suitable conditional flow, then trains with a tractable conditional loss. The construction is valid under stated regularity and integrability conditions, while simpler interpolants may require additional guarantees.

  • Conditional-flow construction: Conditional Flow Matching designs conditional paths and velocities by constructing a conditional flow ψ_t(·|x_1) satisfying the required boundary conditions.The resulting conditional density and velocity are obtained from the flow and its time derivative.
  • Conditional-flow construction: The loss minimizer is the conditional expectation of the conditional velocity given the current state.This gives the marginal velocity its least-squares interpretation.
  • Conditional-flow construction: The conditional Flow Matching loss provides a tractable Monte Carlo objective for regressing the learnable velocity to conditional target velocities.The package exposes probability-path sampling and an example PyTorch training loop implementing the loss as a mean-squared error for Euclidean distance.
  • Validity conditions: Under the corollary’s smoothness, positivity, bounded-support, and integrability assumptions, the marginal velocity generates the marginal path interpolating source and target distributions.The result applies the conditional-flow construction to establish generation of the marginal probability path.
  • Validity conditions: Target-, source-, and two-sided conditioning choices are essentially equivalent when the conditional flow is a diffeomorphism in the relevant variables.Without those conditions, extra assumptions are needed; simple interpolation alone is not sufficient.
  • Optimal Transport and linear conditional flow: Optimal Transport promotes straight trajectories with constant velocity, which can make ODE sampling easier and allows a single Euler step for the target sample in the stated setting.The linear conditional flow minimizes a kinetic-energy bound among conditional flows, and the single-point-target case has an analytic OT solution.

4.8 Affine conditional flows

Affine conditional flows generalize linear conditional paths through scheduler functions controlling signal and noise. They support multiple trainable parameterizations, scheduler conversion, and theoretical equivalences, while endpoint singularities can remain practically problematic.

  • Affine-flow construction: Affine conditional flows use ψ_t(x_0|x_1) = α_t x_1 + σ_t x_0, with schedulers satisfying endpoint and monotonicity conditions.The derivative constraints make α_t strictly increasing and σ_t strictly decreasing.
  • Affine-flow construction: With bounded target support, a smooth positive source density, finite second moments, and independent coupling, the affine marginal velocity generates a path interpolating source and target distributions.Theorem 6 formally states this guarantee for the affine construction.
  • Parameterizations: The flow_matching library implements affine paths, schedulers, Conditional Optimal Transport paths, x_1-prediction training, and conversion back to velocity prediction.The examples include path sampling, Conditional Matching loss computation, and an ODE solver.
  • Parameterizations: Affine marginal velocities admit velocity, x_1-prediction, and x_0-prediction parameterizations that can be learned with matching losses and converted between one another.The conditional and matching losses have the same gradients for arbitrary functions of the coupled endpoint samples.
  • Practical caveats: Velocity-parameterization coefficients can have endpoint singularities that cancel theoretically but remain problematic when continuous models cannot perfectly regress their conditional targets.For Gaussian-source examples, endpoint velocity expressions can be used to correct these singularities.
  • Scheduler transformations: An affine velocity field can be transformed between schedulers using a closed-form scale-time transformation, enabling post-training scheduler changes.The text states that all schedulers theoretically produce the same sampling at t = 1, subject to the stated endpoint assumption.
  • Gaussian paths: For Gaussian paths, the marginal velocity is a gradient and therefore kinetic optimal for the fixed marginalized Gaussian probability path.This conclusion follows from the Gaussian-path velocity representation.

4.9 Data couplings

Data couplings determine how source and target samples are paired when training Flow Matching, enabling applications beyond independent noise-to-data generation and influencing trajectory geometry. Guidance extends the same framework to conditional target distributions, while classifier-free guidance remains widely used despite an unknown sampling distribution.

  • Data couplings: Flow Matching can use independent, paired, or data-dependent couplings between source and target distributions.Independent samples preserve marginals without dependencies, whereas paired data and learned couplings encode task-specific relationships.
  • Data couplings: Image in-painting uses dependent couplings to map masked images to compatible filled images, addressing an ill-defined inverse task.The source is obtained by masking a target image, with added noise supporting diversity and valid source densities.
  • Trajectory geometry: Straight probability paths yield smaller ODE simulation errors, motivating coupling designs that induce straighter trajectories.Optimal Transport couplings and multisample constructions connect transport cost with trajectory geometry.
  • Multisample couplings: For k > 1, multisample couplings reduce transport cost relative to independent couplings and approach Optimal Transport cost as k →∞ under quadratic cost.The method preserves marginals while introducing dependencies through an implicit joint distribution.
  • Guidance: Guided Flow Matching learns conditional distributions q(x1|y) using labeled target samples and a neural velocity field conditioned on the guidance variable.The guided probability path aggregates conditional paths while retaining the Flow Matching training blueprint.
  • Guidance: Classifier-free guidance is popular for conditional models, but the exact distribution it samples from is unknown.Guidance effectiveness is reported as greater when many target samples share a signal than when the signal is complex and non-repeating.

5 Non-Euclidean Flow Matching

Riemannian Flow Matching extends velocity-field generative modeling from Euclidean spaces to smooth manifolds, using tangent-space velocities and simulation-free training when suitable geometric maps are available. Geodesic and alternative-distance conditional flows provide constructions, but general manifolds can make geodesics difficult and introduce singularities or boundary bias.

  • 5 Non-Euclidean Flow Matching: Riemannian Flow Matching extends Flow Matching to general Riemannian manifolds for data represented on non-Euclidean spaces.The framework targets a simulation-free training algorithm, unlike earlier manifold flow training methods requiring expensive ODE simulations.
  • Geometric foundations: Manifold flows use time-dependent vector fields on tangent spaces, with probability densities and flows defined relative to the manifold geometry.The manifold metric supplies tangent-space inner products and volume-based probability integration.
  • Riemannian Flow Matching: The Riemannian Flow Matching loss compares tangent-space velocities through a Bregman divergence while enforcing source and target marginal constraints.Conditional manifold paths are averaged to obtain the marginal velocity field.
  • Riemannian Flow Matching: Theorem 11 states that the gradients of the Riemannian Flow Matching and Riemannian Conditional Flow Matching losses coincide.This provides the manifold analogue of the conditional-training simplification used in Euclidean Flow Matching.
  • Geodesic conditional flows: Geodesic conditional flows use exponential and logarithmic maps with a monotone scheduler κ(t) satisfying κ(0)=0 and κ(1)=1.The construction pushes each starting point toward its conditioning endpoint along a manifold geodesic.
  • Premetrics and limitations: Closed-form exponential and logarithmic maps enable simulation-free manifold training, whereas general manifolds may make geodesics difficult to compute or bias probability toward boundaries.Alternative smooth premetrics can address computational difficulty or introduce desired implicit bias.
  • Premetrics and limitations: Geodesic and premetric conditional flows can have singularities, such as nondifferentiability at antipodal points on the 2-sphere.The problematic set is generally small and usually has zero volume, so the paper reports no practical problems in known use cases.

6 Continuous Time Markov Chain Models

Continuous Time Markov Chains provide a discrete-state counterpart to flow models by prescribing transition rates between states and evolving probability mass through the Kolmogorov equation. Their mass-conservation theory supports constructing paths and adding divergence-free velocities without changing marginals.

  • 6 Continuous Time Markov Chain Models: CTMCs model time-continuous Markov processes on finite discrete state spaces and underpin Discrete Flow Matching.They provide a generative-model alternative to continuous flows for token-like or other discrete data.
  • CTMC definition: A CTMC transition kernel is specified infinitesimally by rates ut(y,x), which determine probability movement from state x to state y.The kernel combines the identity transition with a first-order rate term and an o(h) remainder.
  • CTMC definition: Rate conditions are required to keep small-time transition probabilities nonnegative and normalized.Violating them can make probabilities negative or cause their sum to differ from one.
  • Marginal dynamics: CTMC marginal probabilities satisfy the Kolmogorov equation, the discrete analogue of the continuity equation.The equation expresses marginal evolution through probability flux and its divergence.
  • Existence and uniqueness: Continuous rates guarantee a unique Kolmogorov-equation solution on t∈[0,1) from the initial distribution.The theorem requires rates continuous in time and imposes no additional conditions unlike the nonlinear flow case.
  • Mass conservation: Discrete Mass Conservation characterizes when rates generate a prescribed probability path, provided the rates satisfy the CTMC conditions.The result is the principal construction tool for CTMC probability paths and velocities.
  • Mass conservation: Divergence-free velocities can be added during sampling without changing the marginal probability.The added velocities must satisfy the rate conditions and the divergence-free velocity equation.

7 Discrete Flow Matching

Discrete Flow Matching extends the Flow Matching recipe from continuous distributions to CTMCs on discrete state spaces. It uses learnable velocities and Bregman-divergence losses, with factorization making sequence-scale models tractable.

  • Discrete Flow Matching framework: DFM interpolates source and target probability mass functions with a CTMC whose learnable velocity generates the resulting probability path.Training minimizes a Bregman divergence defining the Discrete Flow Matching loss.
  • Discrete Flow Matching framework: Source and target states may use independent or general couplings, including paired translation data.The coupling specifies how source and target samples are related before constructing the probability path.
  • Marginalization and losses: Under positivity and regularity assumptions, conditional velocities can be marginalized into a velocity that generates the marginal discrete path.The discrete Marginalization Trick transfers the conditional construction to the marginal CTMC.
  • Marginalization and losses: The Discrete Flow Matching and Conditional Discrete Flow Matching losses have identical gradients, and the conditional-loss minimizer is the marginal velocity.This equivalence enables training with conditional quantities while targeting the marginal model.
  • Factorized velocities: Factorized velocities connect states differing in at most one token, reducing the model output dimension from K^d to d · K.The factorized model predicts per-coordinate rates and supports coordinate-wise CTMC sampling.
  • Implementation: The guide supplies standalone pure-PyTorch and library implementations of discrete flow matching, including arbitrary data coupling.These examples support practical adoption of the discrete framework.

8 Continuous Time Markov Process Models

This section generalizes Flow Matching to arbitrary state spaces and continuous-time Markov processes. It introduces generators as learnable first-order descriptions of transitions and develops Generator Matching across process classes.

  • General framework: The unified framework targets general state spaces and Markov processes, covering Euclidean, discrete, manifold, multimodal, flow, diffusion, and CTMC settings.The state space is treated abstractly, with examples spanning several data modalities and process types.
  • General state spaces: General state spaces use a reference measure ν, with densities represented by Radon–Nikodym derivatives; PDFs and PMFs are special cases.Lebesgue measure yields densities on R^d, while counting measure yields discrete probability masses.
  • General state spaces: The framework also accommodates distributions without densities, including point-mass paths and distributions over trajectory spaces.Such cases are represented directly as probability measures p(dx).
  • Continuous-time Markov processes: A CTMP is memoryless: conditioned on the present, its future is independent of the past and is specified by a transition kernel and the initial distribution.The transition kernel gives the conditional law of X_t+h given X_t.
  • Generators: The generator is the first-order time derivative of the transition kernel and provides a simpler object to parameterize than the full kernel.The informal model uses p_t+h|t(·|x) := δ_x + hL_t(x) + o(h).
  • Formal generator definition: Because transition kernels act on distributions rather than vector spaces, test functions are used to define derivatives and characterize distributions through expectations.The formal development imposes Feller-process regularity, including strong continuity and no return from infinity.
  • Generator theory: The section states general mass-conservation and universal generator-characterization theorems under regularity assumptions.These results support a common generator-based formulation for CTMP generative models.

9 Generator Matching

Generator Matching (GM) extends Flow Matching to arbitrary state spaces and general Markov processes by learning generators that produce probability paths from source to target distributions. Its marginalization and loss results provide a unified, scalable training framework, with constructions for flows, diffusions, jumps, model combinations, and multimodal generation.

  • Framework: GM unifies generative models over arbitrary data modalities and general Markov processes, including diffusion, discrete diffusion, and Flow Matching variants.It constructs CTMP generative models through process generators.
  • Probability paths: GM constructs a marginal probability path by sampling a latent state and then sampling from a conditional path, while preserving p0 = p and p1 = q.The conditional construction generalizes the path-design recipe used in Flow Matching.
  • Generator parameterization: Linear generator parameterizations learn only a function F_t while keeping the operator K fixed, covering flow, diffusion, and jump generators.For flows, setting Kf = ∇f and F_t = u_t recovers parameterization through vector fields.
  • Marginalization: The General Marginalization Trick supplies the generator for a marginal path and turns its approximation into a neural-network training target.Earlier marginalization results are special cases of this theorem.
  • Training objective: The Generator Matching and Conditional Generator Matching losses have identical gradients, and the conditional-loss minimizer is the marginal generator parameterization.The conditional objective is tractable when conditional generators can be derived analytically.
  • Extensions: GM supports combining flow, diffusion, and jump dynamics and constructing multimodal models by combining modality-specific models through factorized conditional probability paths.The framework also permits adding divergence-free dynamics without changing the prescribed probability path.

10 Relation to Diffusion and other Denoising Models

The paper frames diffusion and related denoising models as particular Flow Matching constructions, distinguished by their probability paths, time parameterizations, and stochastic or deterministic sampling processes. Under shared conditions, diffusion training and sampling can be recovered within the FM framework, while time-reversal and probability-path generation are equivalent formulations.

  • Diffusion as Flow Matching: Diffusion models construct FM-compatible probability paths through forward noising SDEs whose closed-form marginals support score-based parameterization.The forward process is initialized from data, and its marginal approaches a Gaussian only asymptotically rather than at finite diffusion time.
  • Diffusion as Flow Matching: Diffusion training is equivalent to FM training with an affine Gaussian probability path, independent coupling, and score parameterization after reparameterizing time.The score and x0-prediction parameterizations are theoretically equivalent, although both introduce a singularity near t = 0.
  • Sampling Correspondence: Under a Gaussian source, independent coupling, fixed αt and σt, and score parameterization, probability-flow ODE sampling equals FM sampling, while stochastic SDE sampling equals the corresponding GM sampling.These equivalences hold under the same specified construction conditions.
  • Time-Reversal and Generalization: Finding a backward process with forward marginals, generating the forward probability path, and solving the corresponding Fokker–Planck or continuity equation are equivalent problems.This reframes denoising and time-reversal procedures as instances of the broader Generator Matching and Flow Matching objectives.
  • Sampling Correspondence: The probability-flow ODE follows the same marginals as the diffusion process without being its time-reversal, and it is used for low-evaluation sampling.The paper states that an exact time-reversal can be harder to solve and may be suboptimal for generative modeling.
  • Time-Reversal and Generalization: FM principles extend beyond Euclidean deterministic flows to continuous and discrete Markov processes, manifolds, and denoising models inspired by diffusion.Discrete Flow Matching applies to continuous-time Markov chains and supports discrete generative tasks such as language modeling.

A.1 Discrete Mass Conservation

The appendix establishes when discrete-state dynamics preserve probability mass and when a rate field generates a valid probability-mass path. It proves equivalence between the Kolmogorov equation with rate conditions and the existence of a compatible transition kernel.

  • Mass Conservation: A solution of the discrete Kolmogorov equation remains a probability mass function when initialized as one and when transition rates satisfy the stated continuity and rate conditions.The proof establishes both nonnegativity and preservation of total mass.
  • Mass Conservation: Theorem 13 characterizes discrete mass conservation by equivalence between the Kolmogorov equation with rate conditions and a valid PMF evolution.The result applies to continuous-in-time rates and PMFs on discrete state spaces.
  • Transition Kernels: The proof rules out negative off-diagonal rates and nonzero total outgoing-rate sums because either would violate transition-kernel normalization or nonnegativity.Both contradictions arise from the short-time expansion of the transition kernel.
  • Transition Kernels: A compatible transition kernel can be recovered from a Kolmogorov-equation solution, and its semigroup property follows from uniqueness and composition of intermediate kernels.The construction uses the law of total probability and uniqueness of Kolmogorov-equation solutions.

A.2 Manifold Marginalization Trick

The manifold marginalization trick shows that conditional velocity fields generating conditional probability paths can be integrated into a marginal velocity field generating the marginal path. The proof relies on regularity, integrability, and positivity conditions that justify exchanging differentiation and integration.

  • Statement: Under Assumption 2, a conditionally integrable velocity field that generates each conditional probability path also generates the corresponding marginal probability path.This is the manifold analogue of aggregating conditional dynamics into marginal dynamics.
  • Proof: The proof obtains the marginal continuity equation by differentiating the integrated conditional path and substituting the conditional continuity equation.Leibniz-rule conditions justify exchanging differentiation with integration and divergence operations.
  • Proof: Strict positivity of the marginal density permits the marginal velocity formula, while bounded support or compactness and local Lipschitz arguments provide the needed regularity.These conditions ensure the resulting marginal velocity is integrable and locally Lipschitz.

A.3 Regularity assumptions for KFE

The framework assumes regularity conditions ensuring uniqueness and well-posedness of the Kolmogorov equation. Such results are available in several important settings, but no general theorem covers arbitrary state spaces and Markov processes.

  • Scope of the Assumptions: Assumption 5 is considered relatively weak, but its regularity requirements are stated as an assumption because general-state-space results are unavailable.The paper notes that practitioners’ state spaces are expected to satisfy the assumption.
  • Known Settings: Uniqueness results are cited for flows on Euclidean spaces and manifolds, diffusion on Euclidean spaces and manifolds, and general Itô SDEs in Euclidean spaces.These references delimit settings where Kolmogorov-equation well-posedness has established support.
  • Known Settings: For discrete state spaces, the Kolmogorov equation is a linear ODE with a unique solution when its coefficients are continuous.This gives a direct regularity guarantee for the discrete setting.
Loading 2412.06264v1…