Research map / Probabilistic and scientific modeling / Variational and generative sampling
Optimal Flow Matching: Learning Straight Trajectories in Just One Step
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Optimal Flow Matching (OFM), which restricts flow matching to vector fields given by gradients of convex potentials (ICNNs) so that one FM training step recovers the exact quadratic-cost OT map with straight trajectories.
Problem
Flow matching models produce curved trajectories that need many ODE steps; existing straightening methods (Rectified Flow, OT-CFM) are iterative, accumulate error, or rely on biased minibatch OT, so they fail to recover the quadratic-cost OT displacement.
Approach
Restrict the FM loss to optimal vector fields generated by gradients of convex functions parametrized by ICNNs; the resulting loss has the same minimizers as the dual OT loss for any initial transport plan, giving exactly straight OT paths in a single FM objective and simulation-free inference.
Key results
- On the W2 benchmark, OFM MB achieves L2-UVP of 0.15% at D=2 and 0.52% at D=4, outperforming OT-CFM (0.16%, 0.73%) and RF on most dimensions among FM-based methods.
- OFM MB reaches FID 11.0 on Adult-to-Child translation in ALAE latent space, versus 12.9 for OT-CFM, 13.5 for c-RF, and 21.0 for RF.
- MLP-based amortized OT solver still beats OFM on the W2 benchmark, and OFM with an MLP instead of ICNN failed to converge.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use OFM instead of OT-CFM | Flow Matching solvers · Wasserstein-2 benchmark D=256 | L²-UVP (%): 27.5 -> 2.9 | Helped |
| Use OFM instead of c-RF | Flow Matching solvers · Wasserstein-2 benchmark D=256 | L²-UVP (%): 76.48 -> 2.9 | Helped |
| Use minibatch OT plan instead of independent plan | OFM · Wasserstein-2 benchmark D=64 | L²-UVP (%): 8.35 -> 7.2 | Helped |
| Use minibatch OT plan instead of independent plan | OFM · Adult→Child translation | FID: 11.8 -> 11.0 | Helped |
| Use OFM MB instead of OT-CFM | Flow Matching solvers · Adult→Child translation | FID: 12.9 -> 11.0 | Helped |
| Use OFM MB instead of c-RF | Flow Matching solvers · Adult→Child translation | FID: 13.5 -> 11.0 | Helped |
| Use minibatch OT plan instead of independent plan | OFM · Wasserstein-2 benchmark | L²-UVP (%) | Helped |
| Replace ICNN with MLP (ablation) | OFM · Wasserstein-2 benchmark | convergence | Hurt |
Surprising findings
- OFM with an MLP parametrization instead of an ICNN fails to converge, indicating the convexity constraint is essential in practice.
- Rectified Flow performs worse than even the linear baseline on high-dimensional W2 benchmarks.
Limitations
- MLP-based parametrization fails, so the method depends on ICNN architectures whose expressivity is limited in practice.
- Minibatch OT plans with large variance can converge to slightly worse solutions, and in high dimensions OFM still trails the MLP-based amortized OT solver.
Experimental setup
Datasets: Wasserstein-2 benchmark FFHQ
Benchmarks: Wasserstein-2 benchmark (L2-UVP, cosine similarity) FFHQ Adult-to-Child translation (FID)
Metrics: L2-UVP cosine similarity FID
Baselines: MMv1 Amortized OT ICNN Amortized OT MLP Linear map OT-CFM Rectified Flow c-Rectified Flow
Builds on: Flow Matching OT-CFM Rectified Flow c-Rectified Flow Input Convex Neural Networks
Closest papers by meaning
- On the Relation between Rectified Flows and Optimal Transport (NeurIPS 2025 · similarity 0.91)
- Value Gradient Guidance for Flow Matching Alignment (NeurIPS 2025 · similarity 0.90)
- Switched Flow Matching: Eliminating Singularities via Switching ODEs (ICML 2024 · similarity 0.90)
- Beckmann Transport Models: From Autonomous Flows to One-Step Maps (NeurIPS 2026 · similarity 0.90)
- On the Guidance of Flow Matching (ICML 2025 · similarity 0.90)
- Improving Flow Matching by Aligning Flow Divergence (ICML 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.