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

## Abstract

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 .

## 1 Introduction

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. Notable examples include autoregressive models such as GPT (Radford et al., 2018) and ChatGPT (Achiam et al., 2023) for natural language generation, and diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021) for image synthesis (e.g., Rombach et al., 2022). In addition, other prominent generative modeling approaches include Generative Adversarial Networks (GANs) (Goodfellow et al., 2020), Normalizing Flows (Rezende & Mohamed, 2015), Rectified Flow (Liu et al., 2022), Flow Matching (Lipman et al., 2022), and Stochastic Interpolant (Albergo et al., 2023).

This work focuses on improving a wide range of Flow-based Generative Models (abbreviated as FGMs hereafter), including Flow Matching (Lipman et al., 2022), shortcut model (Frans et al., 2025), MeanFlow (Geng et al., 2025), Live Reflow (Frans et al., 2025), but excluding continuous normalizing flows (CNF) (Chen et al., 2018); see Sec. [3.1](https://arxiv.org/html/2510.15038v1#S3.SS1) for a detailed specification of scope. FGMs focus on learning a time-dependent vector field, approximated by a Neural Network (NN) whose integration gives a trajectory that transforms a randomly sampled noise to a generated datum.

Despite their ability to generate high-quality samples, FGMs (like other related methods such as diffusion models) can require substantial computational cost in their inference step (i.e. sampling), where an Ordinary Differential Equation (ODE) is numerically integrated and each integration step involves at least one forward pass of a large NN. Consequently, generating a single sample requires multiple NN evaluations. This cost is typically measured by the Number of Function Evaluations (NFEs), which denotes the total number of forward passes during ODE integration. For example, in vanilla Flow Matching (Lipman et al., 2022), the NFEs is often greater than 100. The straightness of the ODE trajectory is closely related to NFE, as a straighter trajectory is typically easier to integrate, resulting in fewer NFEs required. See Apdx. [C.2](https://arxiv.org/html/2510.15038v1#A3.SS2) for the inference procedure of FGMs.

Before diving into how to improve trajectory straightness, we first review the training procedure of FGMs. Each training iteration of FGMs generally involves three main steps: (1) sampling noise and data points; (2) constructing the target vector field that a NN is trained to approximate; and (3) updating the model parameters. For implementation details, see Algorithms [1](https://arxiv.org/html/2510.15038v1#algorithm1) and [4](https://arxiv.org/html/2510.15038v1#algorithm4). In this work, we focus on improving the first step. While many State-of-the-Art (SOTA) methods introduce various target vector fields to promote straighter flow trajectories in step 2, noise and data points are still typically sampled independently in step 1. However, such independence has been shown to inherently induce curved flow trajectories and lead to high NFEs during sampling. In other words, the random matching of noise and data points implicitly encourages non-straight generative paths, motivating carefully designed couplings (e.g., Liu et al., 2022; Hertrich et al., 2025).

One family of approaches for making the trajectories straighter is based on using Optimal Transport (OT) to couple noise and data more effectively
(Tong et al., 2023; Pooladian et al., 2023; Kornilov et al., 2024; Cheng & Schwing, 2025). From a theoretical perspective, OT gives a shortest path between the noise and data distributions (Eq. ([4](https://arxiv.org/html/2510.15038v1#S3.E4))), yielding inherently straight mappings suitable for coupling. Despite the appealing theoretical properties of OT, scaling OT-based FGMs to large-scale problems remains challenging (Sec. [2](https://arxiv.org/html/2510.15038v1#S2)). Utilizing discrete OT, Tong et al. (2023) and Pooladian et al. (2023) couple the noise samples and data samples within each minibatch. However, discrete OT plans between minibatches of samples are known to provide biased  (Kornilov et al., 2024) or misspecified (Nguyen et al., 2022) approximation due to a fundamental challenge, namely that an accurate estimation of the OT plan in this setting requires the number of noise samples to grow exponentially with the data dimensionality, known as the curse of dimension (Dudley, 1969; Genevay et al., 2019). In contrast to discrete OT-based methods, Kornilov et al. (2024) uses a continuous OT approach by representing the OT map using a Brenier potential parametrized by an Input Convex Neural Network (ICNN) (Amos et al., 2017). Their algorithm jointly optimizes the transport map and the Flow Inversion Map. Nonetheless, the optimality and convergence of the learned transport map are not established.

We therefore propose a different approach, AlignFlow, based on the following observation: to train a FGM, there are only finite amount of training data, corresponding to a discrete empirical distribution, but we can (and typically do) use a noise distribution that is continuous. Therefore, we seek to optimally couple discrete distribution with continuous distribution, and the OT problem corresponding to this case is known as Semi-Discrete Optimal Transport (SDOT). This is different from the aforementioned discrete OT approaches that use both data and noise samples to estimate an OT map or the continuous OT approach that introduces extra learning component with inductive bias (ICNN). Nevertheless, as a special case of OT, SDOT still inherits its desirable theoretical properties: it defines the most direct and shortest connection between noise samples and data points, thereby guiding the target vector field learned by a NN. Moreover, the SDOT map can be solved efficiently with guaranteed convergence, whose quality of the resulting SDOT map can be evaluated with low computational cost (Liu et al., 2021). In addition, SDOT also has a clear geometric interpretation: the SDOT map partitions the space into Laguerre cells (Fig. [1](https://arxiv.org/html/2510.15038v1#S3.F1)), where each cell corresponds to a data point, and all noise samples residing in that cell are optimally aligned to the data point.

SDOT map is a deterministic map from noise to data, unlike a general coupling which can match one noise sample to different data probabilistically (or vice versa). Such determinism is a necessary condition for achieving optimal coupling and ensuring consistent noise-data matching regardless of batch size (Sec. [4](https://arxiv.org/html/2510.15038v1#S4)) in training FGMs. Motivated by the favorable properties of the determinism of SDOT, we call special couplings that are given by a deterministic mapping from the noise space to a dataset Noise–Data Alignment (NDA) (Def. [1](https://arxiv.org/html/2510.15038v1#Thmdefinition1)). While many prior works leverage OT to improve FGMs, the key innovation of AlignFlow is its use of SDOT to explicitly construct an NDA.

The training procedure of AlignFlow follows a two-stage approach: in the first stage, the SDOT map is computed. In the second stage, a general FGM is trained using this precomputed SDOT map. The key advantages for AlignFlow are summarized below:

- •
AlignFlow is a *plug-and-play* method that can be seamlessly integrated into existing FGMs. It readily combines with SOTA FGMs to further enhance performance.
- •
AlignFlow has its NDA component separated from FGM training, and therefore enjoys the provable convergence of SDOT in (Liu et al., 2021). In contrast, previous OT-based FGMs lack convergence guarantees for the computation of coupling between noise and data.
- •
AlignFlow *bypasses the curse of dimensionality* (Sec. [4.1](https://arxiv.org/html/2510.15038v1#S4.SS1)) that arises when using discrete OT for continuous distributions (Genevay et al., 2019), enabling it to scale effectively to large models and high-dimensional datasets.
- •
The SDOT map defines a *deterministic* transport path, meaning that each noise sample is consistently matched to a specific data point, regardless of batch size in FGM training. This property makes AlignFlow particularly well-suited for training large models where small batch sizes are necessary due to memory constraints.
- •
The additional cost of computing the SDOT map is minimal (less than 1% extra cost in practice) (Sec. [4.3](https://arxiv.org/html/2510.15038v1#S4.SS3)).

## 2 Related Work

A central factor in reducing the NFEs during sampling in FGMs is the straightness of the generative trajectory. Since the inference process involves integrating a learned vector field, straighter trajectories are easier to integrate accurately, thereby reducing NFEs. Numerous methods have been proposed to encourage straighter trajectories, which can be broadly categorized into three main approaches:

Target vector field: This class of methods aims to straighten generative trajectories by guiding the NN to learn a smoother or more linear target vector field. Approaches such as consistency models (Song et al., 2023) and shortcut models (Frans et al., 2025), inspired by techniques from diffusion models (Yang et al., 2024), enforce penalties on inconsistencies between the forward and backward segments of the trajectory. MeanFlow (Geng et al., 2025) introduces a stronger regularization term based on the Jacobian-vector product. These advanced loss functions have proven effective, reducing the NFEs to as low as 4 and even 1 in class-conditioned generation tasks on the ImageNet dataset.

Distillation: Recent works aim to reduce NFEs by distilling a trained FGM into a more efficient model (Boffi et al., 2025; Dao et al., 2025). Distillation techniques have also shown success in diffusion models (Salimans & Ho, 2022; Song et al., 2023; Kim et al., 2023). However, distillation involves an additional training stage built upon a pre-trained model, and thus can be viewed as complementary to our approach.

Coupling: Coupling is a joint distribution over two marginals. In standard FGMs, noise and data are independently paired during training. Several recent works aim to improve FGMs by introducing better coupling strategies by selecting noise and data from a carefully designed joint distribution:

- •
Tong et al. (2023) and Pooladian et al. (2023) employ the Sinkhorn algorithm (Peyré et al., 2019, Sec. 4.2) to compute the OT plan between i.i.d. sampled noise and data points in each FGM training iteration. However, this approach is sensitive to batch size: large batches incur high computational cost, as Sinkhorn is an iterative algorithm with per-batch complexity of $\mathcal{O}(B^{2}\log B)$ where $B$ is the batch size, while small batch size leads to poor estimates of the optimal coupling. Additionally, computing discrete OT in batches is known to produce biased transport plans (Kornilov et al., 2024), or misspecified mappings (Nguyen et al., 2022).
- •
Kornilov et al. (2024) train an Input Convex Neural Network (ICNN) (Amos et al., 2017) to serve as the Brenier potential (Peyré et al., 2019, Thm. 2.1), thereby providing the OT map between the noise and real data distributions. The ICNN and the Flow Inversion Map (Sec. C.a in Kornilov et al. (2024)) are optimized alternatively. However, the convergence property and optimality of the implicit OT plan approximated by ICNN remain unestablished.
- •
Liu et al. (2022) propose to disentangle crossing trajectories in the learned vector field to promote straighter generative trajectories. This approach has recently been scaled to larger models by Esser et al. (2024), although the scaling process is non-trivial. While the method does not directly solve an OT problem, its underlying formulation is closely related to OT, as discussed in Theorem 3.5 of their paper.

A comparison is provided in Table [1](https://arxiv.org/html/2510.15038v1#S2.T1). While conceptually related, AlignFlow differs fundamentally from existing coupling-based methods: it computes an SDOT map by considering the entire noise distribution and all empirical data points, resulting in optimal NDA for training FGMs.

**Table 1: A comparison between coupling methods and AlignFlow, a Noise-Data Alignment method.**
|  | Algorithm | Scale to large models | Entire noise distribution | Explicit OT optimization |
| --- | --- | --- | --- | --- |
|  | Tong et al. (2023) | ✗ | ✗ | ✓ |
| Coupling | Pooladian et al. (2023) | ✗ | ✗ | ✓ |
|  | Liu et al. (2022) | ✓ | ✗ | ✗ |
|  | Kornilov et al. (2024) | ✗ | ✗ | ✗ |
| Noise-Data Alignment | AlignFlow (ours) | ✓ | ✓ | ✓ |

## 3 Methodology

Mathematically, the generative modeling task can be formulated as follows: given a dataset $\{x_{1}^{i}\}_{i\in I}$, assumed to consist of i.i.d. samples from an unknown probability distribution $\tilde{p}_{1}$ on a space $\mathcal{X}$(^1^11In many settings, $\mathcal{X}$ is chosen as a latent space, with data obtained via encoding through a Variational Autoencoder (VAE).), where $I$ is the index set of the dataset (^2^22Without loss of generality, we assume $I=\{1,2,...,N\}$, where $N$ is the dataset size.). The objective is to generate new samples that follow the same distribution $\tilde{p}_{1}$.

#### Notation

In this paper, we establish the following notational conventions. We denote by $p_{0}$ the noise prior distribution (e.g., a standard Gaussian distribution) from which sampling is straightforward. Conversely, $p_{1}$ represents the empirical data distribution, defined as a Dirac mixture $p_{1}=\sum_{i\in I}b_{i}\delta_{x^{i}_{1}}$. In our FGM training setup, the weights are uniform, i.e., $b_{i}=1/|I|$.

### 3.1 Flow-based Generative Models (FGMs)

We begin by summarizing a general FGM framework in Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1). In each training iteration, a set of noise $x_{0}$, data $x_{1}$ and time $t$ are drawn. $x_{t}$ are computed as the interpolation between $x_{0}$ and $x_{1}$. Finally, a NN $u(x_{t},t;\theta)$ is trained to approximate the target vector field $\operatorname{TargetVectorField}$ and the NN parameters $\theta$ are updated accordingly.

Figure: Algorithm 1 Flow-based Generative Models (Training)

#### Examples:

(Vanilla) Flow Matching (Lipman et al., 2022) defines the target vector field as $\operatorname{TargetVectorField}(x_{0},x_{1})=x_{1}-x_{0}$ and uses the loss function $\operatorname{Loss}\left(\left\{\hat{v}^{j},v^{j}\right\}_{j=1}^{B}\right)=\frac{1}{B}\sum_{j}{\left\|\hat{v}^{j}-v^{j}\right\|}_{2}^{2}$. Several existing methods, including the shortcut model (Frans et al., 2025), MeanFlow (Geng et al., 2025), Consistency Training (Frans et al., 2025), Live Reflow Frans et al. (2025), conform to this general FGM framework. See Sec. [C.1](https://arxiv.org/html/2510.15038v1#A3.SS1) for more details.

In line [1](https://arxiv.org/html/2510.15038v1#algorithm1) and [1](https://arxiv.org/html/2510.15038v1#algorithm1) of Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1), noise and data points are sampled independently, and thus paired randomly, leading to curved trajectories intrinsically (Liu et al., 2022; Hertrich et al., 2025).

### 3.2 Coupling between noise and data

In fact, the loss function Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1) provides an empirical estimate of the following expectation (^3^33We here choose the loss function to be $\operatorname{Loss}\left(\left\{\hat{v}_{t}^{i},v_{t}^{i}\right\}_{i=1}^{B}\right):=\frac{1}{B}\sum_{i}{\left\|\hat{v}_{t}^{i}-v_{t}^{i}\right\|}_{p}^{p}$ for simplicity.):

$$ $L(\theta)\approx\mathcal{L}(\theta)=\mathbb{E}_{t\sim\text{Unif}[0,1],x_{0}\sim p_{0},x_{1}\sim p_{1}}{\left\|u(x_{t},t;\theta)-\operatorname{TargetVectorField}(x_{0},x_{1})\right\|}_{p}^{p}$ (1) $$

This formulation implies that $x_{0}$ and $x_{1}$ are sampled independently from $p_{0}$ and $p_{1}$, respectively, as used in Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1); that is $(x_{0},x_{1})\sim p_{0}\times p_{1}$. Recent works, such as (Pooladian et al., 2023; Liu et al., 2022; Tong et al., 2023), aim to construct more informative joint distributions by sampling $(x_{0},x_{1})$ from a coupling $\gamma\in\Gamma(p_{0},p_{1})$:

$$ $\mathcal{L}_{\gamma}(\theta)=\mathbb{E}_{t\sim\text{Unif}[0,1],(x_{0},x_{1})\sim\gamma}{\left\|u(x_{t},t;\theta)-\operatorname{TargetVectorField}(x_{0},x_{1})\right\|}_{p}^{p}$ (2) $$

where $\Gamma(p_{0},p_{1})$ denotes the set of all couplings (i.e., joint distributions) with marginals $p_{0}$ and $p_{1}$:

$$ $\Gamma:=\left\{\gamma\in\mathcal{P}(\mathcal{X}\times\mathcal{X}):\int\gamma(x_{0},x_{1})\,\mbox{d}x_{0}=p_{1}(x_{1})\ ~~\forall x_{1},\int\gamma(x_{0},x_{1})\,\mbox{d}x_{1}=p_{0}(x_{0})\ ~~\forall x_{0}\right\}$ (3) $$

As evident from the definition, $\Gamma$ is a vast set. Although training with any valid coupling theoretically yields a correct vector field, the straightness of the resulting trajectories and the efficiency of the training process can vary significantly depending on the choice of coupling. This naturally raises the question: which coupling should we choose? A growing body of work suggests that OT offers a principled way to guide this selection (Pooladian et al., 2023; Kornilov et al., 2024).

### 3.3 Semi-Discrete Optimal Transport

The Optimal Transport (OT) problem seeks to compute the optimal coupling between two probability distributions by minimizing a given cost function $c:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$, (see, e.g., Peyré et al. (2019) for a comprehensive overview):

$$ $\gamma_{*}:=\arg\min_{\gamma\in\Gamma(q_{1},q_{2})}\left(\int_{\mathcal{X}\times\mathcal{X}}c(y_{1},y_{2})\,\mbox{d}\gamma(y_{1},y_{2})\right),$ (4) $$

where we choose $c(y_{1},y_{2}):={\left\|y_{1}-y_{2}\right\|}^{2}$ throughout the paper. (^4^44The minimum objective is $W_{2}^{2}(q_{1},q_{2}):=\min_{\gamma\in\Gamma(q_{1},q_{2})}\left(\int_{\mathcal{X}\times\mathcal{X}}{\left\|y_{1}-y_{2}\right\|}^{2}\mbox{d}\gamma(y_{1},y_{2})\right)$.)

Figure: (a) Dataset and noise distribution
Refer to caption: x1.png

A discrete distribution contrasts with a continuous distribution (such as the normal distribution), in that the associated random variable assumes only a finite (or countable) set of values. For example, the empirical data distribution can be expressed as $p_{1}=\frac{1}{|I|}\sum_{i\in I}\delta_{x^{i}_{1}}$, where $\delta_{x^{i}_{1}}$ denotes the Dirac measure centered at data point $x^{i}_{1}$. An OT problem between a continuous and a discrete distribution is referred to as Semi-Discrete Optimal Transport (SDOT) (Peyré et al., 2019, Sec. 5), which will serve as our primary technique for aligning noise with data.

The transport map of an SDOT problem can be computed using a $|I|$-dimensional vector $\textbf{g}=[g_{i}]_{i\in I}$, referred to as the dual weight, where $|I|$ is the number of points in the discrete distribution $p_{1}$. Given the dual weight g, we compute $\varphi(\cdot;\textbf{g}):\mathcal{X}\to I$, which maps noise to data index:

$$ $\varphi(x_{0};\textbf{g}):=\arg\min_{i\in I}~~c(x_{0},x_{1}^{i})-g_{i}$ (5) $$

Given $\varphi(\cdot;\textbf{g}):\mathcal{X}\to I$, the SDOT map from the noise distribution to the discrete data points is directly obtained.

In fact, $\mathcal{X}$ is partitioned into cells by the SDOT map in the sense that cell $\mathsf{L}_{i}$ contains the points transported to the $i$-th point in the dataset. Each cell is referred to as the Laguerre cell, defined below:

$$ $\mathsf{L}_{i}(\textbf{g}):=\left\{x\in\mathcal{X}:c(x,y_{i})-g_{i}\leq c(x,y_{j})-g_{j},\forall j\right\}$ (6) $$

Fig. [1](https://arxiv.org/html/2510.15038v1#S3.F1) visualizes the partition in a 2-dim space.

We proceed with the computation of the dual weight g. SDOT, similar to general OT problems in Eq. ([4](https://arxiv.org/html/2510.15038v1#S3.E4)), is a minimization problem. By analyzing its dual problem, the dual weight can be solved by maximizing the following objective function utilizing the Laguerre cells (see e.g., Eq. 5.7 in Peyré et al. (2019)):

$$ $\mathcal{E}(\textbf{g}):=\sum_{i\in I}\int_{\mathsf{L}_{i}(g)}\left(c(x,y_{j})-g_{j}\right)\mbox{d}p_{0}(x)+\langle\textbf{g},\textbf{b}\rangle$ (7) $$

whose gradient is given by $\nabla\mathcal{E}(\textbf{g})_{i}=-\int_{\mathsf{L}_{i}(\textbf{g})}\,\mbox{d}p_{0}+b_{i}$,
where $b_{i}$ is the probability for each data point. In our cases, $b_{i}\equiv\frac{1}{|I|}$.

To solve this maximization problem, we employ the Adam optimizer (Kingma & Ba, 2014). While this approach to the SDOT problem is not novel (e.g., Peyré et al. (2019)),
we introduce a new and efficient Exponential Moving Average (EMA)-based estimation method for computing the $\operatorname{MRE}$ and the $\operatorname{L_{1}}$ distance (further discussed in Sec. [A](https://arxiv.org/html/2510.15038v1#A1)) to evaluate the quality of the resulting dual weights. This estimation facilitates hyperparameter tuning.

Figure: Algorithm 2 SDOT dual weight optimization

#### Computational cost of Algo. 2

: The computation of SDOT dual weight has a complexity of $\mathcal{O}(|I|^{3})$ when $\epsilon=0$ (Liu et al., 2021), and is known can be further reduced by setting the entropic regularization $\epsilon>0$, as a positive $\epsilon$ improves the smoothness of the SDOT objective. Specifically, the outer loop (line [2](https://arxiv.org/html/2510.15038v1#algorithm2)) in Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) haa a per iteration cost of $\mathcal{O}(|I|)$, and the SDOT objective error decreases in the order $\mathcal{O}(1/k)$ (Taşkesen et al., 2023).

### 3.4 The proposed AlignFlow

OT optimization remains a core challenge in previous OT-based FGMs. In this subsection, we propose using SDOT for Noise-Data Alignment (NDA) in training FGMs.

#### Idea behind AlignFlow

The critical insight behind AlignFlow is to use the SDOT to compute a transport map between the noise distribution $p_{0}$ and the known empirical data distribution $p_{1}$. This is because 1)
SDOT considers all empirical data points during optimization, rather than relying on samples
drawn from both the source and target distributions as required by discrete OT; 2) the SDOT objective is convex, and the time complexity for optimizing SDOT is $\mathcal{O}(|I|^{3})$ (Liu et al., 2021); and 3) the optimality of the resulting SDOT map can be estimated with sample complexity of $\mathcal{O}(|I|)$. Importantly, the SDOT map between $p_{0}$ and $p_{1}$ preserves the desirable properties of OT, such as inducing straight transport trajectories. Therefore, unlike previous OT-based FGMs, AlignFlow solves the SDOT map with a convergence guarantee, leading to an explicit coupling of noise and data.

Once the SDOT map is obtained, it can be integrated into the general FGM framework. This leads to the development of AlignFlow (Algo. [3](https://arxiv.org/html/2510.15038v1#algorithm3)), which leverages the SDOT map to align noise with data. Notably, Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) incorporates all data points during optimization and inherently avoids the curse of dimensionality.

Figure: Algorithm 3 AlignFlow: Noise-Data Alignment using SDOT (Training)

### 3.5 Additional techniques

###### Remark 1 (Noise storage) .

In lines [3](https://arxiv.org/html/2510.15038v1#algorithm3) and [3](https://arxiv.org/html/2510.15038v1#algorithm3) of Algo. [3](https://arxiv.org/html/2510.15038v1#algorithm3), a large number of noise samples need to be generated prior to training an FGM. Storing all these samples in memory, or even on disk, quickly becomes impractical. (^5^55Based on our experiments, storing 10 epochs worth of noise samples for ImageNet training in latent space (see Sec. [5.3](https://arxiv.org/html/2510.15038v1#S5.SS3) and [5.2](https://arxiv.org/html/2510.15038v1#S5.SS2)) would require terabytes of disk space. In addition to storage demands, I/O throughput would pose a significant bottleneck.) To address this, we only store the random seed used to generate each noise sample, so each noise-data pair is represented as a (seed, index) tuple.

Such an approach requires a deterministic mapping from the seed to random matrices (supported by JAX) and loading the entire ImageNet latent representations into memory for efficient random access (automatically optimized by the PyTorch dataloader). Using this strategy, storing the (seed, index) pairs for 500 epochs of ImageNet training requires only about 1 GB of disk space.

###### Remark 2 (Data augmentation) .

Data augmentation is a critical component for achieving optimal performance in image-related tasks. However, incorporating complex augmentation techniques, such as random cropping or rotation, directly into the SDOT map formulation can be challenging.

Fortunately, for most image generation tasks, effective augmentation is often limited to random horizontal flips. This particular case can be handled elegantly without requiring complex modifications to the SDOT framework: simply redefine the dataset as the union of the original images and their horizontally flipped versions.

###### Remark 3 (Class-conditioned generation) .

For class-conditioned tasks, as discussed in Sec. [5.2](https://arxiv.org/html/2510.15038v1#S5.SS2) and [5.3](https://arxiv.org/html/2510.15038v1#S5.SS3), we denote the class-conditioned data distribution by $p_{1,c}$ for the $c$-th class. The procedure involves computing the SDOT map between the noise distribution $p_{0}$ and each class-specific distribution $p_{1,c}$. Once noise samples are drawn, AlignFlow generates corresponding (noise, data) pairs and performs the required Rebalance (Sec. [B](https://arxiv.org/html/2510.15038v1#A2)) operation independently for each class.

###### Remark 1 (Noise storage) .

###### Remark 2 (Data augmentation) .

###### Remark 3 (Class-conditioned generation) .

## 4 Advantages of AlignFlow

Inspired by the properties of the SDOT map, we introduce the concept of Noise-Data Alignment (NDA):

###### Definition 1 (Noise-Data Alignment) .

Given a dataset and a noise distribution, noise-data alignment is a deterministic map that maps noise to a data point in the dataset.

A NDA is a deterministic coupling between noise and data, and we highlight determinism as the key distinction from existing coupling methods, based on the following intuitions:

- •
A deterministic coupling ensures consistent noise-data matching regardless of batch size in training FGMs.
- •
If $q_{1}$ in Eq. [4](https://arxiv.org/html/2510.15038v1#S3.E4) is continuous, then the OT map from $q_{1}$ to $q_{2}$ must be deterministic (see e.g., Peyré et al., 2019, Remark 2.24), i.e., determinism is a necessary condition for optimality.

###### Definition 1 (Noise-Data Alignment) .

### 4.1 Bypass the Curse of dimensionality

The sample complexity of OT measures the accuracy of an OT plan estimated by the discrete-OT plan between samples. However, the following theorem shows the number of sample required is exponentially dependent on the dimension:

###### Theorem 1 (Sample complexity of OT (Informal version for Thm. 1 in Fournier & Guillin ( 2015 ) )) .

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}$.

The squared Wasserstein distance $W_{2}^{2}$ corresponds to the minimum value in Eq. ([4](https://arxiv.org/html/2510.15038v1#S3.E4)). Thm. [1](https://arxiv.org/html/2510.15038v1#Thmtheorem1) shows that even estimating an OT plan between a continuous distribution and its samples can be challenging, requiring an exponential number of samples. Existing algorithms may suffer from this curse of dimensionality in two cases, potentially both at once: (1) when using samples from the noise distribution to approximate the underlying continuous noise distribution; and (2) when using the dataset or a batch from the dataset to approximate the unknown real data distribution. Many works (implicitly) assume that the empirical distribution approximates the continuous one, e.g., Assumption A2 and A3 in Pooladian et al. (2023), Thm.1 in Kornilov et al. (2024).

As a NDA approach, AlignFlow circumvents the curse of dimensionality by computing a SDOT map between $p_{0}$ and $p_{1}$. Since $p_{1}$ is fully specified by the dataset, the SDOT map can, in principle, be computed with zero estimation error, thereby avoiding the curse of dimensionality.

###### Theorem 1 (Sample complexity of OT (Informal version for Thm. 1 in Fournier & Guillin ( 2015 ) )) .

### 4.2 Deterministic alignment

The SDOT map is theoretically fully deterministic: a given sample from the noise distribution is consistently mapped to a fixed data point.

Intuitively, this determinism provides a significant advantage in convergence speed: to determine the target vector field $v_{t}$ at some $t$ and $x_{t}$, the standard approach using the random coupling (as in Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1)) requires iterating over the entire dataset:

$$ $u(x_{t},t;\theta)=\mathbb{E}_{x_{1}\sim p_{1}}\operatorname{TargetVectorField}(x_{0},x_{1})p_{0}(x_{0}),\quad x_{0}=x_{1}-(x_{1}-x_{t})/t$ (8) $$

However, the fixed coupling in AlignFlow bypasses this estimation process, directly providing the target vector field for the NN to learn as:

$$ $u(x_{t},t;\theta)=\operatorname{TargetVectorField}(x_{0},x_{1})p_{0}(x_{0}),\quad x_{0}=x_{1}-(x_{1}-x_{t})/t,~~x_{1}=\varphi(x_{0})$ (9) $$

This crucial difference demonstrates that fixed coupling significantly simplifies the estimation of the target vector field, thereby leading to the accelerated convergence observed empirically with AlignFlow.

### 4.3 Low computational cost

AlignFlow directly computes the SDOT map in Stage 1. In contrast to prior methods that estimate the OT plan indirectly through batch-based approaches (e.g., Reflow operation in Liu et al. (2022), ICNN in Kornilov et al. (2024), and Minibatch OT in Tong et al. (2023)). AlignFlow’s direct computation of SDOT is both more accurate and efficient, and Stage 1 adds negligible overhead empirically, accounting for less than 1% additional training time. Further implementation details are provided in Sec. [A](https://arxiv.org/html/2510.15038v1#A1).

Upon completion of Stage 1, the SDOT map is fully computed. Consequently, the only additional overhead in Stage 2 arises from the generating the training noise-data pairs (Lines [3](https://arxiv.org/html/2510.15038v1#algorithm3) and [3](https://arxiv.org/html/2510.15038v1#algorithm3) in Algo. [3](https://arxiv.org/html/2510.15038v1#algorithm3)). This process is highly efficient, runs quickly on modern GPUs, and incurs an almost negligible cost, typically less than 0.1% of the total training time.

## 5 Experiments

Figure: (a) CIFAR-10 results using U-Net, generated by the adaptive ODE integrator DOPRI5 (see Sec. [5.1](https://arxiv.org/html/2510.15038v1#S5.SS1)).
Refer to caption: x3.png

### 5.1 Unconditional Image Generation on CIFAR-10 with U-Net

Following (Tong et al., 2023, Sec. 5.3), we train a unconditional U-Net (Ronneberger et al., 2015) on the CIFAR-10 dataset. In this setup, the FGMs are trained directly in the pixel space. The comparative training curves are shown in Fig. [2(a)](https://arxiv.org/html/2510.15038v1#S5.F2.sf1), and the FID-50k scores for various ODE integrators are detailed in Table [2](https://arxiv.org/html/2510.15038v1#S5.T2). Compared to the coupling estimated via the standard Minibatch OT (Tong et al., 2023), AlignFlow demonstrates faster convergence and achieves better FID scores across all tested ODE integrators. All experiments were performed using the official code released by Tong et al. (2023).

**Table 2: Comparison of FID-50k scores between Minibatch OT and AlignFlow for U-Net trained on CIFAR-10, evaluated across different ODE integrators. Results are averaged over 5 independent runs. AlignFlow consistently achieves better performance than Minibatch OT under all tested ODE integrators.**
|  | Euler (100 steps) | Euler (1000 steps) | DOPRI5 |
| --- | --- | --- | --- |
| Minibatch OT (Tong et al., 2023) | 4.80 | 3.92 | 3.82 |
| AlignFlow (ours) | 4.72 | 3.79 | 3.71 |

### 5.2 ImageNet256 Generation Using DiT with Shortcut Model

AlignFlow can be seamlessly integrated with modern SOTA network architectures and scales effectively to large-scale datasets. We train AlignFlow in conjunction with Diffusion Transformers (DiT, Peebles & Xie (2023)) on the class-conditioned ImageNet dataset at $256\times 256$ resolution (ImageNet256). The FGMs operate in the latent space of size $28\times 28\times 4$, obtained from a pretrained VAE. All model hyperparameters are adopted directly from (Frans et al., 2025, Tables 1 and 3) without any modification or tuning. A comparison of training curves for the shortcut model with and without AlignFlow is shown in Fig. [2(b)](https://arxiv.org/html/2510.15038v1#S5.F2.sf2), while improvements across additional models using AlignFlow are reported in Table [3](https://arxiv.org/html/2510.15038v1#S5.T3).

**Table 3: Evaluation of AlignFlow on DiT-B/2 for ImageNet256 using FID-50k demonstrates consistent performance improvements across all tested NFE configurations.**
| Algorithm | AlignFlow? | NFE=4 | Difference | NFE=1 | Difference |
| --- | --- | --- | --- | --- | --- |
| Flow Matching | ✓ | 93.16 | $\downarrow$ 32.46 | 276.18 | $\downarrow$ 28.86 |
| ✗ | 125.62 |  | 305.04 |  |  |
| Consistency Training | ✓ | 103.14 | $\downarrow$ 8.70 | 64.33 | $\downarrow$ 12.04 |
| ✗ | 111.84 |  | 76.37 |  |  |
| Live Reflow (Frans et al., 2025) | ✓ | 60.23 | $\downarrow$ 34.52 | 47.06 | $\downarrow$ 12.81 |
| ✗ | 94.75 |  | 59.87 |  |  |
| Shortcut Models (Frans et al., 2025) | ✓ | 30.31 | $\downarrow$ 2.80 | 43.92 | $\downarrow$ 2.73 |
| ✗ | 33.11 |  | 46.65 |  |  |

### 5.3 ImageNet256 Generation using SiT with MeanFlow

AlignFlow further enhances the one-step generative model MeanFlow (Geng et al., 2025), which is based on Scalable Interpolant Transformers (SiT, Ma et al. (2024)) and trained on class-conditioned ImageNet256. The FGMs operate in the latent space of size $28\times 28\times 4$, obtained from a pretrained VAE. For our experiments, we use a non-official PyTorch implementation (Zhu, 2025), which reliably reproduces the reported results on GPU. All hyperparameters are kept identical to the official configuration (Geng et al., 2025, Section A), without any additional tuning. The training curve is shown in Fig. [2(c)](https://arxiv.org/html/2510.15038v1#S5.F2.sf3), and the corresponding FID scores are reported in Table [4](https://arxiv.org/html/2510.15038v1#S5.T4). AlignFlow consistently improves both performance and convergence speed across all cases, demonstrating its scalability to large models. Example generated images are provided in Fig. [4](https://arxiv.org/html/2510.15038v1#A6.F4) in the appendix.

**Table 4: We evaluate AlignFlow on ImageNet256 using MeanFlow (NFE=1), showing consistent performance improvements across all model sizes.**
| Backbone | # params | w/ AlignFlow | w/o AlignFlow | Difference |
| --- | --- | --- | --- | --- |
| SiT-B/4 | 131M | 13.75 | 15.53 | $\downarrow$ 1.78 |
| SiT-B/2 | 131M | 5.60 | 6.17 | $\downarrow$ 0.57 |
| SiT-L/2 | 459M | 3.51 | 3.84 | $\downarrow$ 0.33 |
| SiT-XL/2 | 676M | 3.23 | 3.43 | $\downarrow$ 0.20 |

## References

- Achiam et al. (2023)
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Albergo et al. (2023)
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden.
Stochastic interpolants: A unifying framework for flows and diffusions.
*arXiv preprint arXiv:2303.08797*, 2023.
- Altschuler et al. (2022)
Jason M Altschuler, Jonathan Niles-Weed, and Austin J Stromme.
Asymptotics for semidiscrete entropic optimal transport.
*SIAM Journal on Mathematical Analysis*, 54(2):1718–1741, 2022.
- Amos et al. (2017)
Brandon Amos, Lei Xu, and J Zico Kolter.
Input convex neural networks.
In *International conference on machine learning*, pp. 146–155. PMLR, 2017.
- Boffi et al. (2025)
Nicholas Matthew Boffi, Michael Samuel Albergo, and Eric Vanden-Eijnden.
Flow map matching with stochastic interpolants: A mathematical framework for consistency models.
*Transactions on Machine Learning Research*, 2025.
- Chen et al. (2018)
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.
Neural ordinary differential equations.
*Advances in neural information processing systems*, 31, 2018.
- Cheng & Schwing (2025)
H. K. Cheng and A. G. Schwing.
The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation.
In *Proc. ICCV*, 2025.
- Dao et al. (2025)
Quan Dao, Hao Phung, Trung Tuan Dao, Dimitris N Metaxas, and Anh Tran.
Self-corrected flow distillation for consistent one-step and few-step image generation.
In *Proceedings of the AAAI Conference on Artificial Intelligence*, volume 39, pp. 2654–2662, 2025.
- Dudley (1969)
Richard Mansfield Dudley.
The speed of mean glivenko-cantelli convergence.
*The Annals of Mathematical Statistics*, 40(1):40–50, 1969.
- Esser et al. (2024)
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al.
Scaling rectified flow transformers for high-resolution image synthesis.
In *Forty-first international conference on machine learning*, 2024.
- Fournier & Guillin (2015)
Nicolas Fournier and Arnaud Guillin.
On the rate of convergence in wasserstein distance of the empirical measure.
*Probability theory and related fields*, 162(3):707–738, 2015.
- Frans et al. (2025)
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel.
One step diffusion via shortcut models.
*The Thirteenth International Conference on Learning Representations*, 2025.
- Genevay et al. (2019)
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré.
Sample complexity of sinkhorn divergences.
In *The 22nd international conference on artificial intelligence and statistics*, pp. 1574–1583. PMLR, 2019.
- Geng et al. (2025)
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He.
Mean flows for one-step generative modeling.
*arXiv preprint arXiv:2505.13447*, 2025.
- Goodfellow et al. (2020)
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.
Generative adversarial networks.
*Communications of the ACM*, 63(11):139–144, 2020.
- Hertrich et al. (2025)
Johannes Hertrich, Antonin Chambolle, and Julie Delon.
On the relation between rectified flows and optimal transport.
*arXiv preprint arXiv:2505.19712*, 2025.
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Kim et al. (2023)
Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon.
Consistency trajectory models: Learning probability flow ode trajectory of diffusion.
*arXiv preprint arXiv:2310.02279*, 2023.
- Kingma & Ba (2014)
Diederik P Kingma and Jimmy Ba.
Adam: A method for stochastic optimization.
*arXiv preprint arXiv:1412.6980*, 2014.
- Kornilov et al. (2024)
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Aleksandr Korotin.
Optimal flow matching: Learning straight trajectories in just one step.
*Advances in Neural Information Processing Systems*, 37:104180–104204, 2024.
- Lipman et al. (2022)
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
*arXiv preprint arXiv:2210.02747*, 2022.
- Liu et al. (2021)
Huidong Liu, Ke Ma, Lei Zhou, and Dimitris Samaras.
Efficient semi-discrete optimal transport using the maximum relative error between distributions.
2021.
- Liu et al. (2022)
Xingchao Liu, Chengyue Gong, and Qiang Liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
*arXiv preprint arXiv:2209.03003*, 2022.
- Ma et al. (2024)
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie.
Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers.
In *European Conference on Computer Vision*, pp. 23–40. Springer, 2024.
- Nguyen et al. (2022)
Khai Nguyen, Dang Nguyen, Tung Pham, and Nhat Ho.
Improving mini-batch optimal transport via partial transportation.
In *Proceedings of the 39th International Conference on Machine Learning*, 2022.
- Peebles & Xie (2023)
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pp. 4195–4205, 2023.
- Peyré et al. (2019)
Gabriel Peyré, Marco Cuturi, et al.
Computational optimal transport: With applications to data science.
*Foundations and Trends® in Machine Learning*, 11(5-6):355–607, 2019.
- Pooladian et al. (2023)
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen.
Multisample flow matching: Straightening flows with minibatch couplings.
*arXiv preprint arXiv:2304.14772*, 2023.
- Radford et al. (2018)
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al.
Improving language understanding by generative pre-training.
2018.
- Rezende & Mohamed (2015)
Danilo Rezende and Shakir Mohamed.
Variational inference with normalizing flows.
In *International conference on machine learning*, pp. 1530–1538. PMLR, 2015.
- Rombach et al. (2022)
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 10684–10695, 2022.
- Ronneberger et al. (2015)
Olaf Ronneberger, Philipp Fischer, and Thomas Brox.
U-net: Convolutional networks for biomedical image segmentation.
In *International Conference on Medical image computing and computer-assisted intervention*, pp. 234–241. Springer, 2015.
- Salimans & Ho (2022)
Tim Salimans and Jonathan Ho.
Progressive distillation for fast sampling of diffusion models.
*arXiv preprint arXiv:2202.00512*, 2022.
- Sohl-Dickstein et al. (2015)
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In *International conference on machine learning*, pp. 2256–2265. PMLR, 2015.
- Song et al. (2021)
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
*ICLR*, 2021.
- Song et al. (2023)
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.
Consistency models.
2023.
- Taşkesen et al. (2023)
Bahar Taşkesen, Soroosh Shafieezadeh-Abadeh, and Daniel Kuhn.
Semi-discrete optimal transport: Hardness, regularization and numerical solution.
*Mathematical Programming*, 199(1):1033–1106, 2023.
- Tong et al. (2023)
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio.
Improving and generalizing flow-based generative models with minibatch optimal transport.
*arXiv preprint arXiv:2302.00482*, 2023.
- Yang et al. (2024)
Ling Yang, Zixiang Zhang, Zhilong Zhang, Xingchao Liu, Minkai Xu, Wentao Zhang, Chenlin Meng, Stefano Ermon, and Bin Cui.
Consistency flow matching: Defining straight flows with velocity consistency.
*arXiv preprint arXiv:2407.02398*, 2024.
- Zhu (2025)
Yu Zhu.
Meanflow: Pytorch implementation.
[https://github.com/zhuyu-cs/MeanFlow](https://github.com/zhuyu-cs/MeanFlow), 2025.
PyTorch implementation of Mean Flows for One-step Generative Modeling.

## Appendix A Efficient SDOT algorithm for large-scale datasets

### A.1 Performance Metric for Algo. 2

We revisit the definitions of Maximum Relative Error ($\operatorname{MRE}$) and $\operatorname{L_{1}}$ distance as proposed in (Liu et al., 2021), which serve as metrics for evaluating the quality of the computed dual weights:

$$ $\operatorname{MRE}=\max_{i\in I}\frac{{\left|p_{i}-b_{i}\right|}}{b_{i}},\quad\operatorname{L_{1}}(\textbf{g})=\sum_{i\in I}{\left|p_{i}-b_{i}\right|},\quad p_{i}:=\int\mathbb{L}_{i}(\textbf{g})\,\mbox{d}p_{0}$ (10) $$

In Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2), they can be estimated efficiently by $\widetilde{\operatorname{MRE}}={\left\|\nabla\mathcal{E}_{\texttt{ema}}\right\|}_{\infty}$ and $\widetilde{\operatorname{L_{1}}}={\left\|\nabla\mathcal{E}_{\texttt{ema}}\right\|}_{1}$.

Mathematically, $\operatorname{MRE}$ and $\operatorname{L_{1}}$ serve as estimators for the $\ell_{\infty}$-norm and $\ell_{1}$-norm of the gradient of the objective, specifically ${\left\|\nabla\mathcal{E}\right\|}_{\infty}$ and ${\left\|\nabla\mathcal{E}\right\|}_{1}$, respectively. The critical importance of $\operatorname{MRE}$ lies in its ability to quantify the imbalance of the transported probabilities. Importantly, an outcome of $\operatorname{MRE}<1$ guarantees that each target data point receives non-zero probability transported from the source distribution. $\operatorname{MRE}=0$ signifies that the SDOT solution is perfectly optimized: the objective function in Eq. ([4](https://arxiv.org/html/2510.15038v1#S3.E4)) is minimized, and, crucially, the measure preservation constraint is satisfied uniformly, ensuring each target receives an equal probability mass. (See also Sec. [B](https://arxiv.org/html/2510.15038v1#A2) for details on rebalancing.)

### A.2 Entropic regularization for SDOT

Discrete OT problems are intrinsically non-smooth. A possible approach to mitigate this is the introduction of regularization terms to induce a smoother objective landscape. For SDOT problems, similar techniques can also be applied. Instead of solving the problem Eq. ([4](https://arxiv.org/html/2510.15038v1#S3.E4)), SDOT with entropic regularization is solving (Altschuler et al., 2022)

$$ $\min_{\gamma\in\Gamma(p_{0},p_{1})}\int_{\mathcal{X}\times\mathcal{X}}c(x_{0},x_{1})\mbox{d}\gamma(x_{0},x_{1})+\epsilon\operatorname{KL}(\gamma||p_{0}\otimes p_{1})$ (11) $$

where $\epsilon>0$ is the regularization strength and $\operatorname{KL}$ denotes the Kullback-Leibler divergence. While the inclusion of this auxiliary term introduces a controlled bias into the solution, it significantly enhances the smoothness and tractability of the optimization landscape.

### A.3 Hyperparameters tuning for SDOT optimization

Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) involves the tuning of three primary hyperparameters: the entropic regularization strength $\epsilon$, the EMA parameter $\beta$, and the Adam optimizer’s learning rate $lr$. Throughout the maximization process in Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2), the $\operatorname{L_{1}}$ metric is expected to exhibit a continuous decay when the hyperparameters are set appropriately.

- •
Entropic regularization strength $\epsilon$ dictates the fundamental trade-off between bias and optimization difficulty. A larger $\epsilon$ induces a greater bias in the resulting transport plan, while setting $\epsilon$ to zero or a very small value significantly increases the non-smoothness and complexity of the optimization landscape.
- •
The parameter $\epsilon$ should remain fixed throughout the entire optimization procedure. When optimizing the problem, consider increasing the batch size and/or decreasing the learning rate when $\operatorname{L_{1}}$ plateaus.
- •
The optimization stops when $\operatorname{MRE}$ satisfies the required performance threshold. For robust performance in downstream tasks, such as FGMs, we recommend maintaining $\operatorname{MRE}$ below 0.2.
- •
Contrary to common practices in standard deep learning optimization, where learning rates around $0.001$ are preferred, the computation of the SDOT map (Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2)) often requires a relatively high initial learning rate. A value of $lr=10$ serves as a recommended starting point for tuning.
- •
For the management of large datasets, it is advisable to increase the batch size and/or adjust the EMA parameter $\beta$ closer to 1 (e.g., from $0.99$ to $0.999$) to ensure a smoother and more stable gradient update over a larger number of iterations.

### A.4 Hyperparameters and computational cost

On CIFAR-10, we use the training set for training (50000 images with shape $32\times 32\times 3$).

- •
Normalize the whole dataset with mean = $(0.5,0.5,0.5)$ and std = $(0.5,0.5,0.5)$.
- •
Concatenate the dataset with the augmented (horizontally flipped) dataset.
- •
Compute the SDOT map with Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) and hyperparameters in Table [5](https://arxiv.org/html/2510.15038v1#A1.T5)

**Table 5: SDOT hyperparameters for CIFAR-10 (unconditional) generation. Computation takes 8 minutes and 30 seconds on an NVIDIA L40S GPU.**
| # step | learning rate | batch size | EMA parameter $\beta$ | entropic reg $\epsilon$ | $\operatorname{MRE}$ | $\operatorname{L_{1}}$ |
| --- | --- | --- | --- | --- | --- | --- |
| 0-1000 | 10 | 1024 | 0.99 | 1 | 3.4 | 0.29 |
| 1000-6000 | 0.1 | 4096 | 0.999 | 1 | 0.27 | 0.045 |
| 6000-11000 | 0.1 | 16384 | 0.999 | 0.01 | 0.11 | 0.019 |

On ImageNet, we use the full training set consisting of 1,281,167 images across 1,000 classes. In both the shortcut model w/ AlignFlow (Sec. [5.2](https://arxiv.org/html/2510.15038v1#S5.SS2)) and MeanFlow w/ AlignFlow (Sec. [5.3](https://arxiv.org/html/2510.15038v1#S5.SS3)), the SDOT map is computed in the latent space of size
$28\times 28\times 4$. Each image is augmented via horizontal flipping, resulting in approximately 2,600 images per class after augmentation. In our experiments, the SDOT map computation on the ImageNet dataset is cheaper than on the CIFAR-10 dataset.

**Table 6: SDOT hyperparameters for ImageNet256 (class-conditioned, latent space) generation. Computation takes less than 10 seconds per class on an NVIDIA L40S GPU.**
| # step | learning rate | batch size | EMA parameter $\beta$ | entropic reg $\epsilon$ | $\operatorname{MRE}$ | $\operatorname{L_{1}}$ |
| --- | --- | --- | --- | --- | --- | --- |
| 3000 | 10 | 4096 | 0.99 | 0.01 | $\sim$ 0.08 | $\sim$ 0.016 |

## Appendix B Rebalance: handling non-perfect SDOT map

Algo. [3](https://arxiv.org/html/2510.15038v1#algorithm3) deviates from the standard practice of sampling i.i.d. data pairs. Instead, it only samples noise $x_{0}\sim p_{0}$ and subsequently feeds the NN with data generated by the SDOT map
(Line [3](https://arxiv.org/html/2510.15038v1#algorithm3)).

This specific sampling scheme introduces a potential issue: if the SDOT map is not computed perfectly, the data $\varphi(x_{0})$ fed to the NN will be biased. Specifically, when the gradient of the energy $\nabla\mathcal{E}$ is non-zero, the set of mapped samples, represented by the set of indices $\{\varphi(x_{0})\mid x_{0}\overset{\text{i.i.d.}}{\sim}p_{0}\}$ becomes biased from the empirical data distribution $p_{1}$. (See also $\operatorname{MRE}$ defined in Eq. ([10](https://arxiv.org/html/2510.15038v1#A1.E10)) in Sec. [A](https://arxiv.org/html/2510.15038v1#A1)). This bias means the model will be trained with an uneven exposure to the dataset, with some data points being seen more frequently than others.

To mitigate the effects of imperfect SDOT convergence, a particular concern for large, non-conditional datasets where running Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) to $\operatorname{MRE}\to 0$ is computationally prohibitive, we introduce the rebalance operation in Line [3](https://arxiv.org/html/2510.15038v1#algorithm3), defined as:

$$ $\operatorname{rebalance}\left(\{m_{j}\}_{j=1}^{M}\right):=\arg\max_{\left\{\tilde{m}_{j}\right\}}\left\{\sum_{j}\mathbbm{1}(\tilde{m}_{j}=m_{j}):{\left|\max_{i\in\mathcal{I}}\sum_{j=1}^{M}\mathbbm{1}_{i}(\tilde{m}_{j})-\min_{i\in\mathcal{I}}\sum_{j=1}^{M}\mathbbm{1}_{i}(\tilde{m}_{j})\right|}\leq 1\right\}$ (12) $$

Intuitively, the rebalance operation determines the minimal perturbation to the set of target indices $\{m_{j}\}$ such that the empirical frequency of each data point in the modified set $\{\tilde{m}_{j}\}$ is nearly uniform. This operation forces the model to be exposed to an effectively unbiased dataset during the FGM training phase, even if the SDOT map is not fully converged. This comes at the cost of introducing a minimal amount of controlled randomness in the noise-data coupling.

In our experiments, particularly on the CIFAR-10 dataset (Sec. [5.1](https://arxiv.org/html/2510.15038v1#S5.SS1)), the rebalance operation did not yield a noticeable difference in performance. This is likely because Algo. [2](https://arxiv.org/html/2510.15038v1#algorithm2) achieved a sufficiently high degree of convergence (evidenced by the fact that over $85\%$ of the data assignments remained unchanged after the rebalance operation). Nonetheless, the rebalance mechanism is applied consistently across all our experiments to ensure that the data fed into the FGMs is unskewed.

## Appendix C More details for Flow-Based Generative Models

### C.1 More examples for FGM framework in Algo. 1

#### Shortcut model

In this model, an auxiliary input $d$ is introduced to the NN, i.e., $u=u(x,t,d;\theta)$ and $d^{j}$ is i.i.d. sampled from $\mathcal{D}(\cdot|t^{j})$. Given hyperparameter $\kappa$, choose $\operatorname{TargetVectorField}(x_{0}^{j},x_{1}^{j})=x_{1}^{j}-x_{0}^{j}$ for $j=1,...,\kappa$, and $\operatorname{TargetVectorField}(x_{0}^{j},x_{1}^{j})=\operatorname{StopGrad}(s_{t}^{j}+s_{t+d}^{j})$ for $j=\kappa+1,...,B$, where $s_{t}^{j}:=u(x_{t}^{j},t^{j},d^{j})$, $x_{t+d}^{j}:=x_{t}^{j}+s_{t}d^{j}$, $s_{t+d}^{j}:=u(x_{t^{j}+d^{j}}^{j},t^{j}+d^{j},d^{j})$. Together with $\operatorname{Loss}\left(\left\{\hat{v}^{j},v^{j}\right\}_{j=1}^{B}\right):=\frac{1}{B}\sum_{j}{\left\|\hat{v}^{j}-v^{j}\right\|}_{2}^{2}$, Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1) recovers the shortcut model in Frans et al. (2025).

#### MeanFlow

In this model, an extra input $r$ is introduced to the NN, i.e., $u=u(x,t,r;\theta)$. By choosing $\operatorname{TargetVectorField}(x_{0},x_{1})=\operatorname{StopGrad}(v_{t}-(t-r)v_{t}\partial_{x}u+\partial_{t}u)$ and $\operatorname{Loss}\left(\left\{\hat{v}^{i},v^{i}\right\}_{i=1}^{B}\right):=\frac{1}{B}\sum_{i}{\left\|\hat{v}^{i}-v^{i}\right\|}_{p}^{p}$, Algo. [1](https://arxiv.org/html/2510.15038v1#algorithm1) recovers the MeanFlow in Geng et al. (2025).

### C.2 Sampling / inference procedure of FGMs

To better elucidate the advantages conferred by the specially-designed target vector field, we review the inference procedure for FGM, as outlined in Algo. [4](https://arxiv.org/html/2510.15038v1#algorithm4). The inference process entails integrating the ODE $\partial_{t}x_{t}=u(x_{t},t;\theta)$ over the time interval $t\in[0,1]$, where the initial condition $x_{0}$ is sampled from the noise distribution $p_{0}$, and the final state $x_{1}$ constitutes a new data sample. While various ODE integrators may be employed, the computational difficulty of this integration is fundamentally governed by the complexity of the learned vector field $u(x,t;\theta)$. Intuitively, consider two hypothetical trajectories: an almost straight and a highly complex, curvilinear path. The almost straight path permits straightforward computation, possibly achievable in a single step via the forward Euler method, $x_{1}=x_{0}+u(x_{0},0;\theta)$. Conversely, the complex path necessitates a high-order or small-step ODE solver. This intuition highlights a key insight: the straightness of the integrated trajectory is paramount. A straighter path implies easier numerical integration, resulting in a lower NFE of the NN (Liu et al., 2022). Algorithms trained with a carefully-designed target vector field benefit directly from this enhanced straightness. For instance, the MeanFlow (Geng et al., 2025) is capable of generating high-quality samples with an NFE as low as $1$, whereas the vanilla Flow Matching (Lipman et al., 2022) typically requires over $100$ NFEs to achieve comparable results.

Figure: Algorithm 4 FGM / AlignFlow (Sampling)

## Appendix D Synthetic experiment: Checkerboard

In this section, we analyze and visualize the characteristics of different learned trajectories by training the FGMs algorithm on a synthetic two-dimensional data distribution. Following the precedent set by (Lipman et al., 2022, Fig. 4), the target data distribution is defined as the checkerboard distribution in $[-2,2]\times[-2,2]$. The noise distribution $p_{0}$ is the standard normal distribution.

Crucially, our experimental setting differs from prior work (Lipman et al., 2022; Tong et al., 2023), which assumed access to an infinite data stream during training.(^6^66In the checkerboard experiments of Lipman et al. (2022); Tong et al. (2023), fresh training data is sampled for every minibatch.) In our experiments, we sample data from the checkerboard distribution, fix the samples, and use them as the training set. While this fixed-data approach may result in a learned distribution with less spatial smoothness compared to the infinite-data setting, it provides a more faithful simulation of real-world machine learning tasks where data availability is inherently limited.

Fig. [3](https://arxiv.org/html/2510.15038v1#A4.F3) illustrates the density evolution as time $t$ progresses from 0 to $1$, charting the transformation from the Gaussian distribution $p_{0}$ to the checkerboard distribution. The visualizations indicates that AlignFlow generates a significantly straighter transport trajectory when compared to trajectories learned via Minibatch OT and Vanilla Flow Matching, as depicted in the corresponding panels of Fig. [3](https://arxiv.org/html/2510.15038v1#A4.F3). All comparative experiments utilize identical hyperparameters to ensure a fair comparison of the learned vector fields.

Figure: (a) AlignFlow (Alog. [3](https://arxiv.org/html/2510.15038v1#algorithm3))
Refer to caption: img/checkerboard_sdotfm.jpg

## Appendix E Capability of generalization

A natural question arises from the observation that the SDOT map already constitutes a mapping from the source space to the target space ($\mathcal{X}\to\mathcal{X}$). Since the integrated vector field of FGM also provides an $\mathcal{X}\to\mathcal{X}$ transformation, why is the subsequent FGM training in Stage 2 still necessary? The fundamental limitation of the SDOT map is its inability to generalize: it is constrained to only map noise samples to the fixed points in the dataset.
Consequently, it cannot generate new data instances. Therefore, a powerful NN must be trained in Stage 2 of Algo. [3](https://arxiv.org/html/2510.15038v1#algorithm3) atop the SDOT alignment to acquire the necessary generalization capability. The observed generalization performance of standard FGMs suggests that AlignFlow’s generalization performance can be achieved from the inductive bias and regularization inherent in the NN used to learn the vector field.

## Appendix F Image samples generated by AlignFlow

Figure: Figure 4: Images generated by MeanFlow+AlignFlow trained on ImageNet256 (FID=3.23, NFE=1).
Refer to caption: img/more_imgs/000029.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: The authors emphasize that unlike for minibatch OT, for which the OT plan is subject to the curse of dimensionality, semi-discrete OT bypasses the curse of dimensionality. However, that is with respect to the OT plan between $p_0$ and the empirical distribution $\tilde{p}_1$. However, arguably the OT plan that we want to be close to is the one between $p_0$ and the population distribution $p_1$, and I assume semi-discrete OT is still curse wrt this one. Is that so? If so, what are the advantages of not being cursed for the plan between $p_0$ and the empirical distribution $\tilde{p}_1$?
soundness: 4
strengths: - The paper presents a new and relevant method to pair data and noise samples when training FGMs. The proposed method is similar in spirit to mini batch OT, but instead of computing the coupling between the empirical data distribution and an empirical version of a Gaussian distribution, it computes the coupling between the empirical data distribution and the Gaussian, using Semi-discrete OT.

- The experimental results show an improvement over the training performance of mini batch OT, and the proposed method is faster as well.
summary: The paper AlignFlow: Improving Flow-Based Generative Models with Semi-Discrete Optimal Transport proposes a method to enhance flow-based generative models (FGMs) by explicitly aligning noise and data samples through Semi-Discrete Optimal Transport (SDOT). Traditional FGMs sample noise and data independently, leading to curved generative trajectories and high computational cost. AlignFlow introduces a two-stage training process: first computing a deterministic SDOT map that partitions the noise space into Laguerre cells matched to data points, and then training the FGM using these fixed correspondences. This approach bypasses the curse of dimensionality inherent to standard Optimal Transport methods, scales efficiently to large datasets, and adds negligible computational overhead. Experiments on CIFAR-10 and ImageNet show that AlignFlow consistently improves generation quality and convergence speed across multiple state-of-the-art models—including Flow Matching, Shortcut Models, and MeanFlow—while maintaining compatibility as a plug-and-play module for modern generative frameworks.
weaknesses: As in minibatch OT, it is unclear how the proposed approach can be adapted to settings like text-to-image, where we have one or few samples per “class”/prompt, as the plan must have Gaussian marginal for each class/prompt, but this is statistically hard to enforce when there are few samples. Yet, most large scale FGM are text-to-“modality” instead of class-based. To my understanding, this limits the applicability of the method. Could the authors comment on this? Have they found a way to apply their method in text-to-“modality” settings?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. At the first glance the optimization problem for rebalacing in Equation 12 has a large search space. Could the authors provide more details on how this is implemented in practice?

Some typos:
1. Line 131. The sentence seems to miss a word.
2. Line 208. smaple $\rightarrow$ sample.
3. Line 398. I believe Semi-Discrete Optimal Transport was meant instead of Stochastic Denoising Optimal Transport.
soundness: 3
strengths: The paper is well-written. The motivation is clear. The proposed approach is simple and can be used as a plug-and-play method to improve flow-based models. AlignFlow consistently improves sampling efficiency and convergence over the baselines.
summary: The paper proposes AlignFlow, a method that calculates semi-discrete optimal transport (SDOT) between the source and the data distributions prior to the training of a flow-based generative model. It further uses the calculated SDOT plan as coupling in order to minimize the curvature of the resulting sampling trajectories.
weaknesses: 1. The paper lacks discussion to some prior work that also attempted to improve minibatch-based OT coupling in flow-based generative models (see [1, 2, 4]).
2. The paper would benefit from additional analysis of how scalable the Algorithm 2 is with respect of the data size. Besides this, some additional details on how efficient the OT plan calculations are would be appreciated (Equation 5 and steps 6 and 8 in Algorithm 2).
3. The method assumes the dataset is available as whole prior to the training. However, modern generative models trained at scale sometimes deal with streaming data that is not available before the training. Could the authors provide some discussion on this?
4. While the paper suggests a way to deal with class-conditional training, it is not clear how it would extend to more complex conditioning signals, such as text. [3] propose a solution toward this. Could the authors provide some discussion whether this or similar techniques could be applied to AlignFlow?

[1] Davtyan, Aram, et al. "Faster inference of flow-based generative models via improved data-noise coupling." The Thirteenth International Conference on Learning Representations. 2025.

[2] Zhang, Stephen, et al. "On fitting flow models with large sinkhorn couplings." arXiv preprint arXiv:2506.05526 (2025).

[3] Cheng, Ho Kei, and Alexander Schwing. "The curse of conditions: Analyzing and improving optimal transport for conditional flow-based generation." arXiv preprint arXiv:2503.10636 (2025).

[4] Calvo-Ordonez, Sergio, et al. "Weighted conditional flow matching." arXiv preprint arXiv:2507.22270 (2025).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: (1) Does the method resample noise every epoch?
soundness: 3
strengths: (1) Simple, general mechanism: drop-in coupling that works with multiple FGM targets (FM, Shortcut, MeanFlow, etc.). Algorithm 3 is clear.
(2) Practical engineering: seed-based noise storage; class-wise SDOT; lightweight Stage-1 overhead; handy MRE/L1 diagnostics and tuning guidance.
(3) Empirical gains: better FID at fixed NFE across multiple backbones; improvements persist from U-Net/CIFAR10 to DiT/SiT on ImageNet256. Tables 2–4 & Fig. 2.
summary: The paper proposes AlignFlow, a two-stage recipe for FGM training: (1) precompute a semi-discrete OT (SDOT) map from continuous noise to the empirical data distribution by partitioning the noise space into Laguerre cells; (2) train any flow-matching-style model while pairing each sampled noise with its SDOT-assigned datum. The method is “plug-and-play,” deterministic (fixed coupling), and reported to (a) speed convergence and (b) improve FID across several FGM families (Flow Matching, Consistency/Shortcut, Live Reflow, MeanFlow) on CIFAR-10 and ImageNet256 (latent). The paper also argues that SDOT avoids the OT sample-complexity “curse of dimensionality” by optimizing against p_1 rather than the unknown population \tilde p_1.
	​
weaknesses: (1) The paper acknowledges no assumption about SDOT ≈ population OT, but several claims and intuitions (e.g., straighter, “more optimal” paths) implicitly rely on population properties.
(2) Some parts of the presentation are unclear, e.g. the algorithm description.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The paper states that "AlignFlow particularly suitable for large models that require small-batch training". This seems to weaken the contribution of this paper, as methods such as [1] works on small-batch training and [2] works on larger models already.

2. Could authors explain the reasons behind excluding discussions and experiments of AlignFlow on continuous normalizing flows (CNFs)? Demonstrating benefits there would strengthen the claim of generality.

3. Could the SDOT map be updated online to accommodate new data without recomputing from scratch?

##### Reference

[1] Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid RectorBrooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport.

[2] Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling.
soundness: 3
strengths: The paper is written clearly. It explains the problem settings, methodology, and contributions quite well.

Uses SDOT to build a deterministic coupling is a novel yet simple idea that avoids high-dimensional OT pitfalls. Similar ideas are rarely explored and has only one contemporary work published in 2025.

Demonstrated improvement across model sizes and datasets to imply the generality of the algorithm. Wall-clock times for experiment runs on CIFAR and ImageNet are reported to support claims for SDOT cost.
summary: AlignFlow introduces semi-discrete optimal transport (SDOT) as a deterministic way to align each noise-data pairs to ensure more straight flow paths in flow-based generative models (FGMs) training, reducing Number of Function Evaluations (NFE).

A dual-weight vector defines the transport plan between noise and data, corresponding to Laguerre cells that partition the noise space. Computes dual-weight vector by using Adam optimizer in maximizing an objective from the dual problem.

Training proceeds in two stages: (i) compute the SDOT map once (extra cost < 1 %), and (ii) train any FGM while re-using these fixed (noise, data) pairs .

Deterministic alignment yields straighter generative paths, fewer NFE, and bypasses the sample-complexity curse.

Plug-and-play integration with Flow-Matching, Shortcut Models, MeanFlow, etc. Improves FID across CIFAR-10 and ImageNet-256 at low NFE.
weaknesses: 1. Although recommendations for setting SDOT hyper-parameters ($\epsilon$, lr, EMA β) is discussed qualitatively to help practitioners use the algorithm, sensitivity to hyperparameters is not discussed in the experiments. If the algorithm is to be further scaled, existing practices for choosing hyper-parameters could potentially fail.
2. Experiments are done on a rather limited data set size; scalability to larger datasets or infinitely sampled settings (e.g., synthetic data augmentation) is unclear. Paper mentioned that processing time is 8min 30s for CIFAR10, and that the algorithm scales quadratically w.r.t. number of targets. Extrapolating to industrial-sized datasets with typically millions of samples, the pre-processing could be astoundingly time-consuming.
3. The difference of FID-50k with and without AlignFlow for larger models and better FID is quite small according to Table 3. Even though AlignFlow improves FID across sizes, the improvement is not very significant. Perhaps the authors should consider using a new evaluation task or metric to highlight improvement.


Decision: Accept (Poster)