Unifying Distributional Training for One-Step Visual Generation

Chi Zhang; Shi Haoyang; Yueyi Liu; Ruichuan An; Junkang Zhou; Chang Li; Xiuyuan Lu; Yichi Zhang

cs.LGSep 28, 2026

Unifying Distributional Training for One-Step Visual Generation

Authors: Chi Zhang, Shi Haoyang, Yueyi Liu, Ruichuan An, Junkang Zhou, Chang Li, Xiuyuan Lu, Yichi Zhang, +4 more

Organizations: Tsinghua University · Fudan University · Xi’an Jiaotong University · Peking University · Zhejiang University · University of Science and Technology, China · University of California, Berkeley

Abstract

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\times256$ , MGFlow substantially surpasses the FD-Loss baseline, achieving state-of-the-art results with 1.45 $\mathrm{FDr}^6$ 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.

Figures & tables

Figure 1 : A unified view of distributional training. Left: sampling $\mathcal{S}$ , encoding $\mathcal{E}$ , distribution modeling $\mathcal{M}$ , and matching discrepancy $\mathcal{D}$ define the paradigm. Right: FD-Loss, MGFlow, and Gaussian-kernel Drifting use Gaussian, GM, and KDE representations, respectively. Their matching objectives induce feature-space velocity fields through the Wasserstein gradient flow of their respective distributional objectives.

Panel of Table 1 — Table 1 : Distributional training with distribution models and discrepancies. All methods train pMF-B with Inception features for 10 epochs. All methods use the aligned recipe in Appendix E.1 .

Panel of Table 2 — Table 2 : Ablation on component assignment and transportation. JiT-B trains in Inception feature space for 10 epochs with $K=4$ . Timing uses 8 H200 GPUs.

Figure 2 : Toy experiment with component assignment and transport. Score matching with

K=2

moves 2048 particles. Dashed contours show the reference

P

. Shading and contour counts indicate weights. LP fixes component weights, but global velocity collapses modes. Only LP-paired achieves correct transport. Detailed settings are given in Appendix H .

Components	$\text{FDr}^{6}\!\downarrow$	$\text{FDr}^{5}\!\downarrow$	FID $\!\downarrow$
$K=1$	12.73	15.04	1.95
$K=4$	12.28	14.54	1.65
$K=16$	11.32	13.41	1.46
$K=1+4$	11.51	13.59	1.89
$K=1+4+16$	11.05	13.08	1.51

Table 3 : Component-count ablation for MGFlow-KL. All configurations use pMF-B trained in Inception feature space for 10 epochs.

\text{FDr}^{5}

evaluates five held-out encoders.

Panel of Table 4 — Table 4 : Post-training JiT and pMF on ImageNet $\bm{256\times 256}$ . All post-training runs use SIM encoder spaces for 100 epochs. “Model” denotes feature-distribution modeling. Baseline results are from FD-Loss ( Yang et al., 2026a ) , AdvFD ( Gao et al., 2026 ) , and AMFD ( Liu et al., 2026a ) . † denotes the full-CFG NFE upper bound for interval CFG; dashes denote unavailable results. Best and second-best results per backbone are bold and underlined, respectively. More baselines are in Table 10 .

Panel of Table 5 — Table 5 : Text-to-image generation with FLUX.2 [klein] 4B ( Black Forest Labs, 2026 ) at $\bm{512\times 512}$ . We report GenEval ( Ghosh et al., 2023 ) and PickScore ( Kirstain et al., 2023 ) ; PickScore is evaluated on 499 Pick-a-Pic prompts. Higher is better. Bold and underlining mark the best and second-best scores, respectively. Baselines are from AMFD ( Liu et al., 2026a ) ; dashes denote unavailable results.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Table 6 — Table 6: Single-Gaussian FD with fixed reference normalization. $\text{FDr}^{6}\downarrow$ after 100 epochs with SIM encoders. FD-Loss results are from Yang et al. (2026a) ; the aligned runs use our statistics-update recipe.

Panel of Table 7 — Table 7 : ImageNet post-training configurations. Slash-separated pMF sampling values follow B/L/H order.

Panel of Table 8 — Table 8 : Text-to-image post-training configurations. Both variants use reconstructed COCO and GenEval reference images.

Panel of Table 9 — Table 9: Statistics EMA decay. JiT-B trained for 50 epochs with aligned Gaussian-FD, fixed $\mathrm{FD}(R,V)$ normalization, and SIM encoders.

Panel of Table 10 — Table 10 : Additional ImageNet $256{\times}256$ baselines. Published results from FD-Loss ( Yang et al., 2026a ) , AdvFD ( Gao et al., 2026 ) , and AMFD ( Liu et al., 2026a ) . FDr 6 averages normalized Fréchet distances over six encoders; FDr 3 excludes SigLIP, Inception, and MAE. † denotes the full-CFG NFE upper bound for interval CFG; dashes denote unavailable results. JiT/pMF results are in Table 4 .

Encoder	Variant	Input	Dimension	Feature	$b_{e}$
SigLIP2 ( Tschannen et al., 2025 )	ViT-SO400M/16	224	1,152	Attention pool	0.60
Inception-v3 ( Szegedy et al., 2016 )	FID weights	299	2,048	Global average pool	1.68
MAE ( He et al., 2022 )	ViT-L/16	224	1,024	CLS token	0.04
ConvNeXt-v2 ( Woo et al., 2023 )	Base, IN-22K $\to$ IN-1K	224	1,024	Global average pool	56.87
DINOv2 ( Oquab et al., 2023 )	ViT-L/14	256	1,024	CLS token	14.19
CLIP ( Radford et al., 2021 )	ViT-L/14, OpenAI	256	1,024	CLS token	5.60

Appendix

Table 11 : ImageNet representation encoders. Input denotes the resized image side length before patch padding. The first three encoders form SIM; the last three define FDr 3 . The fixed denominator

b_{e}

is used to compute

\text{FDr}_{e}=F_{e}/b_{e}

Figure 3 : Geometry of the ImageNet reference distributions. Within each encoder, all three models share a Fisher projection derived from the $K=16$ reference and the same axis limits. Colors show maximum-posterior assignments in the original feature space and are local to each mixture. Ellipses are 90% probability contours of the projected Gaussian components; the larger dots mark their means.

Encoder	$K$	Mean $q_{c}$	Classes with $q_{c}\geq 0.9$	Image peak
MAE	4	81.24%	398	99.88%
SigLIP	4	91.12%	744	99.96%
Inception	4	82.23%	461	99.97%
MAE	16	64.79%	113	99.79%
SigLIP	16	71.88%	229	99.94%
Inception	16	77.21%	305	99.96%

Appendix

Table 12: Class concentration versus image-level assignment. We use 200 images per ImageNet class. Class concentration

q_{c}

is the mean responsibility of class

c

for its dominant component; “Image peak” averages the maximum responsibility of each image.

Figure 4 : Semantic groups across independently fitted mixtures. Rows show mean posterior responsibilities for 13 label-defined groups covering 417 ImageNet classes, using 200 images per class. We average

r_{ck}

equally over the classes in each group; parentheses give the number of classes. Component indices are local to each GM, not aligned across panels.

Panel of Table 13 — Table 13: Single-encoder ridge selection. pMF-B trained with MGFlow-KL, $K=1$ , for 10 epochs. All metrics are lower-is-better. Bold marks the lowest FDr 6 for each training encoder.

Panel of Table 14 — Table 14: ImageNet covariance eigenvalue quantiles and score ridges. Quantiles of the $K=1$ reference spectra use the inverse empirical CDF. Percentile is the fraction of eigenvalues at or below the training ridge $\lambda_{1}$ .

Panel of Table 15 — Table 15: T2I covariance eigenvalue quantiles and score ridges. The $K=1$ references use 301,706 COCO and GenEval images. Matched denotes the eigenvalue at the encoder-specific ImageNet percentile; $\lambda_{1}$ is the value used in training.

Panel of Table 16 — Table 16: Training time per step. Seconds per optimizer update on eight H200 GPUs with a global batch of 1,024. MGFlow- $W_{2}$ uses $K=1+4$ ; MGFlow-KL uses $K=1+4+16$ . Training states and precision settings are specified below. Lower is better.

Backbone	MGFlow- $W_{2}$	MGFlow-KL	Official FD-Loss
JiT-B	1	1	1
JiT-L	1	1	1
JiT-H	1	1	2
pMF-B	1	1	1
pMF-L	1	1	2
pMF-H	2	2	4

Appendix

Table 17: Gradient accumulation in the timing benchmark. Entries are

a

; the per-GPU microbatch is

128/a

and the global batch remains 1,024. Values above one follow a recorded CUDA OOM at the preceding larger microbatch.

Table 22

Figure 5 : Uncurated ImageNet $\bm{256\times 256}$ samples from JiT-H. FD-Loss ( Yang et al., 2026a ) (left), MGFlow- $W_{2}$ (middle), and MGFlow-KL (right), all with one NFE and identical input noise at corresponding positions.

Figure 6 : Additional uncurated ImageNet $\bm{256\times 256}$ samples from JiT-H. FD-Loss (left), MGFlow- $W_{2}$ (middle), and MGFlow-KL (right), all with one NFE. Corresponding images use identical input noise.

Figure 7 : Uncurated ImageNet $\bm{256\times 256}$ samples from pMF-H. FD-Loss ( Yang et al., 2026a ) (left), MGFlow- $W_{2}$ (middle), and MGFlow-KL (right), all with one NFE and identical input noise at corresponding positions.

Figure 8 : Additional uncurated ImageNet $\bm{256\times 256}$ samples from pMF-H. FD-Loss (left), MGFlow- $W_{2}$ (middle), and MGFlow-KL (right), all with one NFE. Corresponding images use identical input noise.

Panel of Figure 9 — Figure 9 : One-step text-to-image samples. Selected $512\times 512$ images from FLUX.2 [klein] 4B post-trained with MGFlow using joint image–text matching for 1,000 steps.

Explore similar work

One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Diffusion Models Generative Models
Stanford University · ByteDance
Representation Distribution Matching for One-Step Visual Generation
Jul 2, 2026Lan Feng, Wuyang Li, Eloi Zablocki +2Multi-Reference Image Generation Visual Generation
1EPFL, Switzerland · 2Valeo.ai, France · 3Sorbonne Universit´e, France
Perceptual Flow Matching for Few-Step Generative Modeling
Jul 3, 2026Chuyang Zhao, Yifei Song, Hongfa Wang +7Contrastive Flow Matching Few-Step Generation
1Joy Future Academy · 2Fudan University · 3Tsinghua University +1

May 12, 2026, cs.LGJ/K move · Enter open · S save

One-Step Generative Modeling via Wasserstein Gradient Flows

Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu+2

Stanford University · ByteDance

Diffusion Models Generative Models Wasserstein Distance Imagenet Gradient

Open paper AskPDF arXiv

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

\times

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

\times

faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.

Abstract

Explore similar work

One-Step Generative Modeling via Wasserstein Gradient Flows

Representation Distribution Matching for One-Step Visual Generation

Perceptual Flow Matching for Few-Step Generative Modeling