Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On the Relation between Rectified Flows and Optimal Transport

Johannes Hertrich, Antonin Chambolle, Julie Delon

TL;DR

The paper analyzes the relationship between rectified flows, flow matching, and optimal transport. It demonstrates that iterative rectification does not generally yield OT maps and presents counterexamples with disconnected supports and non-rectifiable couplings where zero loss does not imply optimality. It derives affine invariances and explicit Gaussian/GMM velocity fields, and studies gradient vs potential velocity formulations, revealing that gradient constraints are not a reliable route to OT maps. To mitigate fragility, it introduces a smoothed rectification approach and provides numerical evidence suggesting improved convergence toward OT under smoothing, while highlighting remaining theoretical limitations and open questions.

Abstract

This paper investigates the connections between rectified flows, flow matching, and optimal transport. Flow matching is a recent approach to learning generative models by estimating velocity fields that guide transformations from a source to a target distribution. Rectified flow matching aims to straighten the learned transport paths, yielding more direct flows between distributions. Our first contribution is a set of invariance properties of rectified flows and explicit velocity fields. In addition, we also provide explicit constructions and analysis in the Gaussian (not necessarily independent) and Gaussian mixture settings and study the relation to optimal transport. Our second contribution addresses recent claims suggesting that rectified flows, when constrained such that the learned velocity field is a gradient, can yield (asymptotically) solutions to optimal transport problems. We study the existence of solutions for this problem and demonstrate that they only relate to optimal transport under assumptions that are significantly stronger than those previously acknowledged. In particular, we present several counterexamples that invalidate earlier equivalence results in the literature, and we argue that enforcing a gradient constraint on rectified flows is, in general, not a reliable method for computing optimal transport maps.

On the Relation between Rectified Flows and Optimal Transport

TL;DR

The paper analyzes the relationship between rectified flows, flow matching, and optimal transport. It demonstrates that iterative rectification does not generally yield OT maps and presents counterexamples with disconnected supports and non-rectifiable couplings where zero loss does not imply optimality. It derives affine invariances and explicit Gaussian/GMM velocity fields, and studies gradient vs potential velocity formulations, revealing that gradient constraints are not a reliable route to OT maps. To mitigate fragility, it introduces a smoothed rectification approach and provides numerical evidence suggesting improved convergence toward OT under smoothing, while highlighting remaining theoretical limitations and open questions.

Abstract

This paper investigates the connections between rectified flows, flow matching, and optimal transport. Flow matching is a recent approach to learning generative models by estimating velocity fields that guide transformations from a source to a target distribution. Rectified flow matching aims to straighten the learned transport paths, yielding more direct flows between distributions. Our first contribution is a set of invariance properties of rectified flows and explicit velocity fields. In addition, we also provide explicit constructions and analysis in the Gaussian (not necessarily independent) and Gaussian mixture settings and study the relation to optimal transport. Our second contribution addresses recent claims suggesting that rectified flows, when constrained such that the learned velocity field is a gradient, can yield (asymptotically) solutions to optimal transport problems. We study the existence of solutions for this problem and demonstrate that they only relate to optimal transport under assumptions that are significantly stronger than those previously acknowledged. In particular, we present several counterexamples that invalidate earlier equivalence results in the literature, and we argue that enforcing a gradient constraint on rectified flows is, in general, not a reliable method for computing optimal transport maps.

Paper Structure

This paper contains 26 sections, 14 theorems, 66 equations, 17 figures, 1 table.

Table of Contents

  1. Introduction
  2. Outline and Contributions
  3. Related Work
  4. Rectified Flows
  5. Backgrounds
  6. Wasserstein Distance
  7. Rectified Flows
  8. Iterative Rectification
  9. Affine Invariance and Gaussian Case
  10. Rectified Flows and Optimal Transport
  11. Backgrounds on Rectified Flows with Gradient Fields
  12. Velocity Fields with Potential
  13. Counterexamples
  14. Disconnected Supports
  15. Non-Rectifiable Couplings
  16. ...and 11 more sections

Key Result

Theorem 2

Let $(X_0,X_1)$, be a coupling between $\mu_0$ and $\mu_1$, let $v_t=\mathop{\mathrm{arg\,min}}\limits_{w_t}\mathcal{L}(w_t|X_0,X_1)$ be the minimizer of the loss function eq:flow_matching_loss and let $A\in\mathbb{R}^{d\times d}$ be invertible, $b\in\mathbb{R}^d$ and $c\in\mathbb{R}_{>0}$. Then, th If $(X_0,X_1)$ is in addition rectifiable with $(Z_0,Z_1)=\mathcal{R}(X_0,X_1)$, then it holds that

Figures (17)

  • Figure 1: Construction of non-optimal couplings which are fixed points of $\mathcal{R}_p$.
  • Figure 2: Illustration of the non-rectifiable coupling $(X_0,X_1)$ with $X_0=-X_1\sim\mathcal{N}(0,\mathrm{Id})$
  • Figure 3: Iterative rectification for the example from \ref{['eq:counterexample1']} initialized with the optimal coupling $(X_0,X_1)$.
  • Figure 4: Iterative rectification for the example from \ref{['eq:counterexample1']} initialized with the non-optimal fixed point $(\tilde{X}_0,\tilde{X}_1)$ of $\mathcal{R}_p$.
  • Figure 5: Iterative rectification for the example from Section \ref{['sec:disconnected_support']} initialized with the independent coupling.
  • ...and 12 more figures

Theorems & Definitions (34)

  • Definition 1
  • Theorem 2: Affine Transformations
  • Theorem 3: Gaussian Case
  • Proposition 4
  • Theorem 5: Gaussian Mixture Case
  • Definition 6
  • Remark 7: Affine Invariances with Potential
  • Proposition 8
  • Corollary 9
  • Proposition 10
  • ...and 24 more