PaperlayerView exact revision on arXiv

Source-linked AI summary

Unifying Distributional Training for One-Step Visual Generation

Chi Zhang, Haoyang Shi, Yueyi Liu, Ruichuan An, Junkang Zhou, Chang Li, Xiuyuan Lu, Yichi Zhang, Bo Wang, Yuhang Wu, Sen Cui, Miao Liu

arXiv:2609.35763v1cs.LG

TL;DR

Distributional training provides collective supervision by comparing real and generated features in frozen representation spaces. The paper unifies how these distributions are modeled and matched, then introduces MGFlow, which combines Gaussian mixtures with constrained assignment and component matching. MGFlow improves on prior methods across ImageNet and text-to-image generation.

  • Problem

    Distributional training compares collections of real and generated features for collective supervision, but raises questions about how to model feature distributions and use their mismatch to guide learning.

    Source: §1 · p. 2

  • Method

    The paper connects distributional objectives to feature updates through Wasserstein gradient flow and introduces MGFlow, using Gaussian mixtures with mass-constrained assignment and paired component matching.

    Source: §1 · p. 2

  • Results

    Across ImageNet and text-to-image, MGFlow outperforms prior methods, with 1.45 FDr6 on pMF-H and a 23% margin over FD-Loss.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Both modeling granularity and component correspondence matter for one-step generation.

    Source: §6 · p. 10

    For score-based matching, mixture expressivity alone does not ensure effective transport between modes when the score field is weak.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

\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/

1 INTRODUCTION

Distributional training compares collections of real and generated features, raising questions about how to model their distributions and use mismatch to guide learning. The section presents a unified framework, motivates MGFlow, and reports improved generation results.

1 INTRODUCTION

  • Distributional training compares collections of real and generated features for collective supervision rather than assigning each output a fixed target image.Each feature’s gradient depends on shared statistics or interactions with other features.Source: §1 · p. 2
  • The unified framework separates distribution modeling from matching discrepancy and uses Wasserstein gradient flow to connect distributional objectives to pointwise feature updates.It provides a common formulation for global moment losses and local sample interactions.Source: §1 · p. 2
  • A Gaussian model with W2 optimal transport recovers FD-Loss, interpreting its Fréchet discrepancy between global feature statistics as a Gaussian transport cost.At the population level, differentiating through the moments moves features along the associated affine transport field.Source: §1 · p. 2
  • A sample-based density model with KL divergence recovers Gaussian-kernel Drifting, whose mean-shift difference is a density-level KL velocity on kernel density estimates up to a bandwidth-dependent scale.For a Gaussian kernel, each mean shift is proportional to the gradient of the log KDE.Source: §1 · p. 2
  • MGFlow uses Gaussian mixtures at adjustable granularity and combines mass-constrained sample assignment with paired component matching to guide generated features toward reference components.It supports both OT-based and score-based matching; mixture expressivity alone does not ensure mode coverage or correct transportation.Source: §1 · p. 2
  • Figure 2 organizes distributional training into sampling, encoding, distribution modeling, and matching discrepancy, and contrasts Gaussian, GM, and KDE representations.The matching objectives induce feature-space velocity fields through Wasserstein gradient flow.Source: Figure 2 · p. 3
  • 1.45 FDr6 on pMF-H and 1.64 on JiT-H improve over the FD-Loss baseline by 23% and 38%, respectively.On ImageNet 256 × 256, MGFlow achieves state-of-the-art results on both models.Source: §1 · pp. 2–3

2 A UNIFIED VIEW OF DISTRIBUTIONAL TRAINING

The framework separates how feature distributions are modeled from how they are matched, then uses Wasserstein gradient flow to connect global discrepancies to pointwise feature updates. Gaussian optimal transport and kernel-density KL matching recover FD-Loss and Gaussian-kernel Drifting as instances of this formulation.

2 A UNIFIED VIEW OF DISTRIBUTIONAL TRAINING

  • The framework represents distributional training through sampling, frozen encoding, a distribution model, and a matching discrepancy.Each feature update depends on the sampled population through the distribution model, rather than only on an individual target.Source: §2 · p. 3
  • Wasserstein gradient flow connects a global discrepancy to a pointwise descent field for generated features.Under suitable regularity, the induced field moves feature locations to decrease the distributional objective.Source: §2 · pp. 3–4
  • Moment-matched Gaussians with W2 recover FD-Loss, while Gaussian KDEs with KL recover Gaussian-kernel Drifting.The former uses a Gaussian optimal-transport field; the latter evaluates KL velocity on smoothed densities.Source: §2 · p. 4
  • Empirical-measure OT gives W-Flow, and retaining a single Gaussian model while using KL gives Gaussian KL.W-Flow uses a Sinkhorn approximation; Gaussian KL is introduced as a closed-form velocity field.Source: §2 · p. 4
  • All tested alternatives outperform FD-Loss on FDr6 in the aligned comparison.The methods use pMF-B with Inception features for 10 epochs.Source: Table 1 · p. 3, §2 · p. 4

3 MGFLOW

MGFlow places Gaussian-mixture modeling between global moment matching and sample-centered KDE, then combines component allocation with paired updates to preserve correspondence and address mode-transport failures. It supports OT- and score-based matching, with training procedures tailored to each.

3.1 FROM A SINGLE GAUSSIAN TO GAUSSIAN MIXTURES

  • Gaussian mixtures offer adjustable granularity between a single Gaussian and sample-centered KDE, while retaining componentwise means and covariances.The single-Gaussian endpoint models global first and second moments; sample-centered components with common isotropic covariance recover Gaussian KDE at the representation level.Source: §3 · p. 4, §3.1 · p. 5
  • Both W2 and KL matching extend to mixtures, but a global score field can be too weak to correct component-weight mismatches across well-separated modes.The paper notes this issue can also arise for Gaussian-KDE fields when smoothing leaves modes well separated.Source: §3.1 · p. 5

3.2 MASS-CONSTRAINED ALLOCATION AND PAIRED TRANSPORT

  • Generated features are assigned to fixed reference components by a capacity-constrained linear program that allocates Bπk mass to component k in a batch of B.The resulting assignments weight component statistics, which are maintained across batches with an exponential moving average.Source: §3.2 · pp. 5–6
  • LP-based allocation with paired updates correctly transports the toy samples, and the assignment ablation reports uniformly better results than posterior assignments.Paired transport is especially important for multi-step JiT models trained with score-based discrepancies; the authors suggest the large initial-to-target distribution gap may induce abnormal allocation.Source: §3.2 · pp. 6–7, Figure 3 · p. 6

3.3 TRAINING MGFLOW.

  • MGFlow trains OT- and score-based variants and uses fixed, discrepancy-specific encoder normalizers; text-to-image training matches joint image and frozen SigLIP2 text features.The fixed real–validation scales avoid downweighting harder-to-match feature spaces merely because their current discrepancies are large.Source: §3.3 · p. 7

4 EXPERIMENTS

MGFlow improves ImageNet post-training results over FD-Loss, while finer and hierarchical mixture modeling improves generation. On text-to-image generation, joint MGFlow matching outperforms the four-step FLUX.2 [klein] 4B model on GenEval and PickScore.

4.1 INCREASING GAUSSIAN COMPONENTS

  • Increasing mixture size from K = 1 to 16 lowers pMF-B FDr6 from 12.73 to 11.32 and FID from 1.95 to 1.46; hierarchical resolution further lowers FDr6.K = 1 + 4 + 16 combines global moment matching with progressively finer componentwise matching; larger mixtures are constrained by matching cost or component-statistic estimation.Source: §4.1 · p. 7

4.2 CLASS-CONDITIONED IMAGENET GENERATION

  • 1.45 FDr6 on pMF-H and 1.64 on JiT-H are state-of-the-art MGFlow-KL results, surpassing FD-Loss by 23% and 38%, respectively.Both MGFlow-W2 and MGFlow-KL surpass FD-Loss across model sizes; MGFlow-KL also outperforms FD-based baselines on held-out FDr3.Source: Table 4 · p. 8, §4.2 · p. 9

4.3 TEXT-TO-IMAGE GENERATION

  • With COCO references, joint matching reaches 0.900 GenEval versus 0.794 for the four-step backbone and achieves the table’s highest PickScore, 21.98.MGFlow achieves higher GenEval and PickScore with 44% and 33% fewer generated training samples, respectively.Source: Table 5 · p. 9, §4.3 · p. 9

5 RELATED WORK

Related work spans trajectory-based and distillation approaches to one-step generation, as well as distributional supervision through feature discrepancies or sample interactions. The paper presents its framework as unifying these approaches and distinguishes MGFlow by its Gaussian-mixture distribution model.

5 RELATED WORK

  • One-step generation methods learn trajectory endpoints, average velocities, or distill pretrained diffusion models.The paper also describes distributional training as a recent direction for supervising generated populations.Source: §5 · p. 10
  • Recent distributional methods use global feature-space discrepancies or construct updates through attraction–repulsion or transport between batches.Source: §5 · p. 10
  • The proposed framework recovers these related approaches, while MGFlow introduces a Gaussian-mixture distribution approximation distinct from them.Source: §5 · p. 10

6 CONCLUSION

The conclusion frames distributional training through separate modeling and matching choices, with Wasserstein gradient flow linking objectives to feature updates. The related-work discussion distinguishes trajectory generation from output-distribution matching and specifies where the Drifting connection applies.

6 CONCLUSION

  • The paper’s conclusion separates sampling, encoding, distribution modeling, and matching discrepancy, and says both modeling granularity and component correspondence matter for one-step generation.Experiments show improvements across training and held-out representations, with KL matching improving Fréchet metrics without directly optimizing them.Source: §6 · p. 10

A.1 ONE-STEP GENERATION AND DISTRIBUTION MATCHING

  • Trajectory-based methods determine how a generator produces an image in one step, whereas distributional post-training optimizes the distribution of outputs from an already trained one-step model.The paper’s experiments use pMF and a one-step initialization from JiT.Source: §A.1 · p. 16
  • DMD expresses a reverse-KL gradient through real and generated scores at noisy image distributions; MGFlow instead computes KL-matching scores from explicit Gaussian mixtures in frozen representation spaces.DMD uses pretrained and separately trained diffusion models for the scores; DMD2 removes the teacher-pair regression requirement.Source: §A.1 · p. 16
  • Kernel MMD can distinguish distributions beyond their first two moments, while a Gaussian Fréchet objective depends only on means and covariances.The paper says its componentwise moments and assignments allow the GM representation to distinguish modes that a single Gaussian cannot separate.Source: §A.1 · p. 16

A.2 DRIFTING AND ITS GRADIENT-FLOW INTERPRETATIONS

  • Drifting trains features toward detached targets shifted by an attraction–repulsion field during training, while inference uses one generator evaluation.Source: §A.2 · p. 16
  • The paper’s KDE–KL connection is specific to Gaussian-kernel Drifting; it does not identify every Drifting variant with the same KL gradient flow.It contrasts this with W-Flow’s Sinkhorn divergence between empirical measures and notes that kernel choice and normalization matter.Source: §A.2 · p. 16

A.3 REPRESENTATION-BASED DISTRIBUTIONAL POST-TRAINING

The related work contrasts representation-distribution methods and their estimators with MGFlow’s focus on distribution modeling and component updates. It also distinguishes mixture transport and score matching from transport between fitted mixtures and independently weighted scores.

A.3 REPRESENTATION-BASED DISTRIBUTIONAL POST-TRAINING

  • FD-Loss uses a closed-form Gaussian distance between real and generated feature statistics, while changing the discrepancy or using a Gaussian mixture changes the update or distribution representation.Its real statistics may be precomputed, with queues or exponential moving averages for generated statistics beyond a single batch.Source: §A.3 · p. 17
  • RDM methods use feature-distribution discrepancies; the paper focuses on distribution models and updates, comparing Gaussian, GM, and sample-based constructions and addressing GM assignment and correspondence.For text-to-image generation, it follows iRDM in concatenating image and text features for joint matching.Source: §A.3 · p. 17
  • AdvFD adapts a learned representation adversarially, whereas MGFlow keeps encoders fixed and refines the distribution model within each representation.Source: §A.3 · p. 17
  • AMFD amortizes affine denoising operators with neural networks, whereas MGFlow uses explicit full-covariance components and closed-form Gaussian scores or pair costs.MGFlow must assign new generated samples to components while keeping them matched to the reference mixture.Source: §A.3 · p. 17

A.4 MIXTURE TRANSPORT AND COMPONENTWISE SCORES

  • MGFlow’s sample-level LP assigns generated features to components with prescribed reference masses, unlike transporting mass between two already fitted mixtures.For KL matching, it pairs component scores rather than relying on independently weighted marginal scores, which can be insensitive to mixing proportions when components are well separated.Source: §A.4 · p. 17

B DERIVATIONS AND PROOFS

The derivations connect distributional discrepancies to feature-update fields, recovering Gaussian-kernel Drifting and FD-Loss as specific model–discrepancy choices. They also establish mixture matching and detached-target regression results, including a conditional image-score calculation for joint Gaussian features.

B.1 DERIVATION OF THE KDE–KL INTERPRETATION OF DRIFTING

  • Gaussian-kernel Drifting equals the KL Wasserstein velocity on Gaussian-kernel-smoothed real and generated densities, up to the bandwidth factor.The derivation distinguishes this field from varying a composite KL functional through the kernel convolution.Source: §B.1 · pp. 17–18

B.2 DERIVATION OF THE GAUSSIAN OT INTERPRETATION OF FD-LOSS

  • For Gaussian moment models, FD-Loss is a Gaussian W2 transport cost whose feature descent field moves generated features toward the reference Gaussian.For non-Gaussian generated distributions, the same moment gradient matches reference mean and covariance, not the full distribution.Source: §B.2 · p. 19

B.3 GLOBAL KL VELOCITY FOR GAUSSIAN MIXTURES

  • The global KL velocity for Gaussian mixtures is the difference between their mixture scores, computed using each mixture’s own posterior responsibilities.The score calculation assumes positive mixture weights and positive-definite component covariances.Source: §B.3 · p. 20

B.4 COMPONENT MATCHING AND PAIRED UPDATES

  • Paired component transport uses an assignment coupling, and diagonal transport is optimal when each diagonal component cost is no greater than its alternatives.Equal mixture weights alone do not ensure diagonal optimality; swapping equally weighted distinct components can make off-diagonal pairing optimal.Source: §B.4 · p. 20

2. Its steepest-descent velocity minimizes

  • For the paired KL energy, the steepest-descent component velocity does not acquire an extra mixture-weight factor.The derivation relates paired component KL to a marginal KL bound.Source: §2 · p. 21, §B.4 · p. 21

B.5 ENERGY DESCENT AND DETACHED-TARGET REGRESSION

  • The continuous-flow energy identity motivates detached-target regression, but it does not guarantee descent for finite-step parameter optimization.The regression transfers the field through encoder and generator Jacobians without differentiating through the estimated field.Source: §B.5 · pp. 21–22, §2 · p. 22

B.6 CONDITIONAL IMAGE SCORES IN A JOINT GAUSSIAN

  • For joint Gaussian image–text features, image–text cross-covariance affects the image update even though text features receive no gradient.The image score can be obtained by solving block equations rather than forming the full covariance inverse.Source: §B.6 · p. 22

C IMPLEMENTATION

The implementation fits fixed reference distributions and updates generated statistics alongside generator training, using mass-constrained assignments for mixture branches. It specifies loss scaling and task-specific training settings, and reports an EMA-schedule improvement in an aligned single-Gaussian FD run.

C.1 ALGORITHMS

  • The training algorithm reuses fixed LP assignments for component statistics and paired fields, with OT differentiating through candidate batch statistics and KL detaching its regression target.For limited-memory global batches, detached features are assigned once, then generator microbatches are replayed to propagate feature gradients.Source: §C.1 · p. 25

C.2 WEIGHTING MULTIPLE REPRESENTATION LOSSES

  • Encoder loss scales use real-training versus real-validation discrepancies, with one encoder weight shared across its component-count branches.OT uses inverse FD for scaling; KL uses inverse KL between single-Gaussian calibration fits.Source: §C.2 · p. 26

C.3.1 IMAGENET

  • ImageNet post-training uses six JiT/pMF backbones, frozen SigLIP2, Inception-v3, and MAE encoders, and 100 epochs; the additional-baselines table reports FDr6 and FDr3.The supplied table caption defines FDr6 as the mean normalized Fréchet distance over six encoders and FDr3 as excluding the three training encoders.Source: §C.3.1 · p. 27, Table 10 · p. 29

C.3.2 TEXT-TO-IMAGE GENERATION

  • Text-to-image post-training uses MGFlow-KL with equally weighted K = 1 and K = 4 branches, with image-only or joint image–text feature matching.Both variants use a shared reference set of 301,706 images and are evaluated with one Euler step.Source: §C.3.2 · p. 28, Table 8 · p. 28

C.4 STATISTICS EMA SCHEDULE

  • 4.87 versus 4.99 FDr6 after 50 epochs: the scheduled statistics EMA improves on constant 0.999 decay for aligned JiT-B Gaussian-FD training.The schedule starts at 0.995, rises to 0.999 by epoch 32, and then remains fixed.Source: §C.4 · p. 29, Table 9 · p. 29

E.1 THE 10-EPOCH DISTRIBUTION-MODEL COMPARISON

  • The 10-epoch comparison aligns six methods on pretrained pMF-B with frozen Inception features, one-step sampling, and FDr6 evaluation.It includes single-Gaussian FD and KL, MGFlow-W2 and MGFlow-KL, W-Flow, and Gaussian-kernel Drifting.Source: §E.1 · p. 30

C.1 ALGORITHMS

  • Reference Gaussian mixtures are fit offline and held fixed, while generated component statistics are initialized from samples and updated during training.For mixture branches, a capacity-constrained LP assigns batch samples according to prescribed reference weights.Source: §C.1 · pp. 23–24, §E.1 · p. 30

E.2 EVALUATION IN TRAINING AND HELD-OUT REPRESENTATIONS

The section defines evaluation across training and held-out representations, then examines fitted mixture structure, score-ridge choices, and component matching. Its findings show that component assignments vary by class granularity and encoder, while the best ridge setting depends on the training encoder.

E.2 EVALUATION IN TRAINING AND HELD-OUT REPRESENTATIONS

  • FDr6 averages normalized Fréchet distances across six evaluation spaces, while FDr3 averages the three held-out spaces for SIM training.The evaluation uses 50,000 generated images; Inception’s unnormalized Fréchet distance is FID.Source: §E.2 · pp. 30–31

F REFERENCE GMS AND COMPONENT STRUCTURE

  • Mean image-level peak responsibility remains above 99%, while class concentration decreases at K = 16, indicating within-class assignments span multiple components.The class–component association is measured across 200 sampled training images per ImageNet class.Source: F REFERENCE GMS AND COMPONENT STRUCTURE · p. 31, Table 12 · p. 32
  • At K = 4, SigLIP groups furniture and prepared food together but separates them at K = 16; Inception groups dogs with road vehicles while MAE and SigLIP separate them.Figure 4 reports mean posterior responsibilities for 13 label-defined groups; component indices are local to each fitted mixture.Source: F REFERENCE GMS AND COMPONENT STRUCTURE · p. 31, Figure 4 · p. 32

G SCORE RIDGE SETTINGS

  • 5.769 for MAE at λ = 0.001, 5.742 for SigLIP at λ = 0.03, and 10.720 for Inception at λ = 0.03 are the lowest reported FDr6 scores in the ridge sweep.For Inception, FID is lowest at λ = 0.01 rather than the λ = 0.03 that minimizes FDr6; sweep scores are not directly comparable to aligned runs.Source: G SCORE RIDGE SETTINGS · p. 33, Table 13 · p. 33

H COMPONENT MATCHING

  • The assignment ablation compares posterior assignments, capacity-constrained LP assignments with global updates, and LP assignments with paired component updates.All configurations use K = 4 without a K = 1 term.Source: H COMPONENT MATCHING · p. 34

I.1 TRAINING-TIME COMPARISON AND PROTOCOL

The appendix compares MGFlow-KL and FD-Loss step times, then details the computational and sample-support tradeoffs of increasing mixture component counts.

I.1 TRAINING-TIME COMPARISON AND PROTOCOL

  • MGFlow-KL runs faster per step than FD-Loss on all L/H backbones, but slightly slower on the two B backbones.The timing benchmark uses eight H200 GPUs and a global batch of 1,024.Source: §I.1 · p. 34, Table 15 · p. 34
  • The timing protocol measures the slowest rank over complete optimizer steps after warmup, while excluding compilation, initialization, checkpointing, evaluation, and log I/O.It requires at least ten warmup steps and a stable first twenty-step mean; training states and precision settings differ across configurations.Source: §I.1 · pp. 34–35

I.2 SCALING THE NUMBER OF COMPONENTS

  • Both paired variants scale as O(BKd^2 + Kd^3) + TLP(B,K), while KL avoids differentiating the current-batch covariance and saves spectral computation.The stated advantage is a shorter, cheaper backward path rather than a lower asymptotic order.Source: §I.2 · pp. 35–36
  • Larger MGFlow-W2 mixtures take 1.68×/1.95× as long per step on JiT-H and 1.26×/1.40× on pMF-H relative to K = 1 + 4.These costs motivate the default K = 1 + 4 for MGFlow-W2.Source: §I.2 · p. 36, Table 17 · p. 36
  • At K = 64, Inception’s smallest components have only 978–1,358 units of reference responsibility mass in 2,048 dimensions, motivating K = 16 as MGFlow-KL’s largest branch.Small components provide weak support for full-covariance estimates; longer EMA windows also slow adaptation to the changing generator.Source: §I.2 · p. 36
  • For text-to-image generation, K = 16 leaves Inception components with insufficient support for stable full-covariance estimation, so both variants use K = 1 + 4.At K = 4, both weight ratios are below 6.2 and the smallest components each have over 25,000 units of reference responsibility mass.Source: §I.2 · p. 37

J RESULTS IN INDIVIDUAL EVALUATION REPRESENTATIONS

In individual evaluation representations, MGFlow-KL improves over FD-Loss on DINOv2 and CLIP for pMF-L/H, while FD-Loss remains lower on ConvNeXt.

J RESULTS IN INDIVIDUAL EVALUATION REPRESENTATIONS

K QUALITATIVE IMAGENET RESULTS

The appendix presents uncurated one-step ImageNet samples from JiT-H and pMF-H, plus additional text-to-image samples from MGFlow-post-trained FLUX.2 [klein] 4B.

K QUALITATIVE IMAGENET RESULTS

  • JiT-H ImageNet samples compare FD-Loss, MGFlow-W2, and MGFlow-KL at one NFE using identical input noise at corresponding positions.Figure 6 provides additional uncurated samples with corresponding images using identical input noise.Source: Figure 5 · p. 38, Figure 6 · p. 39
  • pMF-H ImageNet samples compare FD-Loss, MGFlow-W2, and MGFlow-KL at one NFE using identical input noise at corresponding positions.Figure 8 provides additional uncurated samples with corresponding images using identical input noise.Source: Figure 7 · p. 40, Figure 8 · p. 41

L ADDITIONAL TEXT-TO-IMAGE EXAMPLES

  • Additional examples show selected 512 × 512 images from FLUX.2 [klein] 4B post-trained with MGFlow for 1,000 steps.Source: Figure 9 · p. 42

Related papers

Loading 2609.35763v1…