Unifying Distributional Training for One-Step Visual Generation: MGFlow

Forum topic · 小凯 · 2026-09-30

Summary

A paper (arXiv:2609.35763) introduces a unified theoretical framework for distributional training in one-step visual generation, separating distribution modeling from matching discrepancy and linking global objectives to pointwise feature updates via Wasserstein gradient flow. Under this framework, FD-Loss and Gaussian-kernel Drifting are recovered as special cases through Gaussian optimal transport and kernel-density-based KL matching. The framework motivates MGFlow, which models feature distributions with Gaussian mixtures at 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 pairwise component updates to address mode collapse. On ImageNet 256x256, MGFlow substantially surpasses FD-Loss baselines, achieving a state-of-the-art 1.45 FDr^6 with pMF-H and 1.64 with JiT-H. For text-to-image generation, post-training FLUX.2 [klein] 4B into a one-step generator outperforms the original four-step model on both GenEval and PickScore.

Paper Overview

Field: Machine Learning Authors: Chi Zhang, Haoyang Shi, Yueyi Liu, Ruichuan An, Junkang Zhou, Chang Li, Xiuyuan Lu, Yichi Zhang, Bo Wang, Yuhang Wu, Sen Cui, Miao Liu Published: 2026-09-28 arXiv: 2609.35763

Abstract

*Distributional training* provides collective supervision for one-step visual generation by matching real and generated features in frozen representation spaces. The authors 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 pairwise component updates to address mode collapse that mixture expressiveness alone cannot resolve.

Results

ImageNet 256x256: MGFlow substantially surpasses FD-Loss baselines, achieving state-of-the-art 1.45 FDr^6 with pMF-H and 1.64 with JiT-H.
Text-to-image: Post-training FLUX.2 [klein] 4B into a one-step generator outperforms the original four-step model on both GenEval and PickScore.