Unifying Distributional Training for One-Step Visual Generation

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 $\textbf{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 $\textbf{1.45}$ $\mathrm{FDr}^6$ on pMF-H and $\textbf{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. Project page: https://shihaoyang0423.github.io/MGFlow-website/

2. Document Verification & Archival Data

• Contributing Researchers: Chi Zhang, Haoyang Shi, Yueyi Liu, Ruichuan An, Junkang Zhou, Chang Li, Xiuyuan Lu, Yichi Zhang, Bo Wang, Yuhang Wu, Sen Cui, Miao Liu
• Submission Date: September 28, 2026
• Full Preprint Document: Download Official PDF
• Permanent Archive Record: arXiv:2609.35763v1

3. Academic Citation Reference

Standard Reference (APA Format):

Chi Zhang, et al. (2026). Unifying Distributional Training for One-Step Visual Generation. arXiv:2609.35763v1. https://arxiv.org/abs/2609.35763v1

Academic Field: Machine Learning | Document Identifier: arXiv:2609.35763v1

BibTeX Entry:

@article{arxiv_2609.35763v1,
  author    = {Chi Zhang and Haoyang Shi and Yueyi Liu and Ruichuan An and Junkang Zhou and Chang Li and Xiuyuan Lu and Yichi Zhang and Bo Wang and Yuhang Wu and Sen Cui and Miao Liu},
  title     = {{Unifying Distributional Training for One-Step Visual Generation}},
  journal   = {arXiv preprint arXiv:2609.35763v1},
  year      = {2026},
  url       = {https://arxiv.org/abs/2609.35763v1}
}