Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
Tsinghua AIAcademy of Artificial IntelligenceAffiliation:ByteDanceFudanPKUXJTUZJUSep 28, 2026arXiv:2609.35763This paper establishes a unified theoretical framework for one-step visual generation by formulating feature-space distributional training as a Wasserstein gradient flow, demonstrating that prior methods like FD-Loss and Gaussian-kernel Drifting are simply Gaussian optimal transport and KDE-based KL matching special cases. To bridge the gap between crude global moments and sample-based representations, the authors introduce MGFlow, which models representations via Gaussian mixtures equipped with mass-constrained sample assignment to prevent mode collapse. The method achieves state-of-the-art ImageNet-256 synthesis (1.45 $\text{FDr}^6$) and distills FLUX.2 [klein] 4B into a one-step text-to-image generator that surpasses the original four-step model on GenEval and PickScore.
Single-step diffusion models can outperform their multi-step teachers when distributional matching is reformulated as a Wasserstein gradient flow over Gaussian mixtures.
\emph{Distributional training} provides collective supervision for one-step visual generation by matching real and generated features in frozen representation spaces. We introduce \emph{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/