PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving and generalizing flow-based generative models with minibatch optimal transport

Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, Yoshua Bengio

arXiv:2302.00482v4cs.LG

TL;DR

CNF training has been limited by simulation-based maximum likelihood and by assumptions such as Gaussian source distributions, while ODE/SDE integration can make inference slow. The paper introduces simulation-free CFM and its OT-CFM variant, showing broader applicability, more efficient flows, and approximation of dynamic OT and Schrödinger bridge problems.

  • Problem

    CNF and diffusion-style generative models face simulation or integration burdens, while earlier flow-matching objectives imposed restrictive source-distribution assumptions.

  • Method

    The paper develops generalized conditional flow matching and OT-CFM, using regression to conditional vector fields and optimal-transport couplings without integrating trajectories during training.

  • Results

    Across single-cell dynamics, image generation, image translation, energy-based models, and Schrödinger bridge inference, CFM methods improve flow-model training or inference, with OT-CFM producing more efficient and accurate OT flows.

  • Takeaways & Limitations

    When transport plans are sampleable, CFM provides simulation-free approximations to dynamic OT and Schrödinger bridge probability flows for arbitrary source and target distributions.

  • Takeaways & Limitations

    CFM requires closed-form conditional flows, and minibatch OT can introduce approximation limitations; the experiments also work in latent rather than image space.

Abstract

from arXiv · show

Continuous normalizing flows (CNFs) are an attractive generative modeling technique, but they have been held back by limitations in their simulation-based maximum likelihood training. We introduce the generalized conditional flow matching (CFM) technique, a family of simulation-free training objectives for CNFs. CFM features a stable regression objective like that used to train the stochastic flow in diffusion models but enjoys the efficient inference of deterministic flow models. In contrast to both diffusion models and prior CNF training algorithms, CFM does not require the source distribution to be Gaussian or require evaluation of its density. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference, as evaluated in our experiments. Furthermore, we show that when the true OT plan is available, our OT-CFM method approximates dynamic OT. Training CNFs with CFM improves results on a variety of conditional and unconditional generation tasks, such as inferring single cell dynamics, unsupervised image translation, and Schrödinger bridge inference.

1 Introduction

The paper introduces CFM as a unified, simulation-free way to train CNFs with arbitrary transport maps, and OT-CFM as a variant that improves flow efficiency and dynamic-OT approximation. Experiments cover multiple generation and inference tasks.

  • Conditional flow matching: CFM unifies flow-matching methods and diffusion-style approaches for training CNFs with arbitrary transport maps.The framework generalizes existing FM and diffusion modeling approaches.
  • Optimal transport CFM: OT-CFM uses optimal-transport couplings to approximate dynamic OT with simulation-free CNF training.The method uses static OT maps between marginal distributions and can solve dynamic OT and Schrödinger bridge problems.
  • Efficiency: OT-CFM improves training efficiency and reduces inference time by producing straighter flows that require fewer network evaluations.The approach targets the long inference times caused by repeatedly integrating neural ODEs or SDEs.
  • Experiments: The authors evaluate CFM and OT-CFM on single-cell dynamics, image generation, unsupervised image translation, energy-based models, and Schrödinger bridge problems.The paper also reports improved and reproducible training practices for high-dimensional image generation.
  • Resources: The paper releases torchcfm, a Python package unifying new and existing flow-training algorithms under a shared interface.The package includes implementations of the main experiments.

2 Background: Optimal transport and neural ODEs

The paper frames CNFs as ODE-defined maps between distributions and motivates simulation-free vector-field regression to avoid costly ODE simulation. It introduces static and dynamic optimal transport as related formulations for constructing efficient transport flows.

  • Neural ODEs: Generative modeling fits and samples from a distribution by learning a mapping from a source distribution q0 to a target distribution q1.The source may be Gaussian or an empirical distribution represented by finite samples.
  • Neural ODEs: Continuous normalizing flows represent the mapping between source and data distributions with a neural ordinary differential equation.The ODE integration map transports samples along a time-varying vector field.
  • Probability flows: A probability flow ODE is a time-varying vector field whose transported density follows a marginal probability path governed by the continuity equation.The integration map pushes the initial density forward to the density at each time.
  • Simulation-free training: Flow matching regresses a neural vector field to a known vector field when the probability path and tractable samples are available.The objective avoids directly simulating the ODE during training.
  • Optimal transport: Static optimal transport minimizes displacement cost over couplings with fixed marginals, whereas dynamic optimal transport optimizes vector fields transforming one measure into another.Under stated assumptions, the static and dynamic 2-Wasserstein formulations are equivalent.
  • Motivation: Prior neural OT models can require tens to hundreds of function evaluations with backpropagation through them, motivating direct simulation-free vector-field regression.The paper also connects optimal transport with Schrödinger bridges, whose probability flow is recovered by a proposed variant.

3 Conditional flow matching: ODEs from static couplings

Conditional flow matching recovers marginal vector fields by regressing on tractable conditional paths, allowing arbitrary source distributions and extensions based on optimal-transport couplings. Its OT and entropic variants target dynamic OT and Schrödinger-bridge probability flows without simulation.

  • Conditional flow matching: Conditional flow matching estimates an intractable marginal vector field by regressing on sampled conditional paths and tractable conditional vector fields.The stochastic objective is unbiased under the stated sampling and computability conditions.
  • Conditional flow matching: Under positivity of the marginal path, the CFM and flow-matching objectives differ only by a θ-independent constant.Thus, minimizing CFM recovers the same learned marginal vector field as the corresponding flow-matching objective.
  • I-CFM: I-CFM uses independent source and target samples, relaxing the Gaussian-source requirement and supporting arbitrary source distributions with intractable densities.As σ →0, its marginal vector field approaches transport from q(x0) to q(x1).
  • OT-CFM: OT-CFM replaces the independent coupling with a joint coupling sampled from a 2-Wasserstein optimal transport plan.This choice is intended to produce an approximate OT flow.
  • OT-CFM: As σ2 →0, the OT-CFM marginal path and vector field minimize the dynamic OT objective under regularity assumptions.When minibatch size equals the support size, the method recovers exact OT and learns exact dynamic OT; smaller batches can still perform well.
  • SB-CFM: SB-CFM uses an entropy-regularized OT coupling and Brownian-bridge conditional paths to train an ODE matching a Schrödinger-bridge probability flow.The resulting marginal vector field generates the same marginal probability path as the Schrödinger-bridge solution, with OT-CFM and I-CFM recovered at limiting regularization values.

4 Related work

Related methods include diffusion models, simulation-free flow objectives, and neural dynamic-OT approaches. The paper positions CFM as a framework that broadens source distributions and avoids architectural constraints or simulation-heavy optimization.

  • Diffusion models: Diffusion models achieve strong generative performance but typically require costly SDE simulation for inference and often use simple Gaussian source processes.General-source diffusion methods can require multiple iterations or additional techniques for effective optimization and inference.
  • Evaluation: Table 2 compares neural OT methods across four distribution pairs using 2-Wasserstein fit, normalized path energy, and runtime.A dash denotes methods requiring a Gaussian source, while training hardware differs between CFM/RF models and other baselines.
  • Flow matching: Prior simulation-free CNF methods correspond to CFM with Gaussian sources or independent samples from q0 and q1.The paper extends this setting to more general conditional couplings and Schrödinger bridges.
  • Dynamic optimal transport: Existing neural dynamic-OT methods use constrained architectures or regularized CNFs that are challenging to optimize.The paper claims OT flows can be achieved without either constraint.

5 Experiments

Experiments evaluate CFM variants and prior methods on optimal transport, generative modeling, single-cell interpolation, Schrödinger bridges, image generation, and translation. OT-CFM generally improves training speed, inference efficiency, and task performance, while its minibatch OT approximation remains effective.

  • Low-dimensional optimal transport: OT-CFM flows generalize better to test data and remain close to dynamic optimal-transport paths by normalized path energy.Small batches were often sufficient to approximate the optimal transport map well.
  • Training and inference efficiency: OT-CFM achieves better validation performance in the same number of training steps than CFM and FM objectives.The authors attribute easier optimization partly to variance reduction in the conditional flow.
  • Training and inference efficiency: OT-CFM requires fewer function evaluations to reach the same sample quality and achieves higher quality at the same evaluation budget.This behavior is consistent with straighter optimal-transport paths.
  • Schrödinger bridges: SB-CFM reproduces Schrödinger bridge flows while being significantly faster than the diffusion Schrödinger bridge method.The comparison uses average 2-Wasserstein distance to ground-truth bridge samples over 18 time steps.
  • Single-cell interpolation: OT-CFM outperforms other methods and baselines on average across all three single-cell datasets.Evaluation uses earth mover’s distance on left-out intermediate timepoints.
  • High-dimensional image generation: On CIFAR-10, OT-CFM outperforms I-CFM and FM across training-step and function-evaluation comparisons.The figure reports FID with dopri5 during training and with Euler integration after 400k training steps.
  • Experimental scope: The authors could not reproduce prior reported results from Lipman et al. using the parameters specified in that paper.They therefore selected different training hyperparameters for the high-dimensional image experiments.
  • Unsupervised translation: The experiments also evaluate CFM for mapping between unpaired high-dimensional datasets after encoding images into 128-dimensional latent vectors.Held-out translated and target embeddings are compared using maximum mean discrepancy.

6 Conclusion

The paper introduces simulation-free conditional flow-matching objectives for continuous-time flows with general source distributions. Optimal-transport conditioning supports dynamic OT and Schrödinger bridge solutions while improving training and inference efficiency.

  • Contribution: CFM provides simulation-free objectives for learning continuous-time flows with a general source distribution.Training does not require integration over time.
  • Contribution: OT-CFM lifts static optimal transport to dynamic OT and Schrödinger bridge problems using simulation-free training.The approach uses static transport information between marginal distributions.
  • Implications: CFM improves training and inference efficiency by lowering objective variance and straightening learned flows.The conclusion identifies these as mechanisms associated with more efficient flow models.
  • Limitations: CFM requires closed-form conditional flows, and minibatch OT can incur error in high dimensions.The authors identify neural-network approximations to OT maps as a direction for addressing the latter limitation.

Contribution statement

The contribution statement assigns conceptual, theoretical, experimental, coding, writing, and project-guidance roles across the authors.

  • Roles: A.T. conceived the idea, while Y.Z., G.H., and N.M. led theoretical development.The statement separately credits high-dimensional experiments and open-source code to A.T. and K.F.
  • Roles: All authors contributed to experiment design, with writing driven by N.M. and A.T.G.W. and Y.B. guided the project.

A Proofs of theorems

The proofs establish that the constructed marginal vector fields generate the intended probability paths and recover dynamic OT or Schrödinger bridge flows under stated assumptions.

  • General flow construction: The marginal vector field in Theorem 3.1 generates the probability path from the initial conditions.The proof verifies the continuity equation for the constructed path and field.
  • Flow-matching equivalence: Under positivity and regularity assumptions, conditional flow-matching and flow-matching objectives have equal gradients or equal objectives up to a parameter-independent constant.Theorem 3.2 assumes sufficient decay and boundedness for the required integrals and derivatives.
  • Conditional paths: For independent endpoint sampling, the conditional paths have Gaussian-smoothed endpoint marginals.Proposition 3.3 identifies the boundary conditions as q1 ∗ N(x | 0, σ^2) and q0 ∗ N(x | 0, σ^2).
  • Dynamic optimal transport: With an optimal transport plan and regularity conditions, the zero-noise limit of the marginal path and field solves dynamic optimal transport.The result requires σ^2 → 0 and assumptions on q0, q1, and the transport plan.
  • Dynamic optimal transport: Brenier’s theorem yields a Monge map whose interpolating paths do not cross and whose associated vector field solves the optimal transport problem.The map is parameterized as the gradient of a convex function.
  • Schrödinger bridges: The vector field defined for the entropic construction generates the same marginal probability path as the Schrödinger bridge solution.The proof connects the construction to Brownian bridges with diffusion scale σ pinned at the endpoints.

B Additional theoretical results

Theoretical results establish equivalence and convergence properties for conditional flow matching, while showing that OT-CFM can reduce regression-target variance and minibatch averaging offers a variance–computation tradeoff.

  • For any σ ∈ R+, conditional flow matching has an equivalent marginal probability flow to Lipman et al.’s flow matching.
  • CFM with the specified conditional-path variance is equivalent to FM with variance (tσ − t + 1)^2.
  • If π is a Monge map, OT-CFM’s objective variance converges to zero as σ → 0.
  • As minibatch size approaches the population size, the conditional vector field converges to the marginal vector field.
  • Averaging over m datapoints reduces target variance but takes roughly m times longer to compute the conditional target.

C.2 Modeling energy functions

The paper extends CFM to settings where endpoint distributions are represented by unnormalized energy functions rather than samples, with a gradient-equivalent training objective and sampling-related limitations.

  • Energy Conditional Flow Matching adapts CFM when endpoint distributions are specified by unnormalized energy functions instead of samples.
  • The energy-based objective can train a flow matching the endpoint energies without access to samples.
  • Under positivity assumptions, the gradients of LFM and LECFM with respect to θ are equal up to a multiplicative constant.
  • Uniform sampling does not scale well to high-dimensional spaces, motivating alternative sampling strategies for practical use.
  • The toy funnel experiment uses a uniform proposal distribution for reweighted importance sampling.

D Additional results

Additional experiments examine entropy-regularized and minibatch Schrödinger bridge objectives, regularized CNFs, and the effects of batch size and bandwidth on fit, path length, and convergence.

  • Minibatch Schrödinger bridge CFM uses Sinkhorn to construct an entropy-regularized transport plan before sampling conditional paths.
  • λe = 0.1 provides a cross-dataset tradeoff between near-optimal path length and good data fit for regularized CNFs.
  • At minibatch size one, OT-CFM is equivalent to CFM; normalized path energy decreases rapidly and plateaus near batch size 64.
  • OT-CFM maintains near-optimal path energies and fit until σ > 1, whereas FM and CFM paths are 2–3x longer than optimal.
  • Diffusion Schrödinger Bridge fits early trajectories better, but fails to fit the bridge after many integration steps.

D.1 Objective variance.

OT-CFM and related variance-reduction strategies lower stochastic objective variance, accelerating convergence and improving efficiency, while energy-based experiments show advantages under fixed integration budgets.

  • OT-CFM remains the fastest-converging method on two additional datasets.
  • OT-CFM and SB-CFM have at least an order of magnitude lower objective variance than CFM and FM across datasets.
  • Lower objective variance correlates with faster training and lower validation error in fewer steps.
  • Batch averaging reaches lower validation error in fewer steps and less walltime, especially for objectives with high variance.
  • MCMC-based target sampling can miss target-density modes when the chain does not mix well.
  • With adaptive integration, OT-CFM requires fewer function evaluations; under a fixed computation budget, it performs significantly better.

E.8 Unconditional CIFAR-10 experiments

The CIFAR-10 experiments held the training setup fixed while varying the probability path, with separate implementation details for reproduced and improved models. An additional figure illustrates OT-CNF image-to-image translation trajectories in CelebA latent space.

  • Experimental setup: All CIFAR-10 methods used the same training setup, differing only in their choice of probability path.The authors note that unreleased code may cause some parameter differences from the reproduced setup.
  • Reproduced configuration: The reproduced configuration used Adam with β1 = 0.9, β2 = 0.999, ϵ = 10−8, and no weight decay.
  • Reproduced configuration: The reproduced models used a UNet with 256 channels, depth 2, four attention heads, batch size 128 per GPU, two GPUs, and 2000 epochs.The learning rate increased linearly for the first 200 epochs and then decayed linearly.
  • Improved configuration: The improved models changed the setup to 128 channels, dropout 0.1, one A100 GPU, 400000 steps, σ = 0, and a constant learning rate of 2 × 10−4.They also used gradient clipping with norm 1.0 and exponential moving-average weights with decay 0.9999.
  • Latent-space translation: The CelebA figure shows OT-CNF trajectories translating latent encodings between negative and positive facial attributes while preserving image structure.Each trajectory starts from a reconstructed negative image encoding and ends at x1 after ten intermediate decoded states.
Loading 2302.00482v4…