Optimal Flow Matching: Learning Straight Trajectories in Just One Step

TL;DR

The novelOptimal Flow Matching (OFM) approach is developed which allows recovering the straight OT displacement for the quadratic transport in just one FM step by employing vector field for FM which are parameterized by convex functions.

Abstract

Over the several recent years, there has been a boom in development of Flow Matching (FM) methods for generative modeling. One intriguing property pursued by the community is the ability to learn flows with straight trajectories which realize the Optimal Transport (OT) displacements. Straightness is crucial for the fast integration (inference) of the learned flow's paths. Unfortunately, most existing flow straightening methods are based on non-trivial iterative FM procedures which accumulate the error during training or exploit heuristics based on minibatch OT. To address these issues, we develop and theoretically justify the novel \textbf{Optimal Flow Matching} (OFM) approach which allows recovering the straight OT displacement for the quadratic transport in just one FM step. The main idea of our approach is the employment of vector field for FM which are parameterized by convex functions.

Figures (10)

• Figure 1: Our Optimal Flow Matching (OFM). For any initial transport plan $\pi$ between $p_0$ and $p_1$, OFM obtains exactly straight trajectories (in just a single FM loss minimization) which carry out the OT displacement for the quadratic cost function.
• Figure 2: Flow Matching (FM) obtains a vector field $u$ moving $p_0$ to $p_1$. FM typically operates with the independent transport plan $\pi=p_0\times p_1$.
• Figure 3: OT-CFM uses minibatch OT plan to obtain more straight trajectories.
• Figure 4: Rectified Flow iteratively applies FM to straighten the trajectories after each step.
• Figure 5: An Optimal Vector Field: a vector field $u^\Psi$ with straight paths is parametrized by a gradient of a convex function $\Psi$.

Theorems & Definitions (14)

• Theorem 1: OFM and OT connection
• Proposition 1: Explicit Loss Gradient Formula
• Proposition 2: Simplified OFM Loss
• Proposition 3: Intractable Distance
• Proposition 4: Tractable Distance For OFM
• proof
• proof
• Lemma 1: Properties of convex functions and their conjugates
• proof
• Lemma 2: Main Integration Lemma