Unifying Distributional Training for One-Step Visual Generation

What it claims

The paper argues that the choice in distributional training reduces to two axes: a distribution model M (single Gaussian, Gaussian mixture, or sample KDE) and a discrepancy D (W2 optimal transport or KL). Wasserstein gradient flow turns any (M, D) pair into a per-feature update field. On ImageNet 256×256, MGFlow, the Gaussian-mixture point on this grid, beats the FD-Loss baseline by 23% FDr6 on pMF-H (to 1.45) and by 38% on JiT-H (to 1.64). One-step FLUX.2 [klein] 4B post-trained with MGFlow reaches 0.900 GenEval and 21.98 PickScore. The authors say this outperforms the original four-step model and all previous methods.

How it works

• The framework's grid. Single Gaussian with W2 gives FD-Loss. Gaussian KDE with KL gives Gaussian-kernel Drifting. Empirical batches with OT gives W-Flow, which uses Sinkhorn. The grid also yields a new closed-form variant, Gaussian KL, with velocity v(z) = Σp⁻¹(μp − z) − Σq⁻¹(μq − z).
• The mixtures. The reference mixture P, with K components, is fit offline by EM and then frozen. The generated mixture Q is tracked with EMA moments, as in FD-Loss, and its weights are set equal to P's (ω = π).
• Why mixtures alone fail. The paper calls this score blindness. When modes are well separated, the density ratio Q/P is roughly constant inside each mode, so the KL velocity is close to zero even when mode weights are badly wrong. In their toy example, posterior assignment stays stuck at an 85.1% / 14.9% split.
• Fix 1: LP allocation. Each batch of B features is assigned to reference components by a capacity-constrained linear program. The cost is −log(πk pk(z)) and component k must receive exactly Bπk. The assignments R* are detached.
• Fix 2: paired matching. Generated component k is matched only to reference component k. For W2, the loss is Σk πk W2²(qk, pk), which upper-bounds W2²(Q, P) and needs K pair costs instead of K². The motivation is empirical: in a pMF-H run, the optimal component transport plan was diagonal at every refresh. For KL, the authors say pairing is needed for correctness, not only for speed.
• Encoders and sampling are FD-Loss's: three frozen encoders.

Evidence

Grid comparison (Table 1): pMF-B, Inception features, 10 epochs, aligned recipe. FDr6, lower is better.

Every alternative beats FD-Loss, even though none of them directly optimises the Fréchet distance.

Assignment ablation (Table 2): JiT-B, Inception features, 10 epochs, K=4, timed on 8 H200 GPUs. Values are FDr6 / seconds per step.

Without pairing, KL matching collapses. For W2, pairing roughly halves step time compared with LP-global.

How to use it

If you already train with FD-Loss, the paper's own grid gives three cheaper experiments, in increasing effort:

1. Swap in Gaussian KL. It is closed-form and keeps the single-Gaussian setup.
2. Move to a GM reference with LP allocation and paired W2.
3. Use paired KL with a larger K.

If you use KL with mixtures, use LP-paired. The other assignments are unusable with KL in Table 2.

The page gives a project page (https://shihaoyang0423.github.io/MGFlow-website/). It states no code, weights, licence, or install instructions.

Limits

• This is a preprint under review.
• The fetched text stops partway through Section 3.2, so the paired score-matching details and the later experiments are not visible. That includes the FLUX.2 training setup and cost, and the full ImageNet tables.
• Paired W2 relies on the diagonal-plan observation from one pMF-H run. The paper gives a sufficient condition only in an appendix that was not fetched.
• The P mixture must be fit with EM in advance, and K differs by discrepancy (4 for W2, 16 for KL). The visible text does not say how sensitive results are to K.
• Grid results are short 10-epoch runs on B-size models with Inception features only.