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

ICLR 2026 Conference Submission12695 Authors

Keywords: Flow-based generative model; flow matching; Semi-discrete optimal transport

Abstract:

Flow-based Generative Models (FGMs) effectively transform noise into a data
distribution, and coupling the noise and data in the training of FGM by Optimal
Transport (OT) improves the straightness of the flow paths. However, existing OT-
based couplings are difficult to combine with modern models and/or to scale to
large datasets due to the curse of dimensionality in the sample complexity of (batch) OT.
This paper introduces AlignFlow, a new approach using Semi-Discrete Optimal
Transport (SDOT) to enhance FGM training by establishing explicit alignment
between noise and data pairs. SDOT computes a transport map by partitioning
the noise space into Laguerre cells, each mapped to a corresponding data point.
During the training of FGM, i.i.d.-sampled noise is matched with corresponding
data by the SDOT map. AlignFlow bypasses the curse of dimensionality and
scales effectively to large datasets and models. Our experiments demonstrate that
AlignFlow improves a wide range of state-of-the-art FGM algorithms and can be
integrated as a plug-and-play solution with negligible additional cost.

Primary Area: generative models

Submission Number: 12695