# Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching

arXiv: https://arxiv.org/abs/2505.23346 · HTML: https://mlsanity.com/paper/2505.23346/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 22 of its 27 references that are in the map.
- First posted on arXiv 29 May 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 22 of 22)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Improving the Training of Rectified Flows](https://mlsanity.com/paper/2405.20320/) (2024)
- [Flow Matching for Scalable Simulation-Based Inference](https://mlsanity.com/paper/2305.17161/) (2023)
- [Rectified Diffusion: Straightness Is Not Your Need in Rectified Flow](https://mlsanity.com/paper/2410.07303/) (2024)
- [Equivariant Flow Matching with Hybrid Probability Transport](https://mlsanity.com/paper/2312.07168/) (2023)
- [Relational Diffusion Distillation for Efficient Image Generation](https://mlsanity.com/paper/2410.07679/) (2024)
- [Towards Hierarchical Rectified Flow](https://mlsanity.com/paper/2502.17436/) (2025)

## Built on by (most cited first, 15 of 15)

- [When Safety Collides: Resolving Multi-Category Harmful Conflicts in Text-to-Image Diffusion via Adaptive Safety Guidance](https://mlsanity.com/paper/2602.20880/) (2026)
- [Delving into the Temporal Challenges of Unified Video Protection Against Image-to-Video and Fine-Tuning-based Customization](https://mlsanity.com/paper/2607.13336/) (2026)
- [SubFlow: Sub-mode Conditioned Flow Matching for Diverse One-Step Generation](https://mlsanity.com/paper/2604.12273/) (2026)
- [Explorative Modeling: Unlocking a Third Pretraining Axis and End-to-End Generation](https://mlsanity.com/paper/2607.27372/) (2026)
- [COT-FM: Cluster-wise Optimal Transport Flow Matching](https://mlsanity.com/paper/2603.13395/) (2026)
- [Flow Matching Calibration for Simulation-Based Inference under Model Misspecification](https://mlsanity.com/paper/2509.23385/) (2025)
- [Intellectual Property Protection for 3D Gaussian Splatting Assets: A Survey](https://mlsanity.com/paper/2602.03878/) (2026)
- [Geometry-Aware Image Flow Matching](https://mlsanity.com/paper/2605.25294/) (2026)
- [Optimal Transport Flow Matching by Design](https://mlsanity.com/paper/2606.04092/) (2026)
- [Reweighted Flow Matching via Unbalanced OT for Label-free Long-tailed Generation](https://mlsanity.com/paper/2509.25713/) (2025)
- [One-Sided Quantile Coupling for Flow Matching](https://mlsanity.com/paper/2608.00978/) (2026)
- [OT-MeanFlow3D: Bridging Optimal Transport and Meanflow for Efficient 3D Point Cloud Generation](https://mlsanity.com/paper/2509.22592/) (2025)
- [One-Step Generative Modeling via Training Dynamics Action](https://mlsanity.com/paper/2610.00518/) (2026)
- [Scalable Minimal-Change Learning for Controllable Image Editing](https://mlsanity.com/paper/2610.06021/) (2026)
- [Co-Evolving Paths and Flows via Path-Flow Alignment](https://mlsanity.com/paper/2610.08717/) (2026)

## Abstract (arXiv)

Flow Matching (FM) is an effective framework for training a model to learn a vector field that transports samples from a source distribution to a target distribution. To train the model, early FM methods use random couplings, which often result in crossing paths and lead the model to learn non-straight trajectories that require many integration steps to generate high-quality samples. To address this, recent methods adopt Optimal Transport (OT) to construct couplings by minimizing geometric distances, which helps reduce path crossings. However, we observe that such geometry-based couplings do not necessarily align with the model's preferred trajectories, making it difficult to learn the vector field induced by these couplings, which prevents the model from learning straight trajectories. Motivated by this, we propose Model-Aligned Coupling (MAC), an effective method that matches training couplings based not only on geometric distance but also on alignment with the model's preferred transport directions based on its prediction error. To avoid the time-costly match process, MAC proposes to select the top-k fraction of couplings with the lowest error for training. Extensive experiments show that MAC significantly improves generation quality and efficiency in few-step settings compared to existing methods. Project page: https://yexionglin.github.io/mac