Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching
Abstract
Flow matching depends critically on the coupling used to pair source and target samples. Optimal transport (OT) couplings can reduce velocity conflict and straighten trajectories; but, in practice, challenges remain: the source comes from a continuous distribution, only a finite number of target samples data are accessible, and the neural network has finite capacity. For these reasons, making the coupling more OT-like may not consistently improve sampling quality, in particular in the low number-of-function-evaluation (NFE) regime. Motivated by this, we introduce ReCouple, a controllable framework that constructs coarse-grained OT couplings, to better understand rollout-free OT-based post-training. We also consider extensions of this basic framework: ReCouple+ extends this toward the continuous-source setting by increasing source resolution, converging to semi-discrete OT in the appropriate limit; and ReCouple++ further adds spatial regularization, motivated by the observation that straight trajectories need not correspond to spatially simple velocity fields. Empirical validation demonstrates the effectiveness of these methods, and it reveals an NFE-dependent interplay between coupling optimality, neural realizability, and generation quality.
Related
- Multiscale Supervised Unbalanced Optimal Transport Flow Matching
- Regularized Coupling Maps on Matrix Polytopes: From Assignment to Graph Matching and Optimal Transport
- Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport
- Expected Batch Optimal Transport Plans and Consequences for Flow Matching
- Statistical Inference for Adversarial Training: Central Limit Theorems via Optimal Transport
BibTeX
@misc{lim2026coarsegraining,
title = {{Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching}},
author = {Soon Hoe Lim and Aditi Gupta and Michael W. Mahoney and N. Benjamin Erichson},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.41269},
eprint = {2610.41269},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.41269},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching
AU - Soon Hoe Lim
AU - Aditi Gupta
AU - Michael W. Mahoney
AU - N. Benjamin Erichson
PY - 2026
DA - 2026/10/30
PB - arXiv
DO - 10.48550/arxiv.2610.41269
UR - https://arxiv.org/abs/2610.41269
AB - Flow matching depends critically on the coupling used to pair source and target samples. Optimal transport (OT) couplings can reduce velocity conflict and straighten trajectories; but, in practice, challenges remain: the source comes from a continuous distribution, only a finite number of target samples data are accessible, and the neural network has finite capacity. For these reasons, making the coupling more OT-like may not consistently improve sampling quality, in particular in the low number-of-function-evaluation (NFE) regime. Motivated by this, we introduce ReCouple, a controllable framework that constructs coarse-grained OT couplings, to better understand rollout-free OT-based post-training. We also consider extensions of this basic framework: ReCouple+ extends this toward the continuous-source setting by increasing source resolution, converging to semi-discrete OT in the appropriate limit; and ReCouple++ further adds spatial regularization, motivated by the observation that straight trajectories need not correspond to spatially simple velocity fields. Empirical validation demonstrates the effectiveness of these methods, and it reveals an NFE-dependent interplay between coupling optimality, neural realizability, and generation quality.
ER -