Flow Matching from Viewpoint of Proximal Operators

TL;DR

It is proved that, for manifold-supported targets, OT-CFM is terminally normally hyperbolic: after time rescaling, the dynamics contracts exponentially in directions normal to the data manifold while remaining neutral along tangential directions.

Abstract

We reformulate Optimal Transport Conditional Flow Matching (OT-CFM), a class of dynamical generative models, showing that it admits an exact proximal formulation via an extended Brenier potential, without assuming that the target distribution has a density. In particular, the mapping to recover the target point is exactly given by a proximal operator, which yields an explicit proximal expression of the vector field. We also discuss the convergence of minibatch OT-CFM to the population formulation as the batch size increases. Finally, using second epi-derivatives of convex potentials, we prove that, for manifold-supported targets, OT-CFM is terminally normally hyperbolic: after time rescaling, the dynamics contracts exponentially in directions normal to the data manifold while remaining neutral along tangential directions.

Figures (6)

• Figure 1: Relationship between a target sample$x_1$, noise sample$x_0$, and the conditional path $x_t$. Here $\phi$ denotes the Aleksandrov--Brenier potential (see Sec. \ref{['sec:AB-potential']}). The target $x_1$ and noise $x_0$ are connected via $x_0 = T(x_1)$, where $T = \partial \phi$ is the (extended) Brenier map. In parallel, $x_t$ is mapped back to target $x_1$ via the denoiser $H_t^{-1}=\nabla\psi_t^\ast$, where $H_t$ is the interpolation operator (see Sec. \ref{['sec:VF_prox']}).
• Figure 2: Example of low-dimensional target probability. Here, $P_1$ is a singular distribution with a density on the 1D manifold, $\mathcal{M}$, shown in orange, while $P_0$ is the normal distribution $N(0,I_2)$ with density in 2D, as shown in the gray density plot. The subdifferential of the Aleksandrov-Brenier potential map, $\partial\phi,$ at $z\in \mathcal{M}$, is normal to $\mathcal{M}$ everywhere in this case, as shown by gray arrows.
• Figure 3: Circle. Left: OT-CFM vector field shown at time $t=0.9$ obtained by a Neural Network, which appears to be an attracting force on the manifold. Right: Sample mean of the eigenvalues of the Jacobian ${\color{black} (1-t)}Dv_t$ at different times over different flow trajectories. The second eigenvalue has a small variance and hence, the terminal Lyapounov exponent is also close to -1, as predicted by the analysis in section \ref{['sec:lyap-prox-epi']}.
• Figure 4: Two moons. Left: Vector field, Right: Eigenvalues of the Jacobian. Similar settings are applied to Fig. \ref{['fig:circle']}.
• Figure 5: Top: Generation with OT-CFM. Middle: Eigenvectors of the Jacobian of the vector field $v_t(x)$ with respect to the 1st, 20th, 40th, 80th, and 160th eigenvalues in the descending order. Bottom: Images perturbed with the eigenvectors in the middle row.

Theorems & Definitions (42)

• Lemma 2.1
• proof
• Corollary 2.2
• Proposition 4.1
• Lemma 4.2
• Lemma 4.3
• Theorem 4.4: Lyapunov exponents
• proof
• Definition A.1: Prox-boundedness
• Definition A.2: Prox-regularity