Unifying Distributional Training for One-Step Visual Generation
MGFlow unifies one-step visual training and beats FD-Loss on ImageNet and four-step FLUX.2.
The paper presents a unified distributional-training framework for one-step visual generation that separates distribution modeling from matching discrepancy using Wasserstein gradient flow. Under the framework, FD-Loss and Gaussian-kernel Drifting emerge as special cases, motivating MGFlow, which models features with Gaussian mixtures and pairs mass-constrained assignment with component updates to limit mode collapse. On ImageNet 256x256, MGFlow reports 1.45 FDr^6 on pMF-H and 1.64 on JiT-H, surpassing the FD-Loss baseline. Post-training FLUX.2 [klein] 4B into a one-step generator outperforms the original four-step model on GenEval and PickScore.
50