# AlignFlow: Improving Flow-based Generative Models with Semi-Discrete Optimal Transport

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 40 references that are in the map.
- First posted on arXiv 16 Oct 2025.

## Builds on (most cited first, 30 of 31)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [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)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [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)
- [Flow map matching with stochastic interpolants: A mathematical framework for consistency models](https://mlsanity.com/paper/2406.07507/) (2024)
- [Consistency Flow Matching: Defining Straight Flows with Velocity Consistency](https://mlsanity.com/paper/2407.02398/) (2024)
- [Optimal Flow Matching: Learning Straight Trajectories in Just One Step](https://mlsanity.com/paper/2403.13117/) (2024)
- [Improving Mini-batch Optimal Transport via Partial Transportation](https://mlsanity.com/paper/2108.09645/) (2021)
- [On the Relation between Rectified Flows and Optimal Transport](https://mlsanity.com/paper/2505.19712/) (2025)
- [The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation](https://mlsanity.com/paper/2503.10636/) (2025)
- [Self-Corrected Flow Distillation for Consistent One-Step and Few-Step Text-to-Image Generation](https://mlsanity.com/paper/2412.16906/) (2024)
- [Asymptotics for semi-discrete entropic optimal transport](https://mlsanity.com/paper/2106.11862/) (2021)

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

- [Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching](https://mlsanity.com/paper/2602.05951/) (2026)
- [Optimal Transport Flow Matching by Design](https://mlsanity.com/paper/2606.04092/) (2026)
- [Gromov-Monge Flow Matching for Equivariant Graph Generation](https://mlsanity.com/paper/2608.26961/) (2026)
- [Dual-guided Hierarchical Edge Localization for Large-scale Optimal Transport Across Dimensions](https://mlsanity.com/paper/2609.13010/) (2026)
- [Beyond Random Couplings: Contrastive Noise Alignment in Generative Flows](https://mlsanity.com/paper/2609.18488/) (2026)
- [Gumbel Straight Flow: Distilling Autoregressive Models into One-step Flow Maps](https://mlsanity.com/paper/2610.00497/) (2026)
- [Scaling Online Optimal Transport Minibatches for Flow Matching](https://mlsanity.com/paper/2610.18838/) (2026)

## Abstract (arXiv)

Flow-based Generative Models (FGMs) effectively transform noise into complex data distributions. Incorporating Optimal Transport (OT) to couple noise and data during FGM training has been shown to improve the straightness of flow trajectories, enabling more effective inference. However, existing OT-based methods estimate the OT plan using (mini-)batches of sampled noise and data points, which limits their scalability to large and high-dimensional datasets in FGMs. This paper introduces AlignFlow, a novel approach that leverages Semi-Discrete Optimal Transport (SDOT) to enhance the training of FGMs by establishing an explicit, optimal alignment between noise distribution and data points with guaranteed convergence. SDOT computes a transport map by partitioning the noise space into Laguerre cells, each mapped to a corresponding data point. During FGM training, i.i.d. noise samples are paired with data points via the SDOT map. AlignFlow scales well to large datasets and model architectures with negligible computational overhead. Experimental results show that AlignFlow improves the performance of a wide range of state-of-the-art FGM algorithms and can be integrated as a plug-and-play component. Code is available at: https://github.com/konglk1203/AlignFlow.