Flow Matching from Viewpoint of Proximal Operators

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.

💡 Research Summary

The paper presents a rigorous reformulation of Optimal Transport Conditional Flow Matching (OT‑CFM), a class of dynamical generative models that transport samples from a simple base distribution to a complex target distribution via an ODE‑driven vector field. The authors show that, even when the target distribution lacks a density (e.g., it lives on a low‑dimensional manifold), the optimal quadratic‑cost transport plan can always be represented by the sub‑differential of a convex function ϕ, which they call the Aleksandrov–Brenier potential. This potential generalizes the classical Brenier map: when a density exists, ϕ is differentiable almost everywhere and the map reduces to the gradient of ϕ; otherwise ϕ may be set‑valued but still encodes the optimal coupling.

Using the potential, the authors define a time‑dependent interpolation operator Hₜ(x)=αₜ∂ϕ(x)+βₜx, where αₜ and βₜ are deterministic schedules satisfying α₀=1, α₁=0, β₀=0, β₁=1. They introduce a strongly convex function ψₜ(x)=αₜϕ(x)+βₜ‖x‖²/2 whose sub‑gradient coincides with Hₜ. By convex duality, the inverse of Hₜ is single‑valued and given by the gradient of the convex conjugate ψₜ*: Hₜ⁻¹(y)=∇ψₜ*(y). Crucially, Lemma 2.1 proves that this inverse can be expressed as a proximal operator: ∇ψₜ*(y)=prox_{λₜϕ}(y/βₜ) with λₜ=αₜ/βₜ. Consequently, the OT‑CFM vector field admits an exact closed‑form: vₜ(y)= (α̇ₜ/αₜ) y + (β̇ₜ/βₜ−α̇ₜ/αₜ) prox_{λₜϕ}(y/βₜ). This representation mirrors the implicit Euler step of a proximal algorithm and provides a deterministic “denoiser” that recovers the target point from any intermediate state y.

The paper also addresses the practical setting where OT‑CFM is trained with minibatch optimal transport. The authors prove that, as the minibatch size grows, the empirical transport plan converges (along a subsequence) to the population optimal plan π*. Therefore the vector field learned from finite batches converges to the exact proximal formulation derived above, justifying the empirical observation that larger batches yield smoother, lower‑variance training signals.

A major theoretical contribution is the stability analysis under the manifold hypothesis. By computing the second epi‑derivative of ϕ, the authors derive the Lyapunov exponents of the terminal‑time dynamics (t→1). After a natural time rescaling, they show that directions tangent to the data manifold have zero exponent (neutral dynamics), while normal directions have strictly negative exponents, implying exponential contraction. This property, termed “terminally normally hyperbolic,” guarantees that the learned flow is robust to small perturbations orthogonal to the manifold while preserving the intrinsic geometry along the manifold.

In summary, the paper makes four key contributions: (1) a novel convex‑analysis formulation of OT‑CFM using the Aleksandrov–Brenier potential; (2) an exact proximal expression for the vector field, valid even for manifold‑supported targets; (3) a convergence proof linking minibatch OT‑CFM to the population formulation; and (4) a rigorous normal‑hyperbolicity analysis showing exponential normal‑direction contraction and neutral tangential dynamics. These results deepen the theoretical understanding of flow‑matching models, bridge optimal transport, proximal algorithms, and dynamical systems, and provide concrete guidance for designing stable, efficient generative flows.