Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Weighted Conditional Flow Matching

Sergio Calvo-Ordonez, Matthieu Meunier, Alvaro Cartea, Christoph Reisinger, Yarin Gal, Jose Miguel Hernandez-Lobato

TL;DR

We address the inefficiency and suboptimal path quality of standard conditional flow matching by introducing Weighted Conditional Flow Matching (W-CFM), which weights training pairs with a Gibbs kernel to approximate entropic optimal transport (EOT) without explicit OT computations. The method yields a loss $L_{\rm W-CFM}(\theta;\varepsilon)$ that corresponds to an EOT-driven prior, recovers the EOT coupling up to marginal tilt, and becomes equivalent to OT-CFM in the large-batch limit when marginals remain essentially unchanged. Theoretical results on marginal tilting, choices of $\varepsilon$, and large-batch equivalence are complemented by extensive experiments on toy transports and unconditional image generation, where W-CFM matches or surpasses baselines in sample quality, path straightness, and diversity while maintaining vanilla CFM efficiency. This offers a scalable, practically effective alternative to minibatch OT for training continuous normalizing flows with high-quality, straight trajectories.

Abstract

Conditional flow matching (CFM) has emerged as a powerful framework for training continuous normalizing flows due to its computational efficiency and effectiveness. However, standard CFM often produces paths that deviate significantly from straight-line interpolations between prior and target distributions, making generation slower and less accurate due to the need for fine discretization at inference. Recent methods enhance CFM performance by inducing shorter and straighter trajectories but typically rely on computationally expensive mini-batch optimal transport (OT). Drawing insights from entropic optimal transport (EOT), we propose Weighted Conditional Flow Matching (W-CFM), a novel approach that modifies the classical CFM loss by weighting each training pair $(x, y)$ with a Gibbs kernel. We show that this weighting recovers the entropic OT coupling up to some bias in the marginals, and we provide the conditions under which the marginals remain nearly unchanged. Moreover, we establish an equivalence between W-CFM and the minibatch OT method in the large-batch limit, showing how our method overcomes computational and performance bottlenecks linked to batch size. Empirically, we test our method on unconditional generation on various synthetic and real datasets, confirming that W-CFM achieves comparable or superior sample quality, fidelity, and diversity to other alternative baselines while maintaining the computational efficiency of vanilla CFM.

Weighted Conditional Flow Matching

TL;DR

We address the inefficiency and suboptimal path quality of standard conditional flow matching by introducing Weighted Conditional Flow Matching (W-CFM), which weights training pairs with a Gibbs kernel to approximate entropic optimal transport (EOT) without explicit OT computations. The method yields a loss that corresponds to an EOT-driven prior, recovers the EOT coupling up to marginal tilt, and becomes equivalent to OT-CFM in the large-batch limit when marginals remain essentially unchanged. Theoretical results on marginal tilting, choices of , and large-batch equivalence are complemented by extensive experiments on toy transports and unconditional image generation, where W-CFM matches or surpasses baselines in sample quality, path straightness, and diversity while maintaining vanilla CFM efficiency. This offers a scalable, practically effective alternative to minibatch OT for training continuous normalizing flows with high-quality, straight trajectories.

Abstract

Conditional flow matching (CFM) has emerged as a powerful framework for training continuous normalizing flows due to its computational efficiency and effectiveness. However, standard CFM often produces paths that deviate significantly from straight-line interpolations between prior and target distributions, making generation slower and less accurate due to the need for fine discretization at inference. Recent methods enhance CFM performance by inducing shorter and straighter trajectories but typically rely on computationally expensive mini-batch optimal transport (OT). Drawing insights from entropic optimal transport (EOT), we propose Weighted Conditional Flow Matching (W-CFM), a novel approach that modifies the classical CFM loss by weighting each training pair with a Gibbs kernel. We show that this weighting recovers the entropic OT coupling up to some bias in the marginals, and we provide the conditions under which the marginals remain nearly unchanged. Moreover, we establish an equivalence between W-CFM and the minibatch OT method in the large-batch limit, showing how our method overcomes computational and performance bottlenecks linked to batch size. Empirically, we test our method on unconditional generation on various synthetic and real datasets, confirming that W-CFM achieves comparable or superior sample quality, fidelity, and diversity to other alternative baselines while maintaining the computational efficiency of vanilla CFM.

Paper Structure

This paper contains 32 sections, 5 theorems, 33 equations, 10 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Entropic Optimal Transport
  4. Conditional Flow Matching
  5. Weighted Conditional Flow Matching
  6. Approximating EOT with Weighted Conditional Flow Matching
  7. Marginal Tilting under W-CFM
  8. On the Choice of ${\varepsilon}$
  9. Equivalence to OT-CFM in the Large Batch Limit
  10. Related Work
  11. Straightening Sample Paths with Optimal Transport.
  12. Other Approaches for Faster Inference.
  13. Biasing Flow Matching via Weighting
  14. Experiments
  15. Experimental Setup
  16. ...and 17 more sections

Key Result

Theorem 1

If $c(x,y) < \infty$$\mu \otimes \nu$-almost-surely, for any ${\varepsilon} > 0$ there exist measurable functions $\phi_{\varepsilon},\psi_{\varepsilon}: \mathbb{R}^d \rightarrow \mathbb{R}$, referred to as EOT potentials, such that the EOT plan is given by

Figures (10)

  • Figure 1: Sample trajectories for $\text{circular MoG} \to 5 \text{ Gaussians}$. Left to right, the models used are trained with: I-CFM, W-CFM (${\varepsilon} = 0.4$), W-CFM (${\varepsilon} = 0.2$), OT-CFM (batch size 16), and OT-CFM.
  • Figure 2: Sample trajectories for moons generation. Source samples are in blue, generated samples are in red. From left to right, we use: I-CFM, W-CFM (${\varepsilon} = 10$), W-CFM (${\varepsilon} = 2$), and OT-CFM. Here, we use a variant of W-CFM where $\hat{f}_{\varepsilon} = \hat{f}_{{\varepsilon}, \mathrm{MC}}, \hat{g}_{\varepsilon} = \hat{g}_{{\varepsilon}, \mathrm{MC}}$ (see Appendix \ref{['app:mc-estimates']}).
  • Figure 3: Contour plots of learned density for moons (using 50,000 generated samples). The leftmost plot corresponds to the true target distribution. Then, from left to right, the models used are trained with: I-CFM, W-CFM (${\varepsilon} = 10$), W-CFM (${\varepsilon} = 2$), and OT-CFM.
  • Figure 4: Sample trajectories on $\text{8 Gaussians} \to \text{moons}$ with different variants of W-CFM. From left to right, the models used are trained with the following values of ${\varepsilon}$: 10,8,6,4,2.
  • Figure 5: Contour plots of learned target density for $\text{8 Gaussians} \to \text{moons}$. The leftmost plot corresponds to the true target distribution. Then, from left to right, the models used are trained with the following values of ${\varepsilon}$: 10,8,6,4,2.
  • ...and 5 more figures

Theorems & Definitions (5)

  • Theorem 1: Theorem 4.2 in nutz2021introduction
  • Proposition 1
  • Proposition 2: Marginal tilting and continuity equation
  • Proposition 3
  • Proposition 4