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

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. Experimental results show that Align-Flow improves the performance of a wide range of state-of-the-art FGM algorithms and can be integrated as a plug-and-play component. A generative model in machine learning is designed to produce new data samples that closely resemble those drawn from a given dataset. This task is of fundamental importance and has seen significant advances over the past decades.