OAT-FM: Optimal Acceleration Transport for Improved Flow Matching

TL;DR

This paper advances flow matching by recasting it as a second-order transport via Optimal Acceleration Transport (OAT) in the product space of samples and velocities. It introduces OAT-FM, a bi-level optimization approach that refines pre-trained velocity fields with a velocity-alignment and acceleration-minimization objective, yielding a practical two-phase FM paradigm. The method demonstrates consistent improvements across low-dimensional OT benchmarks, CIFAR-10, and large-scale ImageNet generation, while remaining computationally efficient and avoiding data distribution drift. The work also outlines limitations and avenues for improvement, including warm-start strategies and faster OAT solvers for scalable, high-dimensional applications.

Abstract

As a powerful technique in generative modeling, Flow Matching (FM) aims to learn velocity fields from noise to data, which is often explained and implemented as solving Optimal Transport (OT) problems. In this study, we bridge FM and the recent theory of Optimal Acceleration Transport (OAT), developing an improved FM method called OAT-FM and exploring its benefits in both theory and practice. In particular, we demonstrate that the straightening objective hidden in existing OT-based FM methods is mathematically equivalent to minimizing the physical action associated with acceleration defined by OAT. Accordingly, instead of enforcing constant velocity, OAT-FM optimizes the acceleration transport in the product space of sample and velocity, whose objective corresponds to a necessary and sufficient condition of flow straightness. An efficient algorithm is designed to achieve OAT-FM with low complexity. OAT-FM motivates a new two-phase FM paradigm: Given a generative model trained by an arbitrary FM method, whose velocity information has been relatively reliable, we can fine-tune and improve it via OAT-FM. This paradigm eliminates the risk of data distribution drift and the need to generate a large number of noise data pairs, which consistently improves model performance in various generative tasks. Code is available at: https://github.com/AngxiaoYue/OAT-FM

Figures (8)

• Figure 1: (a) The principle of OAT-FM and the corresponding two-phase FM paradigm. (b) We train SiT-XL ma2024sit by our two-phase FM paradigm on the ImageNet 256$\times$256 dataset deng2009imagenet.
• Figure 2: An illustration of refining the flow of I-CFM via OAT-FM on the eight Gaussians to the Moons dataset. We conduct this experiment based on the code base provided in tong2024improving.
• Figure 3: The comparison of SiT-XL and SiT-XL + OAT-FM on FID and sFID.
• Figure 4: The visual comparison for SiT-XL and SiT-XL + OAT-FM when CFG is 4.0.
• Figure 5: Stability analysis of OAT-FM fine-tuning on CIFAR-10. Starting from a pre-trained FM model (400K batches), we refine the model using OAT-FM for an additional 1K training samples, recording metrics every 0.1K samples. We track generation quality (FID with 10 and 100 Euler steps) and flow straightness.

Theorems & Definitions (7)

• Proposition 1
• Definition 1: Dynamic Formulation of Optimal Acceleration Transport (OAT) benamou2019second
• Definition 2: Kantorovich formulation of OAT Chen2018MeasureValuedbenamou2019secondBrigati2025KineticOT
• Theorem 2: Straightening Flow via OAT
• Theorem 3: OAT Bound of OAT-FM
• Theorem 4: Straightening a single trajectory via acceleration minimization
• proof