Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Lingkai Kong, Molei Tao, Yang Liu, Bryan Wang, Jinmiao Fu, Chien-Chih Wang, Huidong Liu

TL;DR

AlignFlow addresses the scalability and convergence challenges of coupling noise and data in Flow-based Generative Models by adopting Semi-Discrete Optimal Transport (SDOT) to construct a deterministic Noise–Data Alignment (NDA). The method computes an SDOT map in a Stage 1 offline step, partitioning the noise space into Laguerre cells and linking each cell to a data point, then uses this fixed NDA to train FGMs in Stage 2 with negligible overhead. This setup provides convergence guarantees for the coupling and avoids the dimensionality curse associated with discrete OT on minibatches. Empirically, AlignFlow improves performance and accelerates training across CIFAR-10, ImageNet256 with DiT, and ImageNet256 with SiT MeanFlow, while remaining plug-and-play and scalable to large architectures; code is available at the provided repository.

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.

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

TL;DR

AlignFlow addresses the scalability and convergence challenges of coupling noise and data in Flow-based Generative Models by adopting Semi-Discrete Optimal Transport (SDOT) to construct a deterministic Noise–Data Alignment (NDA). The method computes an SDOT map in a Stage 1 offline step, partitioning the noise space into Laguerre cells and linking each cell to a data point, then uses this fixed NDA to train FGMs in Stage 2 with negligible overhead. This setup provides convergence guarantees for the coupling and avoids the dimensionality curse associated with discrete OT on minibatches. Empirically, AlignFlow improves performance and accelerates training across CIFAR-10, ImageNet256 with DiT, and ImageNet256 with SiT MeanFlow, while remaining plug-and-play and scalable to large architectures; code is available at the provided repository.

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.
Paper Structure (33 sections, 1 theorem, 12 equations, 4 figures, 6 tables, 4 algorithms)

This paper contains 33 sections, 1 theorem, 12 equations, 4 figures, 6 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Notation
  5. Examples:
  6. Coupling between noise and data
  7. Semi-Discrete Optimal Transport
  8. Computational cost of Algo. \ref{['algo_sdot_optim']}
  9. The proposed AlignFlow
  10. Idea behind AlignFlow
  11. Additional techniques
  12. Advantages of AlignFlow
  13. Bypass the Curse of dimensionality
  14. Deterministic alignment
  15. Low computational cost
  16. ...and 18 more sections

Key Result

Theorem 1

In a $d$-dimensional space, the error in estimating the 2-Wasserstein distance $W_2^2(\tilde{q}, q_n)$ between a distribution $\tilde{q}$ and the emperical distribution with $n$-samples $q_n$ with only access to $n$ samples is of order $\sim n^{-1/d}$.

Figures (4)

  • Figure 1: Visualization of Laguerre cells in 2-dim. The noise distribution is the normal distribution (dark shadow in the left figure), and the dataset is the points in the lower left corner. The whole space is partitioned into cells using SDOT, and each region is mapped to the data point with the corresponding color by the SDOT map. The integral of the noise density over each cell $\mathsf{L}_i$ equals the discrete probability mass of the associated data point.
  • Figure 2: Training curves for AlignFlow on different tasks. Each figure illustrates the FID-50k score against the number of training steps. The results demonstrate that AlignFlow provides a consistent and simultaneous improvement over all baseline algorithms shown, enhancing both final performance and training convergence speed.
  • Figure 3: Comparison of generative trajectories in FGMs among various methods. AlignFlow has a straighter trajectory compared to vanilla Flow Matching and has a clearer boundary compared to Minibatch OT (e.g., at $t=0.22$).
  • Figure 4: Images generated by MeanFlow+AlignFlow trained on ImageNet256 (FID=3.23, NFE=1).

Theorems & Definitions (5)

  • Remark 1: Noise storage
  • Remark 2: Data augmentation
  • Remark 3: Class-conditioned generation
  • Definition 1: Noise-Data Alignment
  • Theorem 1: Sample complexity of OT (Informal version for Thm. 1 in fournier2015rate)