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).