Drifting models enable one-step generation by amortizing distribution transport into training, but this efficiency places greater demands on the transport field estimated at each update. In large-scale training, the field is computed from finite mini-batches of generated and real samples, which provide only imperfect approximations to the underlying distributions. Balanced optimal transport enforces exact mass matching within every mini-batch, making the estimated field sensitive to the particular composition of the real-data batch. We find that generated and real samples should be treated asymmetrically: letting the mass assigned to real samples adapt while keeping every generated sample fully transported improves generation across six feature-space metrics in controlled ablations, and is more robust to the relaxation strength than relaxing both marginals simultaneously, which falls below balanced transport under stronger relaxation. Motivated by this observation, we propose Unbalanced Optimal Transport Gradient Flow (UOT-GF), which keeps the generated-sample marginal fixed and relaxes only the real-data marginal. Under identical settings at DiT-B/2 on ImageNet-256, UOT-GF improves Fréchet Inception Distance (FID) from 1.53 to 1.46 over the balanced W-Flow baseline; scaling the same recipe yields 1.34 and 1.22 FID at L/2 and XL/2, the best FID among the one-step models we compare. We further derive the induced UOT transport force, establish a kinetic Vlasov--Fokker--Planck formulation whose overdamped zero-temperature limit recovers the drifting dynamics, and characterize non-target stationary states together with sufficient conditions for convergence.
Figures & tables
Figure 1: Selected 1-NFE samples from UOT-GF-XL/2 trained from scratch on ImageNet-256 (guidance scale 1.14 ).
Figure 2: Fréchet Inception Distance (FID) versus effective sampling compute (parameters × NFE ( Song et al., 2023 ) ) on ImageNet-256. 1-NFE samples are shown in Appendix J .
Figure 3: Oversampled mode in a minibatch. One mode supplies 80% of the targets (red crosses). Balanced OT (a) must deliver full mass to them, pulling sources from distant modes; symmetrized source-fixed UOT (b) reweights them and keeps most couplings local. Gray: sources; blue: remaining targets. More examples in Appendix G.9 .
Figure 4: Mode imbalance. Equal-weight sources (gray dots) against an 80/10/10 target (teal crosses), left. Balanced OT (middle) reassigns source mass across clusters to match the proportions; symmetrized source-fixed UOT (right) keeps sources on nearby clusters.
Figure 5: Sampling efficiency on ImageNet-256, measured on a single Nvidia H20 under identical settings (log scale). (a) Throughput in images per second; (b) compute per image in TFLOPs. UOT-GF uses one network evaluation. Measurement details are in Appendix H .
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Numbers and Arrays
a
A scalar
x
A vector (bold lowercase)
P
A matrix (bold uppercase)
1N
All-ones vector of length N
x⊤,P⊤
Transpose
∥x∥
Euclidean norm
Appendix
Table 6
Ablation (Table 1 )
UOT-GF B/2
UOT-GF L/2
UOT-GF XL/2
Generator Architecture
arch
DiT-B/2
DiT-B/2
DiT-L/2
DiT-XL/2
input size
32×32×4
32×32×4
32×32×4
32×32×4
patch size
2×2
2×2
2×2
2×2
hidden dim
768
768
1024
1152
depth
12
12
24
28
Appendix
Table 5: Configurations and hyperparameters. All rows except the transport rows ( (φ1,φ2) , τ , L , ϵ , symmetrization, diagonal mask) follow Han et al. (2026, Table 5) . CFG sampling entries abbreviate p(w)∝(⋅) .
Figure 6: Forward source-fixed UOT reweights real targets. (a) Balanced OT fixes both marginals; source-fixed UOT fixes the generated marginal and KL-relaxes the real-data marginal. (b) Lowest-mass targets and class-matched controls, where rj is received mass relative to the uniform target marginal. (c) Rankings over 256 targets per class are stable across four contexts ( ρ=0.926 , Jaccard =0.651 ). Lower mass indicates reduced weight in the forward interaction, not lower image quality. Right: final-checkpoint FID and IS.
Figure 7: Wrong-manifold targets. Part of the target mass (red crosses) lies on a corrupted manifold. Balanced OT (left) must transport full mass to it; symmetrized source-fixed UOT (right) reduces the mass sent there. Gray: sources; blue: nominal targets.
Figure 8: Spurious-island targets. A cluster of spurious targets (red crosses) appears away from the data. Balanced OT (left) couples sources to it at full weight; symmetrized source-fixed UOT (right) reduces its influence.
Figure 9: Single-H20 inference efficiency at 256×256 resolution. (c) Mean batch-one generation-and-decoding latency over 200 timed requests. (d) Batch-125 throughput speedup relative to SiT-XL/2. Both panels use logarithmic vertical axes and model-specific high-fidelity sampling settings. The speedup in (d) is not computed from batch-one latency.
Figure 10: Controlled comparison of balanced OT and source-fixed UOT on identical frozen XL sources and real ImageNet targets. Top: received target mass relative to the uniform reference. Bottom: row-normalized positive-forward couplings; the colored strips distinguish 92 clean targets, 20 members of four duplicate families, and 16 foreign-class targets.
Figure 11: Source-fixed UOT update for ImageNet class 178 using the frozen XL checkpoint. (a) Shared PCA projection of the latent particles and one Euler update; arrows connect generated particles to their updated positions. (b–g) Coupling components used to form the update. Computation uses 4096-dimensional global latents, ϵeff=204.8 , τ=0.985 , and 10 Sinkhorn iterations.
Figure 12: The same frozen-checkpoint source-fixed UOT probe for ImageNet class 208. (a) Shared PCA projection and one Euler update. (b–g) Corresponding coupling components. The PCA coordinates are for display only; all transport computations use the full latent representation.
Figure 13: CFG-scale analysis for our B/2 and L/2 models. The left panel shows the FID–IS trade-off, while the right panel reports the corresponding numerical results.
Figure 15: Random class-conditional ImageNet samples generated by our UOT-XL model at 200K training steps with CFG =1.14 . Class indices follow the zero-based ImageNet-1K label mapping.
Figure 16: Random class-conditional ImageNet samples generated by our UOT-XL model at 200K training steps with CFG =1.14 . Class indices follow the zero-based ImageNet-1K label mapping.
Figure 17: Qualitative comparison for lion (ImageNet Class 291): official W-Flow-B ( Han et al., 2026 ) (top) and our UOT-B (bottom). Both models use 200K-step EMA checkpoints, CFG =1.19 , and sampling seed =42 . Corresponding positions use the same initial noise and class label. All 21 samples per model from this seed are shown without filtering.
Figure 18: Qualitative comparison for Pembroke Welsh corgi (ImageNet Class 263): official W-Flow-B ( Han et al., 2026 ) (top) and our UOT-B (bottom). Both models use 200K-step EMA checkpoints, CFG =1.19 , and sampling seed =42 . Corresponding positions use the same initial noise and class label. All 21 samples per model from this seed are shown without filtering.
Figure 19: Qualitative comparison for king penguin (ImageNet Class 145): official W-Flow-B ( Han et al., 2026 ) (top) and our UOT-B (bottom). Both models use 200K-step EMA checkpoints, CFG =1.19 , and sampling seed =42 . Corresponding positions use the same initial noise and class label. All 21 samples per model from this seed are shown without filtering.
Figure 20: Effect of classifier-free guidance scale on samples from our XL model. Both panels show sports cars (ImageNet Class 817), generated using the same 200K-step EMA checkpoint and sampling seed =0 . Each panel contains 21 samples arranged in a 3×7 grid; corresponding positions use the same initial noise.
Teacher-free one-step generative models such as W-Flow learn directly from distributional objectives, with the Sinkhorn divergence providing the strongest empirical performance among the energies studied in W-Flow. However, Sinkhorn divergence is based on balanced optimal transport, whose exact marginal constraints can induce undesirable transport under finite minibatch fluctuations. We introduce UW-Flow, which replaces the balanced Sinkhorn energy with a debiased unbalanced Sinkhorn divergence and derives the corresponding Wasserstein velocity. The resulting dynamics retain W-Flow's cross-minus-self structure, recover the balanced limit, and preserve target identifiability under suitable assumptions. We further characterize the finite-batch velocity error of balanced and unbalanced transport by comparing each minibatch field with its population counterpart. Our analysis identifies conditions under which marginal relaxation yields smaller minibatch-to-population velocity error, and shows that for separated matched modes, balanced sensitivity grows with inter-mode separation while the corresponding Unbalanced Optimal Transport (UOT) sensitivity remains controlled by local transport geometry. Controlled experiments on fixed-support and continuous two-mode distributions support these predictions. On MNIST, CIFAR-10, and ImageNet-100, UW-Flow improves generation under matched training budgets, with the largest gains in unconditional generation and consistent improvements in the class-conditional setting.
Distributional training provides collective supervision for one-step visual generation by matching real and generated features in frozen representation spaces. We introduce a unified theoretical framework that separates distribution modeling from matching discrepancy and connects global objectives to pointwise feature updates through Wasserstein gradient flow. Under this framework, FD-Loss and Gaussian-kernel Drifting are recovered through Gaussian optimal transport and kernel-density-based KL matching, respectively. The framework motivates MGFlow, which models feature distributions with Gaussian mixtures at an adjustable granularity between global moments and sample-based representations. MGFlow supports both optimal transport and score-based matching, and couples mass-constrained sample assignment with paired component updates to address mode collapse that mixture expressivity alone does not resolve. On ImageNet 256×256, MGFlow substantially surpasses the FD-Loss baseline, achieving state-of-the-art results with 1.45 FDr6 on pMF-H and 1.64 on JiT-H. For text-to-image generation, MGFlow post-trains FLUX.2 [klein] 4B into a one-step generator that outperforms the original four-step model on both GenEval and PickScore.
Chi Zhang, Shi Haoyang, Yueyi Liu +9
Tsinghua University · Fudan University · Xi’an Jiaotong University +4
Diffusion models and flow-based methods have shown impressive generative capability, especially for images, but their sampling is expensive because it requires many iterative updates. We introduce W-Flow, a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step. This is achieved in two steps: we first define an evolution from the reference distribution to the target distribution through a Wasserstein gradient flow that minimizes an energy functional; second, we train a static neural generator to compress this evolution into one-step generation. We instantiate the energy functional with the Sinkhorn divergence, which yields an efficient optimal-transport-based update rule that captures global distributional discrepancy and improves coverage of the target distribution. We further prove that the finite-sample training dynamics converge to the continuous-time distributional dynamics under suitable assumptions. Empirically, W-Flow sets a new state of the art for one-step ImageNet 256×256 generation, achieving 1.29 FID, with improved mode coverage and domain transfer. Compared to multi-step diffusion models with similar FID scores, our method yields approximately 100× faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.