PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow Matching for Generative Modeling

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le

arXiv:2210.02747v2cs.LGcs.AIstat.ML

TL;DR

Diffusion-based training confines probability paths and can require long training and specialized sampling, while CNF training lacks tractable target fields. Flow Matching addresses this with simulation-free regression on conditional probability paths, and ImageNet experiments show favorable likelihood, sample quality, and computational trade-offs versus diffusion-based methods.

  • Problem

    Diffusion processes restrict probability paths, contributing to long training times and specialized sampling, while target CNF vector fields are generally unavailable in closed form.

  • Method

    Flow Matching trains CNFs by regressing neural vector fields onto tractable conditional probability paths, including Gaussian paths that subsume diffusion paths.

  • Results

    ImageNet experiments achieve favorable likelihood, sample quality, and computational-cost/sample-quality trade-offs among competing diffusion-based methods.

  • Takeaways & Limitations

    Flow Matching enables scalable CNF training with directly specified probability paths, supporting easier training and sampling and paths designed for faster generation.

  • Takeaways & Limitations

    Diffusion-derived paths do not reach a true noise distribution in finite time, so sampling and likelihood evaluation approximate p0 with a suitable Gaussian.

Abstract

from arXiv · show

We introduce a new paradigm for generative modeling built on Continuous Normalizing Flows (CNFs), allowing us to train CNFs at unprecedented scale. Specifically, we present the notion of Flow Matching (FM), a simulation-free approach for training CNFs based on regressing vector fields of fixed conditional probability paths. Flow Matching is compatible with a general family of Gaussian probability paths for transforming between noise and data samples -- which subsumes existing diffusion paths as specific instances. Interestingly, we find that employing FM with diffusion paths results in a more robust and stable alternative for training diffusion models. Furthermore, Flow Matching opens the door to training CNFs with other, non-diffusion probability paths. An instance of particular interest is using Optimal Transport (OT) displacement interpolation to define the conditional probability paths. These paths are more efficient than diffusion paths, provide faster training and sampling, and result in better generalization. Training CNFs using Flow Matching on ImageNet leads to consistently better performance than alternative diffusion-based methods in terms of both likelihood and sample quality, and allows fast and reliable sample generation using off-the-shelf numerical ODE solvers.

1 INTRODUCTION

The paper introduces Flow Matching as an efficient, simulation-free method for training Continuous Normalizing Flows beyond diffusion processes. It develops conditional training objectives and validates them with Optimal Transport paths on ImageNet, achieving favorable likelihood, sample-quality, and computational trade-offs.

  • Motivation: Diffusion models enable scalable and relatively stable generative-model training, but restricting sampling paths to simple diffusion processes confines the available path space.This motivates seeking more general probability paths for generative modeling.
  • Motivation: Continuous Normalizing Flows can model arbitrary probability paths, including diffusion paths, but scalable training beyond diffusion remains unavailable.Maximum likelihood requires expensive numerical ODE simulations, while existing simulation-free approaches involve intractable integrals or biased methods.
  • Contribution: Flow Matching is an efficient simulation-free approach that trains CNFs using general probability paths without requiring diffusion-process reasoning.The method is designed to enable scalable CNF training beyond diffusion.
  • Contribution: The Flow Matching objective regresses a target vector field generating a desired probability path, while Conditional Flow Matching uses per-example formulations with equivalent gradients.The conditional objective avoids requiring explicit knowledge of the target marginal vector field.
  • Experiments: On ImageNet, Flow Matching with Optimal Transport paths achieves favorable likelihood and sample quality while providing better computational-cost and sample-quality trade-offs than prior diffusion-based methods.The experiments use a large and highly diverse image dataset.

2 PRELIMINARIES: CONTINUOUS NORMALIZING FLOWS

The preliminaries define probability density paths and time-dependent vector fields, whose ODE-generated flows form the basis of Continuous Normalizing Flows. CNFs use neural-network vector fields to transform a simple prior density into a more complex target density, with the continuity equation characterizing valid paths.

  • Core objects: A probability density path p_t is a time-dependent density on R^d, paired with a time-dependent vector field v_t.The density integrates to one at each time, and the vector field maps [0, 1] × R^d to R^d.
  • Core objects: A vector field generates a time-dependent diffeomorphic flow φ_t through an ordinary differential equation.The flow is the mechanism connecting the vector field to evolving probability densities.
  • Continuous Normalizing Flows: Continuous Normalizing Flows parameterize v_t with a neural network and use the resulting flow to reshape a simple prior p_0 into a more complex density p_1.The prior may be pure noise, and the transformation is expressed through a push-forward equation.
  • Probability paths: The continuity equation provides a practical test for whether a vector field generates a probability density path.It is also a key component of the paper’s proofs.

3 FLOW MATCHING

Flow Matching trains continuous normalizing flows by regressing vector fields that generate a chosen probability path from simple noise toward the data distribution. Conditional Flow Matching makes this objective tractable using per-sample paths and vector fields while preserving the original objective’s optimization target.

  • Flow Matching objective: Flow Matching regresses a target vector field u_t with a learnable CNF vector field v_t along a probability path from simple noise toward data.At zero loss, the learned CNF generates the target probability path.
  • Flow Matching objective: Naively computing Flow Matching is intractable because suitable marginal paths and closed-form marginal vector fields are generally unavailable.The method addresses this by constructing paths and vector fields defined per sample.
  • Conditional probability paths: Conditional paths start from the noise distribution and end in distributions concentrated around individual data samples, whose marginal approximates the data distribution.A representative endpoint is p_1(x|x_1) = N(x|x_1, σ^2I) with sufficiently small σ > 0.
  • Conditional probability paths: Marginalizing conditional vector fields produces the correct marginal vector field, which generates the corresponding marginal probability path.This connection decomposes the unknown marginal field into simpler fields depending on single data samples.
  • Conditional Flow Matching: Conditional Flow Matching enables unbiased training estimates by sampling per-sample paths and vector fields instead of evaluating intractable marginal integrals.It requires efficient sampling from p_t(x|x_1) and computation of u_t(x|x_1).
  • Conditional Flow Matching: The FM and CFM objectives have identical gradients with respect to θ, so optimizing CFM is equivalent in expectation to optimizing FM.This trains a CNF toward a marginal path that approximates the unknown data distribution at t = 1 without access to the marginal path or vector field.

4 CONDITIONAL PROBABILITY PATHS AND VECTOR FIELDS

This section defines Conditional Flow Matching over general Gaussian conditional probability paths and derives their canonical vector fields. It recovers diffusion paths while introducing Optimal Transport paths with simpler, straight-line dynamics and a more stable Flow Matching training alternative.

  • General Gaussian paths: Conditional Flow Matching supports any differentiable Gaussian path defined by time-dependent mean µt(x1) and standard deviation σt(x1), subject to shared noise and data boundary conditions.All paths begin at standard Gaussian noise, with µ0(x1)=0 and σ0(x1)=1, and end at a concentrated Gaussian around x1 with µ1(x1)=x1 and sufficiently small σmin.
  • General Gaussian paths: The canonical affine flow map generates each Gaussian conditional path, yielding a unique vector field while avoiding unnecessary distribution-invariant components.The map is ψt(x)=σt(x1)x+µt(x1), and its induced vector field generates pt(x|x1).
  • Diffusion conditional VFs: Diffusion conditional vector fields are recovered as special Gaussian paths, and Flow Matching with them is more stable and robust than existing score-matching approaches.The construction coincides with the deterministic probability-flow vector field when restricted to conditional diffusion processes.
  • Optimal Transport conditional VFs: Optimal Transport paths linearly vary the mean and standard deviation, producing an OT displacement interpolant between the endpoint Gaussians.This construction provides direct control over the path rather than relying on diffusion processes that do not reach true noise in finite time.
  • Optimal Transport conditional VFs: OT trajectories move in straight lines at constant speed, while their conditional vector field has constant direction in time and is simpler to fit than diffusion targets.Diffusion trajectories can overshoot and backtrack, whereas OT paths are guaranteed to stay straight.

5 RELATED WORK

Related work established Continuous Normalizing Flows as continuous-time normalizing flows but highlighted the computational cost of ODE-based maximum-likelihood training. Subsequent simulation-free approaches designed probability paths or diffusion processes, though existing methods faced estimation or gradient limitations that Conditional Flow Matching generalizes.

  • Continuous Normalizing Flows: Continuous Normalizing Flows extend normalizing flows to continuous time, but maximum-likelihood training requires expensive forward and backward ODE simulations.The sequential nature of ODE simulations leads to high time complexity.
  • Simulation-free CNF training: Simulation-free CNF methods explicitly design probability paths and dynamics, but prior approaches faced difficult high-dimensional integrals or biased stochastic-minibatch gradients.Rozen et al. use linear interpolation with difficult-to-estimate integrals, while Ben-Hamu et al. use general probability paths with biased gradients.
  • Diffusion-based training: Diffusion-based simulation-free training indirectly defines target probability paths, and denoising score matching provides unbiased gradients for score matching.Conditional Flow Matching draws inspiration from this conditional-objective result and generalizes it.

6 EXPERIMENTS

Experiments on CIFAR-10 and ImageNet show that Flow Matching, especially with Optimal Transport paths, improves training efficiency, sample quality, and sampling cost over diffusion-based alternatives. OT paths also generate images earlier and achieve strong ImageNet-128 FID performance, while supporting conditional upsampling with similar PSNR and SSIM to prior work.

  • Image quality: On ImageNet-128, Flow Matching with the OT path achieves state-of-the-art FID except for IC-GAN, which uses conditioning with a self-supervised ResNet50 model.IC-GAN is excluded from the comparison table for that reason.
  • Faster training: FM-OT lowers ImageNet-64 FID faster and further than the alternative diffusion-based methods during training.Flow Matching generally converges faster despite diffusion works reporting training runs of 1.3m and 10m iterations.
  • Sample paths: OT paths begin generating recognizable images earlier than diffusion paths, which remain noise-dominated until near the end of sampling.This pattern appears in ImageNet-64 sample paths generated from identical random seeds and similarly in a 2D checkerboard example.
  • Low-cost samples: Roughly 60% of the NFEs are sufficient for FM-OT to reach the same numerical-error threshold as diffusion models in low-cost ImageNet-32 sampling.The comparison uses fixed-step solvers, at most 100 NFEs, 256 random noise seeds, and 1000-NFE solutions as references.
  • Conditional image generation: For 64×64-to-256×256 conditional upsampling, FM-OT achieves PSNR and SSIM values similar to Saharia et al. (2022).The evaluation compares against the original validation-set reference and a regression baseline using FID on upsampled validation images.

7 CONCLUSION

The paper introduces Flow Matching as a simulation-free framework for scaling Continuous Normalizing Flow training to high dimensions, while enabling direct specification of probability paths for faster sampling and improved generation.

  • Conclusion: Flow Matching trains Continuous Normalizing Flows simulation-free through conditional constructions that scale to very high dimensions.The framework is introduced as a new approach to CNF training.
  • Conclusion: Flow Matching reframes diffusion modeling by allowing probability paths to be specified directly rather than relying on stochastic diffusion constructions.This alternative view supports designing paths with targeted generation or sampling properties.
  • Conclusion: Directly specified probability paths can enable faster sampling and improve generation.The conclusion identifies these as practical benefits of moving beyond stochastic diffusion constructions.

8 SOCIAL RESPONSIBILITY

Image generation offers positive applications but also creates risks of harmful use and rising energy demand. The passage emphasizes content controls, image validation, and more update-efficient training as mitigation strategies.

  • Risks: Image generation’s positive applications are accompanied by potential harmful uses.The passage identifies harmful use as a social responsibility concern without specifying particular applications.
  • Mitigations: Content-controlled training sets and image validation or classification can help reduce harmful uses.These measures target both the data used for training and the resulting images.
  • Efficiency: More efficient training with fewer gradient updates or lower image throughput can reduce time and energy consumption.The passage links rapidly increasing training energy demand to the importance of update- and throughput-efficient methods.

A THEOREM PROOFS · B THE CONTINUITY EQUATION · C COMPUTING PROBABILITIES OF THE CNF MODEL

The paper establishes that Flow Matching vector fields generate their marginal probability paths and that, under regularity and positivity conditions, conditional and marginal objectives have identical gradients. It then explains how CNF probabilities and BPD can be computed through ODE integration, Hutchinson trace estimation, and data transformations.

  • A THEOREM PROOFS: Theorem 1 shows that the marginal vector field generates the marginal probability path by satisfying the continuity equation.The proof exchanges integration and differentiation under Leibniz-rule regularity conditions.
  • A THEOREM PROOFS: Under pt(x) > 0, conditional and marginal Flow Matching objectives differ only by a θ-independent constant and have equal gradients.The derivation assumes sufficient decay for q(x) and pt(x|x1), plus bounded ut, vt, and ∇θvt to justify the required integrals and Fubini exchanges.
  • A THEOREM PROOFS: For Gaussian probability paths, the unique vector field defining the corresponding flow map is the conditional field ut(x|x1), which therefore generates the Gaussian path.Invertibility follows from σt(x1) > 0.
  • B THE CONTINUITY EQUATION: The continuity equation is a necessary and sufficient PDE test for whether a vector field generates a probability path.Its divergence is taken with respect to the spatial variable x ∈ Rd.
  • C COMPUTING PROBABILITIES OF THE CNF MODEL: CNF log probabilities are computed jointly with flow trajectories by integrating the instantaneous change-of-variables ODE, solving it in reverse for a given data point.ODE uniqueness identifies the reverse solution and recovers log p1(x1) from the initial prior probability and integrated divergence.
  • C COMPUTING PROBABILITIES OF THE CNF MODEL: The Hutchinson trace estimator replaces costly divergence computation in Rd and yields an unbiased estimator of log p1(x1).The estimator uses z ∈ Rd sampled so that E[zzT] = I, with the resulting ODE solved under controlled numerical error.
  • C COMPUTING PROBABILITIES OF THE CNF MODEL: For transformed data, the model is composed with the data transformation, and the transformed prior pushforward determines the resulting density.For images, pixel values are mapped from [−1, 1] to [0, 256].
  • C COMPUTING PROBABILITIES OF THE CNF MODEL: Bits-Per-Dimension is estimated by averaging the unbiased log-probability estimator over a large transformed test set.The estimator is applied to ϕ−1(x1) before averaging to approximate test-set BPD.

D DIFFUSION CONDITIONAL VECTOR FIELDS

This section derives conditional vector fields for Fokker–Planck probability paths and instantiates the construction for VE and VP diffusion paths. After reversing time to match the noise-to-data convention, the resulting fields coincide with the corresponding diffusion Probability Flow ODE fields.

  • D DIFFUSION CONDITIONAL VECTOR FIELDS: For VE and VP diffusion paths, the derived conditional vector fields coincide with the corresponding fields obtained from Theorem 3.The VE and VP paths are instantiated separately, with time reversal applied because diffusion conventionally runs from data to noise.
  • D DIFFUSION CONDITIONAL VECTOR FIELDS: The derived fields for VE and VP paths also coincide with the vector fields governing the Probability Flow ODE.The correspondence is stated for the VE and VP diffusion paths and the fields defined in the cited equations.
  • D DIFFUSION CONDITIONAL VECTOR FIELDS: Time reversal maps a vector field u_t(x) to ˜u_t(x) = −u_1−t(x), reversing the probability path from data-to-noise into the noise-to-data convention.The reversed field generates ˜p_t(x) = p_1−t(x).
  • Conditional VFs for Fokker-Planck probability paths: The section derives conditional vector fields by rewriting the Fokker–Planck equation as a continuity equation for the associated probability path.The construction begins from an SDE and its Fokker–Planck equation, then expresses the evolution in continuity-equation form.

E IMPLEMENTATION DETAILS · E.1 DIFFUSION BASELINES

The implementation uses shared architectures, preprocessing, hyperparameters, and training duration across FM-OT, FM-Diffusion, and SM-Diffusion, while evaluating three standard diffusion-loss parametrizations and a VP diffusion path. Sampling uses ODE-based score-matching generation or the DDPM update rule with a score-to-noise conversion.

  • E IMPLEMENTATION DETAILS: FM-OT, FM-Diffusion, and SM-Diffusion use the same architecture, hyperparameters, and number of training epochs.The 2D model is a 5-layer MLP with 512 neurons per layer, while image models use the UNet architecture from Dhariwal & Nichol (2021).
  • E IMPLEMENTATION DETAILS: Image preprocessing uses center cropping and resizing, with the Chrabaszcz et al. (2017) preprocessing retained at 32×32 and 64×64 resolutions.
  • E.1 DIFFUSION BASELINES: Three diffusion-baseline losses represent popular parametrizations from Song & Ermon (2019), Song et al. (2021), Ho et al. (2020), and Kingma et al. (2021).The losses are formulated under the general Gaussian path form of equation 10.
  • E.1 DIFFUSION BASELINES: The original Score Matching loss corresponds to Song & Ermon (2019), while λ(t) = β(1 − t) gives the Score Flow loss motivated by an NLL upper bound.The learnable score function is denoted s_t.
  • E.1 DIFFUSION BASELINES: The DDPM Noise Matching loss from Ho et al. (2020) is specified alongside a standard-Gaussian p0(x) = N(x|0, I) and learnable noise function ϵ_t.
  • E.1 DIFFUSION BASELINES: The diffusion path is the standard VP diffusion, with β(s) = βmin + s(βmax − βmin).The implementation sets βmin = 0.1, βmax = 20, and samples time in [0, 1−ϵ] with ϵ = 10−5 for training, likelihood, and sampling.
  • E.1 DIFFUSION BASELINES: Score-matching samples are produced by solving the ODE in equation 1 with the specified vector field.
  • E.1 DIFFUSION BASELINES: DDPM samples use equation 46 after setting s_t(x) = ϵ_t(x)/σ_t, where σ_t = q.

E.2 TRAINING & EVALUATION DETAILS · F ADDITIONAL TABLES AND FIGURES

The paper specifies standardized training, likelihood, solver, and evaluation procedures, then provides additional visual results for FM-OT generation across ImageNet resolutions and sampling budgets.

  • E.2 TRAINING & EVALUATION DETAILS: Training uses full 32-bit precision on CIFAR10 and ImageNet-32, versus 16-bit mixed precision on ImageNet-64/128/256.All trained methods use identical architectures, parameters, and epoch counts.
  • E.2 TRAINING & EVALUATION DETAILS: All trained methods use Adam with β1 = 0.9, β2 = 0.999, weight decay = 0.0, and ϵ = 1e−8.The methods are FM-OT, FM-Diffusion, and SM-Diffusion.
  • E.2 TRAINING & EVALUATION DETAILS: Negative log-likelihood is reported after standard uniform dequantization using an importance-weighted estimate.The associated results are presented in bits per dimension for different K values.
  • E.2 TRAINING & EVALUATION DETAILS: Likelihood computation solves at t = 1 with adaptive-step dopri5 using atol = rtol = 1e-5 through torchdiffeq.Estimated values for different K are reported in Table 4.
  • E.2 TRAINING & EVALUATION DETAILS: CIFAR10 and ImageNet-32/64 FID and Inception scores use TensorFlow GAN library 2, while ImageNet-128 uses Dhariwal and Nichol’s evaluation script.The ImageNet-128 script is used to preserve comparability with Dhariwal & Nichol (2021).
  • F ADDITIONAL TABLES AND FIGURES: Additional figures show non-curated unconditional FM-OT samples for ImageNet-32, ImageNet-64, and ImageNet-128.These results are documented in Figures 11–13.
  • F ADDITIONAL TABLES AND FIGURES: Further figures present conditional 64×64→256×256 FM-OT upsampling and samples generated from fixed initial noise at varying function-evaluation counts.The varying-NFE examples use ImageNet-128 and ImageNet 256×256 models.
Loading 2210.02747v2…