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

2025-10-16

AI Summary

Existing flow-based generative models (FGMs) trained via optimal transport (OT) rely on mini-batch noise–data sampling to estimate OT plans, leading to poor scalability on large-scale, high-dimensional data, curved generation trajectories, and inefficient inference. This work introduces semi-discrete optimal transport (SDOT) into FGMs for the first time. By partitioning the noise space via Laguerre tessellation, SDOT establishes an explicit, provably convergent optimal alignment from the continuous noise distribution to discrete data points—enabling deterministic, point-wise noise-to-data matching. The method supports i.i.d. noise sampling and data pairing, integrates seamlessly into mainstream FGM frameworks, and is plug-and-play. Experiments demonstrate substantial improvements in generation quality and trajectory straightness across multiple state-of-the-art FGMs, significant inference acceleration, and negligible computational overhead—while maintaining efficiency and stability on large-scale datasets.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Planning with Language Models
• Machine Learning: Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

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.

Problem

• Improving flow-based generative models with semi-discrete optimal transport
• Enhancing training by establishing explicit noise-data alignment with convergence guarantees
• Scaling optimal transport methods to large datasets with minimal computational overhead

Innovation

• Leverages Semi-Discrete Optimal Transport for training
• Partitions noise space into Laguerre cells for alignment
• Provides plug-and-play component for flow models