PaperlayerView exact revision on arXiv

Source-linked AI summary

Multisample Flow Matching: Straightening Flows with Minibatch Couplings

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, Ricky T. Q. Chen

arXiv:2304.14772v2cs.LG

TL;DR

Existing Flow Matching methods independently pair noise and data samples, limiting probability-path design. Multisample Flow Matching uses nontrivial couplings between minibatch data and noise samples while preserving correct marginals and retaining simulation-free training. On ImageNet, it reduces sampling cost by 30% to 60% for low FID versus baseline Flow Matching with only 4% more training time, without degrading log-likelihood or sample quality.

  • Problem

    Independent data-noise sampling limits probability-path design and produces nonzero gradient variance even at convergence.

  • Method

    Multisample Flow Matching trains with nonindependent couplings of k-sample empirical distributions using a simulation-free joint Flow Matching objective that requires only sampling from the joint distribution.

  • Results

    30% to 60% lower sampling cost reaches low FID on ImageNet than baseline Flow Matching, with 4% more training time and no degradation in log-likelihood or sample quality.

  • Takeaways & Limitations

    The framework yields straighter flows and improved sample efficiency while preserving correct marginal distributions, with applications to generative modeling and transport maps.

  • Takeaways & Limitations

    The gradient-variance proxy is always smaller than actual gradient variance because it omits variability over x and t, while the transport theorem analyzes only asymptotic properties.

Abstract

from arXiv · show

Simulation-free methods for training continuous-time generative models construct probability paths that go between noise distributions and individual data samples. Recent works, such as Flow Matching, derived paths that are optimal for each data sample. However, these algorithms rely on independent data and noise samples, and do not exploit underlying structure in the data distribution for constructing probability paths. We propose Multisample Flow Matching, a more general framework that uses non-trivial couplings between data and noise samples while satisfying the correct marginal constraints. At very small overhead costs, this generalization allows us to (i) reduce gradient variance during training, (ii) obtain straighter flows for the learned vector field, which allows us to generate high-quality samples using fewer function evaluations, and (iii) obtain transport maps with lower cost in high dimensions, which has applications beyond generative modeling. Importantly, we do so in a completely simulation-free manner with a simple minimization objective. We show that our proposed methods improve sample consistency on downsampled ImageNet data sets, and lead to better low-cost sample generation.

1. Introduction

Existing Flow Matching and diffusion approaches independently pair noise and data samples, limiting probability-path design. Multisample Flow Matching introduces nonindependent minibatch couplings, improving flow straightness and low-cost sampling while preserving marginal distributions.

  • Independent endpoint sampling causes nonzero converged gradient variance, slower training, and restricted probability-path design.
  • Multisample Flow Matching generalizes Flow Matching by coupling k-sample empirical noise and data distributions nonindependently.
  • As k →∞, an appropriate OT-inspired coupling makes sample paths straight, enabling more efficient simulation.
  • 30% to 60% lower sampling cost achieves low FID than baseline Flow Matching on ImageNet, with only 4% more training time.
  • The method regularizes toward the optimal vector field simulation-free, avoids adversarial formulations, and preserves correct marginal distributions.

2. Preliminaries

The preliminaries formulate continuous normalizing flows as vector fields transporting one density to another and review Flow Matching's conditional-path construction. They also connect straight sample paths to optimal transport and simpler simulation.

  • A Continuous Normalizing Flow models a neural vector field whose flow transports an initial density p0 to a target density p1.
  • Flow Matching trains a parametric vector field by regressing toward a target field that generates a probability path with prescribed endpoint marginals.
  • Conditional Flow Matching makes this objective tractable by defining conditional paths whose marginalization recovers the desired probability path and vector field.
  • With standard Gaussian noise, the optimal-transport conditional construction transports each x0 to x1 at t = 1 along a straight path.
  • Optimal transport minimizes a coupling cost while enforcing endpoint marginals, and its dynamic formulation yields straight paths that are easier to simulate.

3. Flow Matching with Joint Distributions

The Joint CFM framework replaces independent endpoint sampling with arbitrary joint distributions that preserve the correct marginals, enabling variance reduction, straighter flows, and near-optimal transport maps. Its objective remains simulation-free and reduces to standard Conditional Flow Matching under independent coupling.

  • Joint CFM construction: Joint CFM generalizes Conditional Flow Matching to arbitrary joint distributions q(x0, x1) satisfying the correct marginal constraints.The construction samples from the joint distribution without requiring explicit endpoint or conditional-path densities.
  • Joint CFM construction: The Joint CFM objective uses samples from q(x0, x1) and reduces to the original CFM objective when q(x0, x1) = q0(x0)q1(x1).The optimal learned vector field still maps between the marginal distributions q0 and q1.
  • Variance reduction: Choosing a suitable joint distribution can reduce gradient variance during training, whereas independent sampling generally cannot achieve zero proxy variance even at the optimum.The average gradient variance at fixed x and t is upper bounded in terms of the Joint CFM objective, which serves as a proxy for total gradient variance.
  • Straight flows: Straightness measures whether marginal trajectories remain constant-velocity paths, and straighter trajectories can be simulated accurately with fewer function evaluations.Independent endpoint sampling generally produces nonzero straightness, while certain joint distributions yield straightness close to zero.
  • Near-optimal transport cost: The marginal flow defines a transport map from q0 to q1, and well-chosen joint distributions can make its transport cost approach the optimal value.This gives Joint CFM a simulation-free use for approximating optimal transport maps in high-dimensional settings beyond generative modeling.

4. Multisample Flow Matching

Multisample Flow Matching constructs joint data–noise distributions from minibatch couplings while preserving the correct marginals. BatchOT and approximations target lower variance, straighter flows, and near-optimal transport at manageable computational cost.

  • Multisample construction: The multisample construction samples batches and uses a doubly-stochastic coupling to define a joint distribution with the correct data and noise marginals.Different couplings induce different joint distributions while preserving the marginal constraints.
  • Multisample construction: The construction includes the independent Flow Matching joint distribution as the uniform coupling π(i, j) = 1/k.This provides a natural reference point for comparing nonindependent couplings.
  • Batch Optimal Transport: BatchOT solves minibatch optimal transport exactly using standard solvers, producing a permutation coupling with runtime complexity O(k^3).The resulting permutation matrix is doubly stochastic and can be incorporated into the training procedure.
  • Batch Optimal Transport: As k →∞, BatchOT makes the optimal Joint CFM objective and flow straightness converge to zero while its transport cost converges to the optimal transport cost.These limits imply reduced gradient variance, easier simulation, and simulation-free approximation of optimal transport maps.
  • Batch Entropic OT: BatchEOT uses entropic regularization and Sinkhorn’s algorithm, with runtime ˜O(k^2/ε), interpolating between BatchOT and independent coupling as ε changes.As ε →0 it recovers BatchOT; as ε →∞ it recovers the independent coupling.
  • Empirical comparison: Figure 2 reports Joint CFM objective values of 0.57 for BatchEOT and 0.24 for BatchOT, versus 10.72 for CondOT.The objective provides an upper bound on the variance of u_t at convergence.
  • Stable and Heuristic Couplings: The stable and heuristic couplings provide lower-cost alternatives, with the heuristic coupling modifying stable matching to discourage increases in matched-pair cost.The ranking-based stable approach costs O(k^2 log(k)) when coupling costs must be sorted.

5. Related Work

The related work contrasts Multisample Flow Matching with prior optimal-transport and generative-modeling approaches. Its distinction is preserving marginal distributions while using minibatch couplings for simulation-free continuous-time modeling.

  • Optimal transport and generative modeling: Optimal transport has been used in generative modeling to stabilize training, regularize continuous normalizing flows, and improve interpretability.These uses connect transport objectives with easier simulation and broader continuous-time modeling goals.
  • Empirical comparison: Figure 3 compares FID with Euler-discretization compute cost measured by NFE, reporting significantly higher low-NFE FID for CondOT than proposed methods.The comparison focuses on the sample-quality versus computation trade-off.
  • Minibatch couplings for generative modeling: Prior minibatch optimal-transport methods can fit the barycentric map rather than correctly matching the true marginal distribution at finite batch sizes.The paper attributes this limitation to static models that parameterize the map directly with a neural network.
  • Minibatch couplings for generative modeling: Multisample Flow Matching preserves marginal constraints in a dynamic neural-ODE setting while using batch couplings within a simulation-free framework.The paper contrasts this construction with prior approaches that lack finite-batch marginal preservation.
  • Empirical comparison: Table 1 indicates that diffusion-based baselines such as ScoreFlow and DDPM require more than 40 NFE to reach the listed FID values.The table derives approximate NFE requirements from Figure 3.

6. Experiments

Experiments show that Multisample Flow Matching improves convergence, low-cost sample quality, consistency, and transport costs while preserving marginal distributions. The method achieves these gains with small runtime overhead and across image, synthetic, and 2D settings.

  • Image Datasets: Multisample Flow Matching retains Flow Matching performance while improving sample quality, compute cost, and variance without observed performance degradation.Runtime increases are 0.8% on ImageNet32 and 4% on ImageNet64.
  • Image Datasets: Multisample Flow Matching generally improves sample quality at fixed, lower numbers of function evaluations than the naive CondOT approach.The experiments use fixed-step solvers and compare FID as sampling cost decreases.
  • Image Datasets: BatchOT produces more consistent individual samples across low- and high-NFE solutions than CondOT despite similar FID values.Consistency is evaluated with a perceptual metric based on pretrained InceptionNet hidden units.
  • Image Datasets: BatchOT coupling yields faster training, consistent with reduced gradient variance from non-trivial joint distributions.Figure 4 compares convergence against Flow Matching with CondOT and diffusion-based methods.
  • Improved Batch Optimal Couplings: Multisample Flow Matching preserves marginal distributions while reducing transport cost relative to the BatchOT oracle at fixed batch sizes.It also converges toward the optimal-transport solution faster than the BatchOT oracle.

7. Conclusion

The conclusion presents Multisample Flow Matching as a simulation-free extension of Flow Matching that uses batch couplings while preserving marginal constraints. It supports efficient sampling and transport-map learning without simulation during training or min-max optimization.

  • 7. Conclusion: Multisample Flow Matching uses complex joint distributions instead of independently sampled data and noise to design probability paths.The framework builds on simulation-free training for continuous normalizing flows.
  • 7. Conclusion: The method improves sample efficiency and quality with low-cost solvers without simulation during training or min-max formulations.The conclusion also identifies marginal preservation as important for generative modeling and transport problems.
  • 7. Conclusion: Batch coupling algorithms construct an implicit joint distribution satisfying marginal constraints, with BatchOT approximating optimal-transport maps.Stable and Heuristic couplings provide lower-complexity alternatives with desired coupling properties.
  • 7. Conclusion: Stable coupling is cheaper than optimal transport but is agnostic to pairing costs and uses sample rankings.The Heuristic coupling modifies reassignment using cost information inspired by cyclic monotonicity.

B.1. Full results on ImageNet data

Appendix results summarize Multisample Flow Matching’s comparison with Flow Matching across sample quality, efficiency, and variance metrics.

  • B.1. Full results on ImageNet data: Multisample Flow Matching improves sample quality and sample efficiency without trading off performance against Flow Matching.Table 6 reports FID, bits per dimension, number of function evaluations, and variance.

B.2. How batch size affects the marginal probability paths on 2D checkerboard data

The appendix compares coupling behavior across batch sizes, sampling costs, and runtimes for ImageNet and 2D marginal paths. Larger coupling sizes improve convergence stability, while the figures organize the relevant efficiency-quality trade-offs.

  • B.2. How batch size affects the marginal probability paths on 2D checkerboard data: Marginal probability paths are shown for batch sizes 64 and 8.The comparison isolates how batch size affects the paths on 2D checkerboard data.
  • B.2. How batch size affects the marginal probability paths on 2D checkerboard data: Sample quality is evaluated as FID versus compute cost measured by NFE using midpoint discretization.The ImageNet32 and ImageNet64 tables compare FID across sampling budgets.
  • B.2. How batch size affects the marginal probability paths on 2D checkerboard data: Runtime comparisons report iterations per second and relative increase with respect to CondOT.The comparison includes CondOT, BatchOT, and Stable matching at standard batch sizes.
  • B.2. How batch size affects the marginal probability paths on 2D checkerboard data: Larger coupling sizes k produce faster and more stable convergence while keeping the training batch size fixed.Only k varies when solving the couplings in the 64-dimensional experiments.

C. Generated samples

The appendix establishes that minibatch-coupled marginal paths preserve the endpoint distributions and converge toward straight, optimal-transport behavior under stated assumptions. The accompanying ImageNet figures compare sample consistency across substantially different function-evaluation budgets.

  • C. Generated samples: Batch optimal couplings produce more consistent ImageNet32 samples across 200, 12, 8, and 6 NFEs from the same noise sample.The samples use midpoint discretization.
  • C. Generated samples: Batch optimal couplings produce more consistent ImageNet64 samples across 200, 12, 8, and 6 NFEs from the same noise sample.The samples use midpoint discretization.
  • D. Theory: The constructed marginal path interpolates between q0 and q1, with the endpoint conditional distribution at t = 1 concentrated at x1.The conditional flow transports points to x1 at time 1.
  • D. Theory: The marginal vector field transports between p0 = q0 and p1 = q1, while the Joint CFM objective shares Flow Matching’s expected gradient and is uniquely minimized by the marginal field.This connects the coupled objective to the original Flow Matching objective without changing the target marginal transport.
  • D.4. Convergence: Under assumptions (A1)–(A3), the minibatch coupling converges weakly to the population optimal transport plan, and the straightness converges to zero as k →∞.For BatchOT, assumption (A3) holds almost surely; the theorem separately states limk→∞Sk = 0.

D.5. Bounds on the transport cost and monotone convergence results

The transport-cost analysis bounds the cost of the marginal flow using the chosen joint coupling and shows monotone convergence behavior for BatchOT. These results provide a weak guarantee that increasing minibatch size does not substantially worsen transport cost.

  • D.5. Bounds: For any joint distribution q(x0, x1), the marginal flow’s transport cost can be upper-bounded using a quantity determined by q.The bound follows from the shared interpolation distribution and Jensen’s inequality.
  • D.5. Monotone convergence: For BatchOT, the upper-bound quantity is monotonically decreasing in minibatch size k.This monotonicity combines with the general transport-cost bound to support the stated convergence guarantee.
  • D.5. Monotone convergence: The BatchOT analysis provides a weak guarantee that transport cost should not get much higher as k increases.The conclusion combines Proposition D.7 with Theorem D.8.

E.1. Image datasets

The ImageNet experiments use controlled architectures and optimizer settings across downsampled datasets, with distinct precision and ODE-solver configurations for ImageNet-32 and ImageNet-64. Coupling algorithms are restricted to individual GPUs.

  • E.1. Image datasets: ImageNet-32 training uses full 32-bit precision, whereas ImageNet-64 uses 16-bit mixed precision.Both use the same Adam parameter settings and matched training epochs across methods.
  • E.1. Image datasets: ImageNet experiments use the Dhariwal–Nichol architecture with substantially lower attention resolution.All trained methods use identical architectures and corresponding parameters for the same number of epochs.
  • E.1. Image datasets: Sampling solves for t = 1 with dopri5 at atol = rtol = 1e−5, using K = 15 for ImageNet32 and K = 10 for ImageNet64.FID is computed with the TensorFlow-GAN library.
  • E.1. Image datasets: Coupling algorithms run only within each GPU because cross-GPU experiments showed no noticeable sample-efficiency gains and slightly worse performance and sample quality.The selected setup therefore uses smaller per-GPU batches.
  • E.1. Image datasets: Synthetic experiments use three datasets in dimensions 2, 32, and 64, each consisting of 50K samples from Gaussian mixtures.The number of mixture centers is specified in Table 11.
Loading 2304.14772v2…