On the Equivalence of Optimal Transport Problem and Action Matching with Optimal Vector Fields
Nikita Kornilov, Alexander Korotin
TL;DR
The paper exposes a fundamental link between optimal transport and vector-field-based transport methods used in generative modeling. By focusing on optimal vector fields, it shows that Action Matching (AM) inherently recovers OT solutions, just as Optimal Flow Matching (OFM) does for FM under the quadratic cost, with the AM loss aligning with the OT dual loss up to constants for any interpolation sequence. The key technical insight is that, for optimal vector fields, the AM functional reduces to expressions involving the Brenier potential $\Psi$ and its convex conjugate, effectively yielding the OT map via $\nabla \Psi^*$. This equivalence suggests a robust, sequence-wide approach to OT in practice, enabling OT solutions through AM formulations that operate on entire distributions trajectories rather than fixed interpolations.
Abstract
Flow Matching (FM) method in generative modeling maps arbitrary probability distributions by constructing an interpolation between them and then learning the vector field that defines ODE for this interpolation. Recently, it was shown that FM can be modified to map distributions optimally in terms of the quadratic cost function for any initial interpolation. To achieve this, only specific optimal vector fields, which are typical for solutions of Optimal Transport (OT) problems, need to be considered during FM loss minimization. In this note, we show that considering only optimal vector fields can lead to OT in another approach: Action Matching (AM). Unlike FM, which learns a vector field for a manually chosen interpolation between given distributions, AM learns the vector field that defines ODE for an entire given sequence of distributions.