Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching

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, the source comes from a continuous distribution, only a finite number of target samples data are accessible, and the neural network has finite capacity. Therefore, making the coupling more OT-like may not consistently improve sampling quality in the low number-of-function-evaluation (NFE) regime. Motivated by this, we introduce ReCouple, a framework that constructs coarse-grained OT couplings, to better understand rollout-free OT-based post-training. ReCouple+ extends this toward the continuous-source setting by increasing source resolution; and it converges to semi-discrete OT in the appropriate limit. ReCouple++ further adds spatial regularization, motivated by the observation that straight trajectories need not correspond to spatially simple velocity fields. Empirical validations demonstrate their effectiveness, while revealing an NFE-dependent interplay between coupling optimality, neural realizability, and generation quality.