Optimal Transport Flow Matching by Design

2026-06-02

AI Summary

This work addresses the computational intractability of high-dimensional optimal transport couplings, which often leads existing flow-matching methods to suffer from either bias or excessive computational cost. The authors propose a novel approach that treats the prior distribution as a designable variable, constructing it via low-frequency image projection to simultaneously satisfy the optimal transport identity coupling and enable efficient sampling. This reformulation reduces the generative task to synthesizing high-frequency details alone. Notably, the method achieves efficient single-step generation without modifying the underlying flow model and further enhances sample quality through Gaussian interpolation. Experimental results demonstrate a more than two-fold reduction in trajectory curvature across all benchmarks, substantially improving performance in few-step and even single-step generation settings.

Abstract

Flow matching models learn to transport samples from a simple prior distribution to a complex data distribution. When prior-data pairs are coupled via optimal transport (OT), the learned trajectories are straight and non-crossing, enabling fast, even single-step, generation. However, computing the OT coupling in high dimensions is intractable, and existing methods attempt to solve the OT problem, at the cost of persistent bias or significant overhead. Rather than solving for the OT coupling, we reformulate the problem. Once the prior is treated as a design choice rather than a fixed input, the OT coupling between prior and data is no longer unique. Many priors admit an OT-optimal identity coupling to the data, leaving us free to choose one that is also tractable to sample. We identify low-frequency projection of natural images as such a choice. The identity coupling between data and its low-frequency representation is empirically OT-optimal, the prior is structured enough to be sampled by a lightweight model at inference, and the remaining flow-matching task reduces to synthesizing high-frequency detail. Interpolating the prior with Gaussian noise further improves generation quality while preserving the OT coupling. The approach requires no modifications to the flow model itself, and integrates naturally with latent-space models, classifier-free guidance, and one-step generation frameworks. Across all benchmarks, our method reduces trajectory curvature by more than $2\times$ compared to existing flow matching methods, yielding better generation quality in the few-step regime.

Problem

• Optimal Transport
• Flow Matching
• High-dimensional Coupling
• Generation Efficiency
• Prior Design

Innovation

• Optimal Transport
• Flow Matching
• Low-frequency Prior
• Identity Coupling
• One-step Generation