AlignFlow: Improving Flow-based Generative Models with Semi-Discrete Optimal Transport
Authors: Lingkai Kong
Molei Tao
Yang Liu
Bryan Wang
Jinmiao Fu
Chien-Chih Wang
Huidong Liu
Venue: ICLR 2026 Poster (2026)
Keywords: Flow-based generative model; flow matching; Semi-discrete optimal transport
Primary area: generative models

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.

Published as a conference paper at ICLR 2026
ALIGNFLOW: IMPROVING FLOW-BASED GENERATIVE MODELS WITH SEMI-DISCRETE OPTIMAL TRANS- PORT
Lingkai Kong", Molei Tao',
Yang Liu², Bryan Wang". Jinmiao Fu²; Chien-Chih Wang?, Huidong Liu?
1 Georgia Institute of Technology, 2 Amazon.com, Inc.
{1kong75, m @gateech edu (yliuu, brywan, jinmiaof, cewang, 1 1iwhuido)@amazzooooooo
ABSTRACT
Flow-based Generative Models (FGMs) effectively transform noise into com- plex data distributions. Incorporating Optimal Transport (OT) 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 FGMs. This paper introduces AlignFlow, novel approach that leverages Semi- Discrete Optimal Transport (SDOT) to enhance the training of FGMs by estab- lishing 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 corresponding data point. During FGM training, i.idd 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 Align- Flow improves the performance of a wide range of state-of-the-art FGM algo- rithms and can be integrated as plug-and-play component. Code is available at: wttp:///..oooooooo//l/o/c/ /i low
INTRODUCTION
A generative model in machine learning is designed to produce new data samples that closely re- semble 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. Ho et al.) 2020 Song et al. 202 for image synthesis (e.g.. Rombach 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.. (022). Flow Matching (Lipman et al. 2022). and
Stochastic Interpolant (Albergo et al. 2023).
This work focuses on improving wide range of Flow-based Generative Models (abbreviated as FGMs hereafter), including Flow Matching (Lipman et al., 2022), shortcut model (Frans et al.- 2023}. MeanFlow (Geng et al., 2025), Live Reflow (Frans et al.] (025), but excluding continuous normalizing flows (CNF) (Chen et al.) 2018): see Sec. [.] for detailed specification of scope. FGMs focus on learning time-dependent vector field, approximated by a Neural Network (NN) whose integration gives trajectory that transforms randomly sampled noise to 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 forward pass of large NN. Consequently, generating single sample requires multiple NN evaluations. This cost is typically measured by the Number of Function Evaluations
"This work primarily conducted during Lingkai Kong's internship at Amazon.

---

Published as a conference paper at ICLR 2026
(NFEs), which denotes the total number of forward passes during ODE integration. For example, in vanilla Flow Matching (Lipman et al. 2022), the NFEs are often greater than 100. The straightness of the ODE trajectory is closely related to NFEs, as straighter trajectory is typically casier to integrate, resulting in fewer NFEs required. See Apdx. D.2 for the inference procedure of FGMs.
Before diving into how 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 [ and I In this work, we focus on improving the first step. While many Sfate-of-the-Aa 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).
(SOTA) methods introduce
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 ct al. 2024 Cheng & Schwing, 2025). From a theoretical perspective, OT gives a shortest path between the noise and data distributions (Eq. (4), 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. [ Utilizing discrete OT, Tong al. (2023) and |Pooladian
ct 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 (Korillov et al. 2024] or misspecified (Nguyen et al. (022) 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 dimensionality (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 distribu- tion, but we can (and typically do) use noise distribution that is continuous. Therefore, we seek to optimally couple discrete distribution with continuous distribution, and the OT problem corre- sponding this case is known 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: defines the most direct and shortest connection between noise samples and data points, thereby guiding the target vector field learned by NN. Moreover, the SDOT map can be solved efficiently with guaranteed convergence, and the quality of the resulting map can be evaluated with low com- putational cost (Liu et al.] 202T). In addition, SDOT also has a clear geometric interpretation: the SDOT map partitions the space into Laguerre cells (Fig. D. 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 necessary condition for achieving optimal coupling and ensuring consistent noise-data matching regardless of batch size (Sec. 4 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. DT. 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 AlignFloww follows two-stage approach: in the first stage, the SDOT map is computed. In the second stage, general FGM is trained using this precomputed SDOT map. The key advantages for AlignFlow are summarized below:
* AlignFloow is a pplug-andplaay method that can be seamlessly integrated into existing FGMs. It readily combines with SOTA FGMs to further enhance performance.
2

---

Published as a conference paper at ICLR 2026
AlignFlow has its NDA component separated from FGM training, and therefore enjoys the provable convergence of SDOT in (Liu et al. (022). In contrast, previous OT-based FGMs lack convergence guarantees for the computation of coupling between noise and data.
AlignFlow bypasses the curse of dimeesiionality (Sec. 4.1) that arises when using discrete OT for continuous distributions (Genevay et al., 2019). enabling it to scale effectively to
large models and high-dimensione datasets.
The SDOT map defines deterministic and optimal transport path without randomness; that is, each noise sample is consistently matched to a fixed data point, independent of batch size. This batch-invariance ensures stable convergence (see Sec. 4.2 even when batch sizes are severely constrained, addressing a critical bottleneck for training large-scale models where memory limitations enforce small-batch regimes.
The extra cost for computing the SDOT map is low (less than 1 extra cost) (Sec. 4.3
2 RELATED WORK
A central factor in reducing the NFEs during sampling in FGMs is the straightness of the gener- ative trajectory. Since the inference process involves integrating a learned vector field, straighter trajectories easier to integrate accurately, thereby reducing NFEs. Numerous methods have been proposed to encourage straighter trajectories, which can be broadly categorized into three main ap- proaches:
Target vector field: This class of methods aims to straighten generative trajectories by guiding the NN to learn smoother more linear target vector field. Approaches such as consistency models (Sog et al. 2023) and shortcut models (Frans et al., 2025), inspired by techniques from diffusion models YYang
2024), enforce penalties 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 in class- conditioned generation tasks on the ImageNet dataset.
Distillation: Recent works aim to reduce NFEs by distilling a trained FGM into more efficient model (Boff et al., 22055: Dao et al., 2025). Distillation techniques have also shown success in diffusion models (Salimans & Hoj 2022; Song al.] 2023 Kim et al., (023). However, distillation
involves an additional training stage built upon pre-trained model, and thus can be viewed as complementary to our approach.
Coupling: Coupling is 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 carefully designed joint distribution:
Tong et al. (2023) and Pooladian et al. (203)) 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 each FGM training iteration. However, this approach is sensitive to batch size: large batches in- cur high computational cost, as Sinkhorn an iterative algorithm with per-batch complex- ity of O(B² log B) where B is the batch size, while small batch size leads to poor estimates of the optimal coupling. Calvo-Ordonez et al. (2025) introduce a weight from Gibbs kernel in the FGM training objective, which generalizes the idea of Tong al. (2023). However, its dependence on minibatch makes have similar weaknesses. Additionally, computing discrete OT in batches is known to produce biased transport plans (Komilov 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 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 unestab- lished.
Zhang al. 2025 provides an approach that scales the Sinkhomn algorithm to large-scale datasets, computing OT couplings between sampled noise and data points in large batches. Moreover, PCA can be used further speed up Sinkhorn computation, which which[Zhang al. (2025) reports does not sacrifice generation quality. [Datyaan et al. (2025) improves upon
3

---

Published as a conference paper at ICLR 2026
Scale to large models
Entire noise distribution
Explicit OT optimization
Algorithm
Tong et al. (2023)
X
X
Coupling
Pooladian et (2023
X
X
Liu et al. 2022
X
Kormlov et (2024)
X
X
Noise-Data
AlignFlow (ours)
Alignment
Table A comparison between coupling methods and AlignFlow, Noise-Data Alignment method.
mini-batch OT by computing the map between a pre-sampled noise set and the whole dataset. Nevertheless, it is limited by its ability handle only finite number of noise samples.
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, discussed in Theorem 3.5 of their paper.
A comparison is provided in Table While conceptually related, AlignFlow differs fundamentally from existing coupling-bart -baed methods: computes an SDOT map by considering the entire noise distribution and all empirical data points, resulting in optimal NDA for training FGMs.
After ICLR decision, we saw the announcement of a nice concurrent work Mousavi Mousavv-Hosseini et al. 2025.. Both their work and this demonstrate the effectiveness and robustness of the idea.
METHODOLOGY
Mathematically, the generative modeling task can be formulated as follows: given a dataset {ai his assumed to consist of i.i.d. samples from unknown probability distribution Ω1 space X where I is the index set of the dataset + The objective to generate new samples that follow the same distribution p).
Notation In this paper, establish the following notational conventions. We denote by Po the noise prior distribution (e.g., standard Gaussian distribution) from which sampling is straight- forward. Conversely, Pi represents the empirical data distribution, defined Dirac mixture P1 = ElEr be0zg- In our FGM training setup, the weights are uniform, i.e., by = 1/
3.1 FLOW-BASED GENERATIVE MODELS (FGMs)
We begin by summarizing a general FGM framework in Algo. In each training iteration, set of noise To, data Dy and time t are drawn. are computed as the interpolation between æo and æ1. Finally, NN u(x,, t; f)) is trained to approximate the target vector field Target (icctrFFccd and the NN parameters are updated accordingly.
(Vanilla) Flow Matching (Lipman et al. 2022) defines the target vector field as Target VectorField((zo. = I1 - To and uses the loss function Loss ((6), vi B
Examples:
B Σ
-
b ||:- Several existing methods, including the shortcut model (Frans al. 2025
Mean Flow (Geng al. 2025). Consistency Training (Frans al. 2025), Live Reflow Frans et conform to this general FGM framework. See Sec. D. for more details.
2025
In many settings, chosen as a latent space, with data obtained via encoding through a Variational Autoencoder (VAE).
Without loss of generality, assume -1,2,.., ..., N}. where N the dataset size.
4

---

Published as a conference paper at ICLR 2026
Algorithm 1: Flow-based Generative Models (Training)
Input: Source distrihution Po, dataset { her, neural network u(z, f; 0) Output: Learned velocity field :(),:;))
Hyperparameters:: batch size B, number of steps K
P Training loop for == 1... do
Sample i.i.d. {- pm
i.idd
sample noise
Sample minibatch (01 )= from dataset {aiher
sample data
8 ~ ( 1) for = I...,
b sample time
for 1... B do
o - (1 - 1 yay Px
10
Thrget Vectorr
u(r; t²;
11
12
L(Ø) Loss ((0), vz) } )=1 j=1)
objective function b optimization
13
Update using VoL
In lines and O of Algo. noise and data points are sampled independently, and thus paired ran- domly, leading to curved trajectories intrinsically (Liu et al. 2022 Hertrich al. 2025).
3.2 COUPLING BETWEEN NOISE AND DATA
In fact, the loss function in Algo. provides empirical estimate of the following expectation [|u(ze, t; 0) - VenteetVecttooieeld((0.... x =1]
L(0) x C()) =
This formulation implies that zo and 21 are sampled independently from Po and pt. respectively, as used in Algo. 1 that is (No, .1) ~ Po P1. Recent works, such (Pooladian et al. 2023 Liu et al. 2022
Tong et al. 2023), aim to construct more informative joint distributions by sampling (To, 21 from coupling y [po. PI):
(0) B..- (1,00000)))1l((1,, t; 0) Target PargetVecttrField(//a =1)
(2)
where [poo P1) denotes the set of all couplings (i.e., joint distributions) with marginals po and pr:
r (EP((Xx)) (200 dro (1(11) Vzl / 7(20 21) dz = po( Yaoo }
(3)
As evident from the definition, i' 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 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 compute the optimal coupling between two prob- ability distributions given a cost function e X' X' + IR, (see, e.g., Peyré et al. (2019) for comprehensive overview):
(/r c(V1, U2) dy))))
== arg min TET(m.
(4)
where we choose c(v1, V2) := ||l| 12 throughout the paper.
A discrete distribution contrasts with a continuous distribution (such as the normal distribution), in that the associated random variable assumes only finite (or countable) set of values. For example,
We here choose the loss function he Loss (10) (ft = E. || I| for simplicity.
*The minimum objective is W (q1, 92) = min min,er(e1.92) (Jxxx ll y2 2
5

---

Published as a conference paper at ICLR 2026
the empirical data distribution can be expressed as p1 1t Eier Suj. where قzg denotes the Dirac measure centered at data point xt. An OT problem between a continuous and a discrete distribution referred to Semi-Discres 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 SDOT problem can be com- puted using a liimmnnnsio vector g = [g.]kez. re- ferred to the dual weight, where [ is the number of points in the discrete distribution P1- Given the dual weight g, we compute p(::g)
I, which
maps noise to data index:
<f g) == arg min c(Fo, 9: ies
(5)
Given ((:g) : X + I, the SDOT map from the noise distribution to the discrete data points is di- rectly obtained.
(a) Dataset and noise (b) Laguerre cells
distribution
In fact, X' is partitioned into cells by the SDOT map in the sense that cell L; contains the points trans- ported to the i-th point in the dataset. Each cell referred to as the Laguerre cell, defined below:
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 L; equals the discrete probability mass of the associated data point.
(g) = {ar E X': c(x, M1 I c(x, 9G. Vj)
(6)
Fig. visualizes the partition in a 2-dim space.
We proceed with the computation of the dual weight SDOT, similar to general OT problems in Eq. (4). minimization problem. By analyzing its dual problem, the dual weight can be solved by maximiz- ing the following objective function utilizing the Laguerre cells (see e.g., Eq. 5.7 in Peyré et al. 2019
((g) Eg))= = (a)==[(
E/ (c(x, U.) - gi) dpo(x) + (g. b)
NEI L:(g)
whose gradient is given by Vg(g); -L dpo b. where b; is the probability for each data point. In our case, b; -
To solve this maximization problem, employ the Adam optimizer (Kingma & Ba 2014). While this approach to the SDOT problem not novel (e.g.. et al. (2019)), we introduce a new and efficient Exponential Moving Average (EMA)-based estimation method for computing the MRE and the L1 distance (further discussed in Sec. B to evaluate the quality of the resulting dual weights. This estimation facilitates hyperparameter tuning.
The computation of SDOT dual weight requires ((I7)) iter- ations to achieve a satisfactory MRE when = (Liu et al.) (021). and can be further reduced by setting the entropic regularization e as a positive improves the smoothness of the SDOT objective. Specifically, the outer loop (line 4h in Algo. 2 has a per iteration cost of O(Z]), and the SDOT objective error decreases in the order ((1/k) (Taşkesen al. 2023).
Computational cost of Algo. 2
3.4 THE PROPOSED ALIGNFLOW
OT optimization remains core challenge in previous OT-based FGMs. In this subsection, propose using SDOT for Noise-Data Alignment (NDA) in training FGMs.
Idea behind AlignFlow The critical insight behind AlignFlow to use the SDOT to compute a transport map between the noise distribution Po and the known empirical data distribution P1. This because 1) SDOT considers all empirical data points during optimization, rather than relying on samples drawn from both the source and target distributions required by discrete OT; 2) the SDOT

---

sample noise Published as a conference paper at ICLR 2026
Algorithm 2: SDOT dual weight optimization
Input: Source distrihution Po. dataset {x") and the corresponding probabilities b = [b'),,,, entropic
regularization strength e, EMA parameter B, batch size B, cost function
Output: Dual weight g [9:kkr
Initialization: V....
g 0. Rana
for 1, do
Sample i.i.d. {- po for j = I.... B
sample noiss
i * 0 then
h Maxier
b SDOT map with current
else
10
p g) ainien c( x[)
bj 1 ((=((x))
12
VE(g) = B Σj hj- -
13
Eemo = WV + (g))(g
b Smoothen by EMA b optimizat.....
14
Update g using (g)
15
Suna /g.. + (1 8)g
16
Return: dual weight B.na
objective is convex, and the number of iterations required optimize SDOT is ((12)) (Liu al. 2021]: and 3) the optimality of the resulting SDOT map can be estimated with sample complexity of O())) Importantly, the SDOT map between Po and P1 preserves the desirable properties of OT, such as inducing straight transport trajectories. Therefore, unlike previous OT 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 obtained, can be integrated into the general FGM framework. This leads to the development of AlignFlow (Algo.B) which leverages the SDOT map to align noise with data. Notably, Algo. incorporates all data points during optimization and inherently avoids the curse of dimensionality.
Algorithm 3: AlignFlow: Noise-Data Alignment using SDOT (Training)
Input: Source distribution po. dataset {ai Jai) neural network t; Output: Learned velocity field u(:):()) Hyperparameters: batch size B, number of steps K
Stage I: compute SDOT map
Run Algo. 2 to get dual weight g.
Stage 2: Train flow dowedased generative model
K- B
b. Total number samples needed
Sample i.i.d. (- ~ my ((x() for 1....,. M 10 ~ ((0, 1) for = 1...., M u (my),- Rebalance (4m
match noise to data b sample time > Only is needed. Sec.
Training lcop
12
for
13
for 1 = 1... ...... do
14
j == (k 16-11 B+
of + (1 (=- + 1n
15
Thrget Vectorffeldd ""))
16
17
18
L(0) Loss ((1), 1d A
(==(-11)--1)
objective funct ion b optimization
19
Update using VoL

---

Published as a conference paper at ICLR 2026
3.5 ADDITIONAL TECHNIQUES
Remark (Noise storage). In lines HandTl of Algo. 3 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. [ To address this, we only store the random seed used to generate each noise sample, so each noise- pair is represented as (seed, index) tuple.
Such an approach requires 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 random cropping rotation, directly into the SDOT map formulation can be challenging.
Fortunately, for most image generation tasks, effective augmentation is often limited to random hor- izontal 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 and 5..3 we denote the class-conditoneed data distribution by e for the e-th class. The procedure involves computing the SDOT map between the noise distribution po and each class- dass-specific distri- bution Pi.e. Once noise samples are drawn, AlignFlow generates corresponding (noise, data) pairs and performs the required Rebalance (Sec. C operation independently for each class.
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 align- ment 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: 1) A deterministic coupling ensures consistent noise-data matching regardless of batch size in training FGMs; 2) If q in Eq. is continuous, then the OT map from di to d must be deterministic (see e.g., Peyré al., Remark 2.24), i.e., determinism is a necessary condition for optimality.
2019
4.1
BYPASS THE CURSE OF DIMENSIONALITY
The sample complexity of OT measures the accuracy of an OT plan estimated by the discrete-O1 plan between samples. However, the following theorem shows the number of samples required is exponentially dependent on the dimension:
Theorem (Sample complexity of OT (Informal version for Thm. 1 in Fournier & Guillin (015
In a d-dimensional space, the error in estimating the 2- Wasserstein distance W2(., qn) between a distribution q and the empirical distribution 9n with only access to n samples is of order ~ n 1/d
The squared Wasserstein distance W2 corresponds the minimum value in Eq. 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. in Kornilov et al. (2024).
Thm.
'Based on our experiments, storing 10 epochs worth of noise samples for ImageNet training in latent space (see Sec. 5 and 5.2) would require terabytes of disk space. In addition to storage demands, 1/O throughput would pose sigmfiicntt bottleneck.
8

---

Published as a conference paper at ICLR 2026
As NDA approach, AlignFlow circumvents the curse of dimensionality by computing SDOT map between Po and p1- Since p1 is fully specified by the dataset, the SDOT map can, in principle, be computed with zero estimation error, thereby avoiding the curse of dimensionality.
4.2 DETERMINISTIC ALIGNMENT
The SDOT map theoretically fully deterministic: a given sample from the noise distribution is consistently mapped to fixed data point.
Intuitively, this determinism provides a significant advantage in convergence speed: to determine the target vector field at some and 24, the standard approach using the random coupling (as in Algo. requires iterating over the entire dataset:
u( 0) Eny~p; Targel VectorFielle Field( æ1 (xo), so (2) - xx)/t
(8)
However, the fixed coupling in AlignFlow bypasses this estimation process, directly providing the larget vector field for the NN to learn as:
u(x,, t; 0) Target peetorField(zz)) 4i (zo)), To (z1 xt)/t, 21 p(x)) (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 esti- mate the OT plan indirectly through batch-based approaches (e.g., Reflow operation in [iu et al. (022). ICNN in Kornilov al. (2024), and Minibatch OT in Tong et al. (2023)). AlignFlow's s di- rect computation of SDOT is both more accurate and efficient, and Stage adds negligible overhead empirically, accounting for less than 1% additional training time. Further implementation details are provided in Sec. B
Upon completion of Stage 1, the SDOT map is fully computed. Consequently, the only additional overhead in Stage 2 arises from generating the training noise-data pairs (Lines and Algo. This process is highly efficient, runs quickly modern GPUs, and incurs an almost negligible cost, typically less than 0.1 % of the total training time.
EXPERIMENTS
AlignFlow
w/ Alignflow AlignFlow reported
Alignflow wio Allanflow
Minibotch OT
60
XL/2
F10-50k
FID 50kk
40
3.5
2 Steps
0.5
1.0
26.65
1.5
2.0
40
60
160
240
x10
Steps
x10*
Epoch
CIFAR-10 results using U-Net, (b) DiT-B/2 ImageNet256 with (c) SiT ImageNe1256 with generated by adaptive ODE in- the shorteut model, generated using MeanFlow, generated using 1-step tegrator DOPRI5 (see Sec. 5.I).
4-step forward Euler.
forward Euler.
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 consistent and simultancous improvement over all baseline algorithms shown, enhancing both final perfor- mance and training convergence speed.
5.1 UNCONDITIONAL IMAGE GENERATION ON CIFAR- 10 WITH U-NET
Following Tong et al. 2023 Sec. 5.3), train unconditional U-Net (Ronneberger et al.) 2015) on the CIFAR-10 dataset. In this setup, the FGMs are trained directly the pixel space. The com-

---

Published as a conference paper at ICLR 2026
parative training curves shown in Fig. 2(a) and the FID-50k scores for various ODE integrators are detailed in Table 2 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
al.
2023
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
Table 2: Comparison of FID-50k scores between Minibatch OT and AlignFllow for U-Net trained on CIFAR-10, 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.
5.2 IMAGENET256 GENERATION USING DIT WITH SHORTCUT MODEL
AlignFlow can be seamlessly integrated with modern SOTA network architectures and scales ef- fectively to large-scale datasets. We train AlignFlow conjunction with Diffusion Transformers (DiT, Peebless & Xie (2023)) the flass-conditione ImageNet dataset at 256 x 256 resolution (ImageNet256). The FGMs operate the latent space size x 28 x 4. obtained from pre- trained VAE. All model hyperparameters
adopted directly from (Frans 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) while improvements across additional models using AlignFlow are reported in Table
Algorithm
AlignFlow?
NFE=4
Difference
NFE=1
Difference
93.16
32.46
276.18
28.86
Flow Matching
X
125.62
305.04
103.14
J 8.70
64.33
12.04
Consistency Training
X
111.84
76.37
Live Reflow Frans
60.23
34.52
47.06
12.81
al. 2025
94.75
59.87
Shortcut Models
30.31
4 2.80
43.92
2.73
(Frans et al. 2025)
33.11
46.65
Table 3: Evaluation of AlignFlow on DiT-B/2 for ImageNe1256 using FID-50k demonstrates con- sistent performance improvements across all tested NFE configurations.
5.3 IMAGENET256 GENERATION USING SIT WITH MEANFLOW
AlignFlow further enhances the one-step generative model MeanFlow (Geng al. 2025). which is based on Scalable Interpolant Transformers (SiT, Ma et al. (024)) and trained on class-conditioned ImageNe1256. The FGMs operate in the latent space of size 28 28 obtained from a pretrained VAE. For our experiments, we use a non-official Py Torch implementation (Zhul 2025). which reli- ably 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. [[c] and the corresponding FID scores are reported Table AlignFlow consistently improves both performance and convergence speed across all cases, demonstrating its scalability to large models. Example generated images are provided Fig. in the appendix.
Backbone
# params
AlignFloww
w/o AlignFlow
Difference
SIT-B/4
131M
13.75
15.53
1.78
SiT-B/2
131M
5.60
6.17
0.57
SiT-L/2
459M
3.51
3,84
0.33
SiT-XL/2
676M
3.23
3,43
0.20
Table 4: We evaluate AlignFlow on ImageNet256 using MeanFlow (NFE=1), showing consistent performance improvements across all model sizes.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
LK thanks Amazon for hosting his internship. LK and MT are grateful for partial supports by NSF Grants DMS- 1847802, DMSS2513699, DOE Grants NA0004261, SC0026274, and Richard Duke Fellowship. We also thank Macro Cuturi for valuable comments during the review process, which significantly improved the quality of this work.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303. 087744 2023.
Michael Albergo, Nicholas Boffi, and Eric Vanden-Fijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint 2rXiv:2303.08897, 2023.
Jason M Altschuler, Jonathan Niles- Niles-Weed, and Austin Stromme. Asymptotics for semidiscrete entropic optimal transport. SIAM Journal on Mathematical Analysis, 54(2):111--1711, 2022.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International confer- machine learning, pp. 146-155. PMLR, 2017.
Nicholas Matthew Boffi, Michael Samuel Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: A mathematical framework for consistency models. Transactions Machine Learning Research, 2025.
Sergio Calvo-Ordonez, Matthieu Meunier, Alvaro Cartea, Christoph Reisinger, Yarin Gal, and Jose Miguel Hernandez-Lobato. Weighted conditional flow matching. arXiv preprint arXiv:2507.222700 2025.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
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.
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. Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 2654-2662, 2025.
Aram Davtyan, Leello Tadesse Dadi, Volkan Cevher, and Paolo Favaro. Faster inference of flow- based generative models via improved data-noise coupling. In The Thirteenth International Con- ference on Learning Representations, 2025.
Richard Mansfield Dudley. The speed of mean glivenko-cantelli convergence. The Annals of Math- ematical Statistics, 40(1):40-55 50, 1969.
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 muchine learning, 2024.
Nicolas Fournier and Arnaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability theory and related fields, 162(3)::07-738, 2015.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. The Thirteenth International Conference on Learning Representations, 2025.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complex- ity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pp. 1574-1583. PMLR, 2019.
11

---

Published as a conference paper at ICLR 2026
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13447, 2025.
lan 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): 63(11):139-144, 2020.
Johannes Hertrich, Antonin Chambolle, and Julie Delon. On the relation between rectified flows and optimal transport. arXiv preprint 2rXiv:2505.19712, 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Dongjun Kim, Chieh-Hsin Lai, Wei- Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning proba- bility flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279, 2023.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Aleksandr Korotin. Optimal flow match- ing: Learning straight trajectories in just one step. Advances in Neural Information Processing Systems, 37:104 1042044 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Huidong Liu, Ke Ma, Lei Zhou, and Dimitris Samaras. Efficient semi-discrete optimal transport using the maximum relative error between distributions. 2021.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022.
Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas M Boffi, Eric Vanden- Fijeen- and Sain- ing Xie. Sit: Exploring flow and diffusion- generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23-40. Springer, 2024.
Alireza Mousavi-Hosseini. Stephen Y Zhang, Michal Klein, and Marco Cuturi. Flow matching with semidiscrete couplings. arXiv preprint wwXiv:2509.255/9 2025.
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.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF international conference on computer vision, pp. 2op.419-4422 2023.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, 11(5-6):355-00,,
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lip- man, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch cou- plings. arXiv preprint 2xXi::2304.14772. 2023.
Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, al. Improving language under- standing by generative pre-training. 2018.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
12

---

Published as a conference paper at ICLR 2026
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234-24--0000 234-241. Springer. 2015.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005522 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265. PMLR, 2015.
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.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
Bahar Taskesen, Soroosh Suuffezzadeh-Abadeee and Daniel Kuhn. Semi-discrete optimal trans- port: Hardness, regularization and numerical solution. Mathematical Programming, 199(1): 1033-1106, 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 arXivv2232.0004 2023.
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.
Stephen Zhang, Alireza Mousavi-Hosseini, Michal Klein, and Marco Cuturi. On fitting flow models with large sinkhorn couplings. arXiv preprint arXiv:2506.0553 2025.
Yu Zhu. Meanflow: Pytorch implementation.
https: ttw..:/.comttooccooo.o.o zhwou//suuuus////// MeanFlow 2025. PyTorch implementation of Mean Flows for One-step Generative Modeling.
13

---

Published as a conference paper at ICLR 2026
A
FUTURE WORK
Many modern tasks involve datasets with complex label structures, such as text-to-image generation where each data point is tuple consisting of image and text. Since text prompts rarely identical across a dataset, standard label-based techniques are insufficient.
Here can provide idea of how AlignFlow may be used address this task: suppose the data is given by tuples (I; y), where x is the image (let's take the image modality as an example) and the text, and the task the text (e.g., via an LLM or extract the text embedding of each text and apply traditional clustering method) and assign a label clustering(y) to each y, making the input data (x,, y, z). For example, V1 ="a dog is swimming" and ="a dog is running" could be clustered into the same cluster. Then, for each cluster, we compute the SDOT map fl to map noise images æ. Training the flow-based generative model with (x]ly) can be guided by the corresponding SDOT map 5/2., where z the class label for y.
train AlignFlow that generates new ā given some new ÿ. We can cluster
B EFFICIENT SDOT ALGORITHM FOR LARGE-SCALE DATASETS
B. PERFORMANCE METRIC FOR ALGO.
We revisit the definitions of Maximum Relative Error (MRE) and Lı distance proposed in (Liu ct al. 202I), which serve as metrics for evaluating the quality of the computed dual weights:
|p:
bil [g = =- 1p: bal,
MRE
ps p=:= [L) dpo
(10)
iET
bi
In Algo. they be estimated efficiently by MRE V llae and L2 ||
Mathematically, MRE and L1 serve as estimators for the fo -nrm and fn-rrmm of the gradient of the objective, specifically ||VE'l and || respectively. The critical importance of MRE lies in its ability to quantify the imbalance of the transported probabilities. Importantly, an outcome of MRE 1 guarantees that each target data point receives non-zero probability transported from the source distribution. MRF = 0 signifies that the SDOT solution is perfectly optimized: the objective function Eq. 4 minimized, and, crucially, the measure preservation constraint satisfied uniformly, ensuring each target receives an equal probability mass. (See also Sec. C for details on rebalancing.)
B.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) SDOT with entropic regularization is solving (Altschuler et al. 2022
min rET(po-P1)
c( an (dy 21) + KT,(T)po ④ pi)
(11)
xxx
where e > 0 the regularization strength and KL denotes the Kullback-Leibler divergence. While the inclusion of this auxiliary term introduces controlled bias into the solution, significantly enhances the smoothness and tractability of the optimization landscape.
B.3 HYPERPARAMETERS TUNING FOR SDOT OPTIMIZATION
Algo. involves the tuning of three primary hyperparameters: the entropic regularization strength the EMA parameter 3, and the Adam optimizer's 's learning rate Ir. Throughout the maximization process in Algo. 2 the L, metric expected exhibit continuous decay when the hyperparameters are set appropriately
Entropic regularization strength e dictates the fundamental trade-off between bias and op- timization difficulty. A larger induces greater bias in the resulting transport plan, while
14

---

Published as a conference paper at ICLR 2026
step
learning rate
batch size
EMA parameter
entropic reg
MRE
L1
0-1000
10
1024
0.99
3.4
0.29
1000-6000
0.1
4096
0.999
0.27
0.045
6000-11000
0.1
16384
0.999
0.01
0.11
0.019
Table 5: SDOT hyperparameters for CIFAR-10 (unconditional) generation. Computation takes minutes and 30 seconds an NVIDIA L40S GPU.
# step
learning rate
batch size
EMA parameter
entropic reg
MRE
L1
3000
10
4096
0.99
0.01
0.08
0.016
Table 6: SDOT hyperparameters for ImageNer256 (elass-conditioned, latent space) generation. Computation takes less than 10 seconds per class on an NVIDIA L40S GPU,
setting to plexity of the optimization landscape.
very small value significantly increases the non-smoothness and com-
The parameter should remain fixed throughout the entire optimization procedure. When optimizing the problem, consider increasing the batch size and/or decreasing the learning rate when L1 plateaus.
The optimization stops when MRE satisfies the required performance threshold. For robust performance in downstream tasks, such FGMs, we recommend maintaining MRE below 0.2.
Contrary to common practices in standard deep learning optimization, where learning rates around 0.001 preferred, the computation of the SDOT map (Algo. often requires relatively high initial learning rate. A value of Ir = 10 serves as a recommended starting point for tuning.
For the management of large datasets, is advisable increase the batch size and/or adjust the EMA parameter 3 closer to (e.g., from 0.99 to 0.999) to ensure a smoother and more stable gradient update over larger number of iterations.
B.4 HYPERPARAMETERS AND COMPUTATIONAL COST
On CIFAR-10, use the training set for training (50000 images with shape 32 x 32 x 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 and hyperparameters in Table
On ImageNet, we use the full training set consisting of 1,281,167 images across 1,000 classes. In both the shorteut model w/ AlignFlow (Sec. 5.2 and MeanFlow w/ AlignFlow (Sec. 5.3). the SDOT map is computed in the latent space size 28 x 28 x 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. Evaluating the SDOT map ip in Eq. (5) (Line in Algo. 3 requires computing the minimum across the entire dataset, incurring a computational cost of OTZ4)) with a small constant factor. While this may initially appear computationally expensive, modern machine learning models are typically over-pane wver-paamettrized, meaning the number of parameters vastly exceeds the number of data points. Consequently, the computational cost is dominated by the forward pass and backward propagation rather than the SDOT map evaluation. This analysis explains why the additional overhead observed our experiments negligible, regardless of dataset size.
C REBALANCE: HANDLING NON-PERFECT SDOT MAP
Algo. 3 deviates from the standard practice of sampling i.idd data pairs. Instead, it only samples noise 2o Po and subsequently feeds the NN with data generated by the SDOT map (Line
15

---

Published as a conference paper at ICLR 2026
This specific sampling scheme introduces a potential issue: if the SDOT map is not computed per- fectly, the data (()) fed to the NN will be biased. Specifically, when the gradient of the energy Ve non-zero, the set of mapped samples, represented by the set of indices (4(xx) Zn iid. po) be-
comes biased from the empirical data distribution p1- (See also MRE defined in Eq. (TO) in Sec. B
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. to MRE + 0 is computationally prohibitive, we in- troduce the rebalance operation in Line defined as:
M
M
rebalance
1 my): = arg max {my
max SEZ
Σ (.))) SEZ
min Σ (m))
S
(12)
Intuitively, the rebalance operation determines the minimal perturbation to the set of target indices {my} such that the empirical frequency of each data point in the modified set {m,} is nearly uni- form. 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 intro- ducing a minimal amount of controlled randomness in the noise-data coupling.
In our experiments, particularly on the CIFAR-10 dataset (Sec. 5.1) J the rebalance operation did not yield noticeable difference in performance. This is likely because Algo. 2 achieved 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 consis- tently across all our experiments to ensure that the data fed into the FGMs is unskewed.
D MORE DETAILS FOR
GENERATIVE MODELS
D.1 MORE EXAMPLES FOR FGM FRAMEWORK IN ALGO.
Shorteut model In this model, an auxiliary input d is introduced to the NN, i.e., u(x, L, dd 0) and d' is i.dd sampled from D(-10)) Given hyperparameter k, choose Target VoctorFicld(ar, a1) = of - zy for j = 1,..,Kk, K. and Target VectorField( m1) = StopGrad(s} (s] + (44 for = + 1..... B, where sf == u( Iletd == of + s,di
444 u(>)) 1 + d, d²). Together with Loss (104 od ))m))) E, 1100 o 10 Algo. recovers the shortcut model in Frans et al. (005)..
MeanFlow In this model, an extra input r introduced to the NN, i..., u u(I, r; r). By choos- ing Target VectorField(zz. .1) StopGrad (t - r) u + 0,u) and Loss (10' v (=)))
Σ e II Algo. recovers the MeanFlow in Geng et al. 2025).
B
D.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 The inference process entails integrating the ODE 0, = u(I,,4 I; 0) over the time interval E [0, I]. where the initial condition 2i is sampled from the noise distribution Po. and the final state 21 constitutes 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, f). Intuitively, consider two hypo- thetical trajectories: an almost straight and highly complex, curvilinear path. The almost straight path permits straightforward computation, possibly achievable a single step via the forward Euler method, 2 = zo + u(zo, O; 0). 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 lower NFEs of the NN (Liu et al., 2022). Algorithms trained with a carefully-desigae target vector field benefit directly from this enhanced straightness. For instance, MeanFlow (Geng al. 2025) is capable
16

---

Published as a conference paper at ICLR 2026
1-1.
1-:11
CH-013
(a) AlignFlow (Alog. t
(b) Minibatch OT Tong et al 2024
(c) Vanilla Flow Matching
Lipman et al. 2022
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 = 0.22).
of generating high-qualing samples with an NFEs as low 1, whereas the vanilla Flow Matching (Lipman al. 2022) typically requires over 100 NFEs to achieve comparable results.
Algorithm 4: FGM/ AlignFlow (Sampling)
Input: Noise distribution Po, neural network a u(x, 0), ODEIntegrator
Output: A new sample from the data distribution
Sample {=(0)) ~
b. sample noise
v(t) = u[x(tt, t;
=(1)= ODEIntegrvaoor((())
Return: new sample x(1)
E SYNTHETIC EXPERIMENT: CHECKERBOARD
In this section, we analyze and visualize the characteristics of different learned trajectories by train- ing the FGMs algorithm on synthetic two-dimensional data distribution. Following the precedent set by (Lipman al. 2022 Fig. 4), the target data distribution is defined as the checkerboard distribution in 2, x 2]. The noise distribution is the standard normal distribution.
Crucially, our experimental setting differs from prior work (Lipman et al. 2022; Tong al. 2023 which assumed access to an infinite data stream during training 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 illustrates the density evolution time f progresses from 1, charting the transformation from the Gaussian distribution po to the checkerboard distribution. The visualizations indicate that AlignFlow generates significantly straighter transport trajectory when compared to trajectories learned via Minibatch OT and Vanilla Flow Matching, depicted in the corresponding panels of Fig. 3 All comparative experiments utilize identical hyperparameters to ensure a fair comparison of the learned vector fields.
'In the checkerboard experiments of L Lipman et al. for every minibatch.
2022 Tong et al. 2023) fresh training data is sampled
17

---

Published as a conference paper at ICLR 2026
F CAPABILITY OF GENERALIZATION
natural question arises from the observation that the SDOT map already constitutes a mapping from the source space the target space (X' + X). Since the integrated vector field of FGM also provides an X' + A' 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 constrained only map noise samples to the fixed points in the dataset. Consequently, it cannot generate new data instances. Therefore, a powerful NN needs be trained Stage 2 of Algo. Algo.3]atop the SDOT align- ment 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.
18

---

Published as a conference paper at ICLR 2026
G IMAGE SAMPLES GENERATED BY ALIGNFLOW
Figure 4: Images generated by McanFlow+AAignFllo trained on ImageNet256 (FID=3.23, NFE=1).
19


Reviews:
[{"review_id": "zi0lYP1UbM", "reviewer": "Reviewer_k6zR", "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.", "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$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n- The experimental results show an improvement over the training performance of mini batch OT, and the proposed method is faster as well.", "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "90qxvaVl1t", "reviewer": "Reviewer_urUF", "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.", "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?\n\nSome typos:\n1. Line 131. The sentence seems to miss a word.\n2. Line 208. smaple $\\rightarrow$ sample.\n3. Line 398. I believe Semi-Discrete Optimal Transport was meant instead of Stochastic Denoising Optimal Transport.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.", "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]).\n2. 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).\n3. 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?\n4. 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?\n\n[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.\n\n[2] Zhang, Stephen, et al. \"On fitting flow models with large sinkhorn couplings.\" arXiv preprint arXiv:2506.05526 (2025).\n\n[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).\n\n[4] Calvo-Ordonez, Sergio, et al. \"Weighted conditional flow matching.\" arXiv preprint arXiv:2507.22270 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TKd1tWyo7f", "reviewer": "Reviewer_GdTu", "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.\n\t​", "questions": "(1) Does the method resample noise every epoch?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) Simple, general mechanism: drop-in coupling that works with multiple FGM targets (FM, Shortcut, MeanFlow, etc.). Algorithm 3 is clear.\n(2) Practical engineering: seed-based noise storage; class-wise SDOT; lightweight Stage-1 overhead; handy MRE/L1 diagnostics and tuning guidance.\n(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.", "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.\n(2) Some parts of the presentation are unclear, e.g. the algorithm description.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7jgEdmExzg", "reviewer": "Reviewer_mCfS", "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).\n\nA 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.\n\nTraining 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 .\n\nDeterministic alignment yields straighter generative paths, fewer NFE, and bypasses the sample-complexity curse.\n\nPlug-and-play integration with Flow-Matching, Shortcut Models, MeanFlow, etc. Improves FID across CIFAR-10 and ImageNet-256 at low NFE.", "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.\n\n2. 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.\n\n3. Could the SDOT map be updated online to accommodate new data without recomputing from scratch?\n\n##### Reference\n\n[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.\n\n[2] Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is written clearly. It explains the problem settings, methodology, and contributions quite well.\n\nUses 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.\n\nDemonstrated 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.", "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)