Unbalanced W-Flow: Unbalanced Wasserstein Drifting under Finite-Batch Marginal Fluctuations

Jongwook Jeon; Mingue Park; Jiwon Choi; Jaihoon Kim; Yunhong Min; Minhyuk Sung

cs.LGOct 13, 2026

Unbalanced W-Flow: Unbalanced Wasserstein Drifting under Finite-Batch Marginal Fluctuations

Authors: Jongwook Jeon, Mingue Park, Jiwon Choi, Jaihoon Kim, Yunhong Min, Minhyuk Sung

Organizations: Yonsei University · KAIST

Abstract

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.

Figures & tables

Panel of Figure 2 — Figure 2: Minibatch velocity error on fixed supports. Each method is evaluated against its own population velocity. Panel (a) sweeps $B$ at $d=4$ , while panel (b) sweeps $d$ at $B=16$ . Within each panel, the left plot shows $B\times\mathrm{MSE}$ and the right plot shows $B\times\mathrm{RelMSE}$ . RelMSE is MSE divided by the corresponding population-field energy.

Table 3

Figure 4: CIFAR-10 training dynamics across transport support sizes. FID versus training epoch for W-Flow and UW-Flow with

\tau\in\{0.1,1,10\}

. FID is computed using 10,000 generated samples at each evaluation. Panels show

B=4

16

, and

64

, with the sample budget per optimizer update fixed at

B_{\mathrm{opt}}=N_{c}B=128

Setting	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
Uncond	W-Flow	55.25	23.89	0.42	0.14
Uncond	UW-Flow	46.17	27.71	0.40	0.29
Cond	W-Flow	9.43	60.31	0.74	0.26
Cond	UW-Flow	9.11	61.34	0.75	0.26

Table 4: ImageNet-100 results at matched budgets.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Full fixed-support velocity-error results. Columns show $B\times\mathrm{MSE}$ , $B\times\mathrm{RelMSE}$ , and cosine similarity. The top and bottom rows vary $B$ and $d$ , respectively.

Panel of Figure 6 — Figure 6: Velocity errors under Gaussian-mixture sampling. Columns show $B\times\mathrm{MSE}$ , $B\times\mathrm{RelMSE}$ , and cosine similarity. The top and bottom rows vary $B$ and $d$ , respectively.

Algorithm 1 UW-Flow Training 1: Generator

f_{\theta}

; distributions

p_{\rm ref},p

; sample sizes

N,M

; step size

\eta

;

\varepsilon,{\color[rgb]{0,0.2539,0.7461}\tau},L

. 2: repeat 3: Sample

Z=\{z_{i}\}_{i=1}^{N}\sim p_{\rm ref}

4: Sample

Z^{\prime}=\{z^{\prime}_{i}\}_{i=1}^{N}\sim p_{\rm ref}

5: Sample

Y=\{y_{j}\}_{j=1}^{M}\sim p

X\leftarrow f_{\theta}(Z)

X^{\prime}\leftarrow\operatorname{sg}(f_{\theta}(Z^{\prime}))

(T^{+},s^{+})\leftarrow\textsc{USinkhorn}(X,Y;\varepsilon,\tau,L)

(T^{-},s^{-})\leftarrow\textsc{USinkhorn}(X,X^{\prime};\varepsilon,\tau,L)

10:

V_{i}\leftarrow s_{i}^{+}(T_{i}^{+}-x_{i})-s_{i}^{-}(T_{i}^{-}-x_{i})

11:

\widetilde{x}_{i}\leftarrow\operatorname{sg}(x_{i}+\eta V_{i})

12: Update

\theta

by minimizing 13:

\mathcal{L}(\theta)=\frac{1}{N}\sum_{i=1}^{N}\|f_{\theta}(z_{i})-\widetilde{x}_{i}\|_{2}^{2}

14: until converged

Algorithm 2 Unbalanced Sinkhorn 1: Particles

X=\{x_{i}\}_{i=1}^{N}

Y=\{y_{j}\}_{j=1}^{M}

;

\varepsilon,\tau,L

. 2:

a_{i}\leftarrow 1/N,\quad b_{j}\leftarrow 1/M

\triangleright

Weights 3:

C_{ij}\leftarrow\tfrac{1}{2}\|x_{i}-y_{j}\|_{2}^{2}

\triangleright

Cost 4:

\phi_{i}\leftarrow 0,\quad\psi_{j}\leftarrow 0

\triangleright

Initialize 5:

\alpha\leftarrow\dfrac{\tau}{\tau+\varepsilon}

\triangleright

Relaxation 6: for

\ell=1,\ldots,L

do 7:

\phi_{i}\leftarrow-\alpha\varepsilon\log\sum_{j}b_{j}e^{(\psi_{j}-C_{ij})/\varepsilon}

\triangleright

Source 8:

\psi_{j}\leftarrow-\alpha\varepsilon\log\sum_{i}a_{i}e^{(\phi_{i}-C_{ij})/\varepsilon}

\triangleright

Target 9: end for 10:

\Pi_{ij}\leftarrow a_{i}b_{j}e^{(\phi_{i}+\psi_{j}-C_{ij})/\varepsilon}

\triangleright

Coupling 11:

r_{i}\leftarrow\sum_{j}\Pi_{ij}

\triangleright

Row mass 12:

s_{i}\leftarrow\dfrac{r_{i}}{a_{i}}

\triangleright

Mass ratio 13:

T_{i}\leftarrow\frac{1}{r_{i}}\sum_{j}\Pi_{ij}y_{j}

\triangleright

Barycenter 14: return

(T,s)

Appendix

Table 8

$N_{c}$	$N_{\mathrm{gen}}$	$N_{\mathrm{pos}}$	Method	FID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
16	4	4	W-Flow	19.70	0.536	0.059
			UW-Flow ( $\tau=10$ )	22.54	0.562	0.081
			UW-Flow ( $\tau=1$ )	17.98	0.464	0.116
			UW-Flow ( $\tau=0.1$ )	18.66	0.419	0.253
8	8	8	W-Flow	13.69	0.461	0.239
			UW-Flow ( $\tau=10$ )	13.50	0.545	0.110

Appendix

Table 5: Full results of the fixed-budget minibatch discrepancy ablation on unconditional MNIST, with fixed

B_{\mathrm{opt}}=N_{c}N_{\mathrm{gen}}=64

$N_{c}$	$N_{\mathrm{gen}}$	$N_{\mathrm{pos}}$	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
32	4	4	W-Flow	30.58	9.75	0.55	0.01
			UW-Flow ( $\tau=10$ )	22.03	9.58	0.57	0.02
			UW-Flow ( $\tau=1$ )	19.41	10.10	0.56	0.02
			UW-Flow ( $\tau=0.1$ )	15.56	10.60	0.52	0.07
16	8	8	W-Flow	24.39	8.99	0.58	0.01
			UW-Flow ( $\tau=10$ )	14.64	10.39	0.59	0.04

Appendix

Table 6: Full results of the fixed-budget minibatch discrepancy ablation on unconditional CIFAR-10, with fixed

B_{\mathrm{opt}}=N_{c}N_{\mathrm{gen}}=128

Hyperparameter	MNIST	CIFAR-10	ImageNet-100
Model and representation
Generator	DiT-based, patch 2	DiT-based, patch 2	DiT-B/2
Training shape / space	$28\!\times\!28\!\times\!1$	$32\!\times\!32\!\times\!3$	$32\!\times\!32\!\times\!4$
Feature encoder (frozen)	ConvNeXt V2–Tiny	ConvNeXt V2–Tiny	Latent-MAE ResNet (width 256)
Feature normalization	Distance-based scaling	Distance-based scaling	Distance-based scaling
Velocity normalization	RMS normalization	RMS normalization	RMS normalization

Appendix

Table 7: Hyperparameters for image-generation experiments.

N_{\mathrm{gen}}

and

N_{\mathrm{pos}}

are the numbers of generated and target samples per transport group; the total generated batch size is

N_{c}N_{\mathrm{gen}}

. The listed

\varepsilon

and

\tau

values are before scaling by the feature dimension

D

Setting	$N_{c}$	$N_{\mathrm{gen}}$	$N_{\mathrm{pos}}$	$B_{\mathrm{opt}}$	Method	FID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
Uncond	64	64	64	4096	W-Flow	55.25	23.89	0.42	0.14
					UW-Flow ( $\tau=10$ )	55.00	25.01	0.40	0.16
					UW-Flow ( $\tau=1$ )	51.08	25.47	0.41	0.19
					UW-Flow ( $\tau=0.1$ )	46.17	27.71	0.40	0.29
Cond	64	64	64	4096	W-Flow	9.43	60.31	0.74	0.26
					UW-Flow ( $\tau=10$ )	9.47	61.08	0.75	0.25

Appendix

Table 8: Full ImageNet-100 generation results following the ImageNet ablation setup and CFG selection protocol of W-Flow ( Han et al., 2026 ) . Both settings use

N_{c}=64

N_{\mathrm{gen}}=N_{\mathrm{pos}}=64

, and

B_{\mathrm{opt}}=4096

. For class-conditional generation, we select the CFG scale with the lowest FID for each method from

\{1.0,1.5,2.0,2.5,3.0,3.5,4.0\}

Sinkhorn iterations $L$
$L$	FID $\downarrow$	IS $\uparrow$	Precision $\uparrow$	Recall $\uparrow$
$5$	5.04	10.39	0.57	0.41
$10$	4.73	10.54	0.58	0.45
$20$	5.16	10.26	0.59	0.40
$30$	4.51	10.37	0.60	0.45
$40$	5.18	10.29	0.57	0.45

Appendix

Table 9: Sinkhorn parameter ablations on CIFAR-10. We fix

N_{c}=2

N_{\mathrm{gen}}=N_{\mathrm{pos}}=64

B_{\mathrm{opt}}=128

, and

\tau=0.1

. The default configuration uses

L=30

and

\varepsilon=0.05

; each table varies one parameter while fixing the other.

Figure 7: Qualitative $B\times\tau$ ablation on checkerboard. Rows vary $B$ ; columns show the target, UW-Flow with increasing $\tau$ , and W-Flow.

Figure 8: Qualitative $B\times\tau$ ablation on two moons. Rows vary $B$ ; columns show the target, UW-Flow with increasing $\tau$ , and W-Flow.

Figure 9: Qualitative $B\times\tau$ ablation on Gaussian ring. Rows vary $B$ ; columns show the target, UW-Flow with increasing $\tau$ , and W-Flow.

One-Step Generative Modeling

One-step generative modeling maps noise or conditions to data in one network evaluation, replacing iterative denoising or numerical integration.

See all 142 papers

One-Step Generative Modeling via Unbalanced Optimal Transport
Sep 26, 2026Yirong Shen, Mengfei Xia, Junpeng Jing +2One-Step Generative Modeling Optimal Transport
Imperial College London · Ant Group · Brunel University of London
One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Wasserstein Gradient Flows Image Generation
Stanford University · ByteDance
Robustness and Structure Preservation in Flow-Based Generative Models via Wasserstein Path-Space Divergences
Oct 2, 2024Ziyu Chen, Markos A. Katsoulakis, Benjamin J. ZhangFlow Matching Equivariant Neural Networks
Department of Mathematics and School of Data and Information Sciences, University of North Carolina at Chapel Hill · Department of Mathematics and Statistics, University of Massachusetts Amherst · Department of Mathematics, Rutgers University–New Brunswick

Sep 26, 2026J/K move · Enter open · S save

One-Step Generative Modeling via Unbalanced Optimal Transport

Yirong Shen, Mengfei Xia, Junpeng Jing, Lu Gan+1

Imperial College London · Ant Group · Brunel University of London

One-Step Generative Modeling Optimal Transport Wasserstein Gradient Flows

Open paper AskPDF arXiv

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.

Abstract

One-Step Generative Modeling

Explore similar work

One-Step Generative Modeling via Unbalanced Optimal Transport

One-Step Generative Modeling via Wasserstein Gradient Flows

Robustness and Structure Preservation in Flow-Based Generative Models via Wasserstein Path-Space Divergences