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.
Figures & tables
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)
✓
✓
✓
Table 1: A comparison between coupling methods and AlignFlow, a Noise-Data Alignment method.
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 Li equals the discrete probability mass of the associated data point.
Figure 2: Training curves for AlignFlow on different tasks. Each figure illustrates the FID-50k score against the number of training steps. The results demonstrate that AlignFlow provides a consistent and simultaneous improvement over all baseline algorithms shown, enhancing both final performance and training convergence speed.
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 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.
Algorithm
AlignFlow?
NFE=4
Difference
NFE=1
Difference
Flow Matching
✓
93.16
↓ 32.46
276.18
↓ 28.86
✗
125.62
305.04
Consistency Training
✓
103.14
↓ 8.70
64.33
↓ 12.04
✗
111.84
76.37
Live Reflow ( Frans et al., 2025 )
✓
60.23
↓ 34.52
47.06
↓ 12.81
✗
94.75
59.87
Table 3: Evaluation of AlignFlow on DiT-B/2 for ImageNet256 using FID-50k demonstrates consistent performance improvements across all tested NFE configurations.
Backbone
# params
w/ AlignFlow
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.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
# step
learning rate
batch size
EMA parameter β
entropic reg ϵ
MRE
L1
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
Appendix
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 β
entropic reg ϵ
MRE
L1
3000
10
4096
0.99
0.01
∼ 0.08
∼ 0.016
Appendix
Table 6: SDOT hyperparameters for ImageNet256 (class-conditioned, latent space) generation. Computation takes less than 10 seconds per class on an NVIDIA L40S GPU.
Figure 3: Comparison of generative trajectories in FGMs among various methods. AlignFlow has a straighter trajectory compared to vanilla Flow Matching and has a clearer boundary compared to Minibatch OT (e.g., at t=0.22 ).
Figure 4: Images generated by MeanFlow+AlignFlow trained on ImageNet256 (FID=3.23, NFE=1).
Diffusion and flow-matching models are typically trained by corrupting data through independently sampled Gaussian noise. While simple and scalable, this forward process induces arbitrary data-noise couplings, forcing the network to learn high-curvature transports between unrelated endpoints. Existing optimal-transport methods reduce this burden by reassigning fixed noise samples to data, but the source noise distribution itself remains passive. To address this, we introduce Contrastive Noise Alignment (CNA), a training-time method that creates dynamic, contrastive couplings by optimizing the noise representations directly. By modeling the noise batch as an interacting particle system, CNA employs a cross-modal InfoNCE objective to align noise particles with their paired data targets. To prevent spatial collapse, this alignment is regularized using an angular entropy term and a radial norm penalty. We show theoretically that this equilibrium asymptotically preserves Gaussian structures, maintaining tractability during inference. Empirically, CNA improves the alignment between noise and data, reduces flow curvature, and provides better generation quality with fewer required sampling steps. For few-step, pixel-space generation (2-4 NFEs), CNA reduces FID by over 50% compared to standard rectified flow, and by at least 24% against Optimal Transport baselines.
Diffusion and flow matching models generate samples by learning time-dependent vector fields whose integration transports noise to data, requiring tens to hundreds of network evaluations at inference. We instead learn the transport map directly. We propose Optimal Transport Neural Flow Matching (OT-NFM), an ODE-free generative framework that parameterizes the flow map with neural flows, enabling true one-step generation with a single forward pass. We show that naive flow-map training suffers from mean collapse, where inconsistent noise-data pairings drive all outputs toward the data mean. We prove that consistent coupling is necessary for non-degenerate learning and address this using optimal transport pairings with scalable minibatch and online coupling strategies. Experiments on synthetic benchmarks and image generation tasks (MNIST and CIFAR-10) demonstrate competitive sample quality while reducing inference to a single network evaluation.
Flow matching models learn to transport samples from a simple prior distribution to a complex data distribution. When prior-data pairs are coupled via optimal transport (OT), the learned trajectories are straight and non-crossing, enabling fast, even single-step, generation. However, computing the OT coupling in high dimensions is intractable, and existing methods attempt to solve the OT problem, at the cost of persistent bias or significant overhead. Rather than solving for the OT coupling, we reformulate the problem. Once the prior is treated as a design choice rather than a fixed input, the OT coupling between prior and data is no longer unique. Many priors admit an OT-optimal identity coupling to the data, leaving us free to choose one that is also tractable to sample. We identify low-frequency projection of natural images as such a choice. The identity coupling between data and its low-frequency representation is empirically OT-optimal, the prior is structured enough to be sampled by a lightweight model at inference, and the remaining flow-matching task reduces to synthesizing high-frequency detail. Interpolating the prior with Gaussian noise further improves generation quality while preserving the OT coupling. The approach requires no modifications to the flow model itself, and integrates naturally with latent-space models, classifier-free guidance, and one-step generation frameworks. Across all benchmarks, our method reduces trajectory curvature by more than 2× compared to existing flow matching methods, yielding better generation quality in the few-step regime.