Mixture Gradient Flow

Unifying Distributional Training
for One-Step Visual Generation

A common gradient-flow view of distributional training.
Gaussian mixtures capture multimodal feature distributions with a compact, adjustable representation.

Explore samples Model weights

Anonymous project page

One step. Many possibilities.

One forward pass. FLUX.2 [klein] 4B post-trained with MGFlow · Selected samples

1 NFE

One-step sampling

1.45 ↓

FDr⁶ · pMF-H · ImageNet

0.900 ↑

GenEval · FLUX.2 [klein] 4B

21.98 ↑

PickScore · FLUX.2 [klein] 4B

01 / Sample explorer

From a single step.

Browse selected generated images.
Choose an image to view it at full resolution.

Text-to-image 20 selected images

FLUX.2 [klein] 4B + MGFlow · 512 × 512 · 1 NFE

Loading samples…

This is a gallery of pre-generated, selected qualitative examples from the paper, not a live inference service or an unbiased sample set.

02 / The idea

One framework.
Different distribution models.

Distributional training matches generated and reference features. Its update depends on how we represent the distributions and how we compare them.

𝒮

Batch sampling

ℰ

Frozen encoders

ℳ

Distribution model

𝒟

Matching objective

A unified distributional training framework: FD-Loss models a single Gaussian, MGFlow uses Gaussian mixtures, and Drifting uses kernel density estimates. Distribution matching supplies a Wasserstein gradient-flow update. — MGFlow uses Gaussian mixtures to capture multimodal structure between a single-Gaussian approximation and sample-based representations.

Distribution models and matching objectives
Objective	Gaussian	Gaussian mixture	Sample-based
OT-based	FD-LossYang et al., 2026 ↗	MGFlow-W₂	W-FlowHan et al., 2026 ↗
Score-based (KL)	Gaussian KL	MGFlow-KL	Gaussian-kernel DriftingDeng et al., 2026 ↗

W-Flow uses empirical measures and Sinkhorn divergence; Gaussian-kernel Drifting uses kernel-smoothed densities. MGFlow-W₂ uses paired Gaussian transport costs, and MGFlow-KL uses paired component score fields.

Match the components,
not just their local scores.

A more expressive density model alone does not ensure stable matching. Soft assignments can lose components, while global mixture scores can miss incorrect mass allocation between separated modes.

MGFlow constrains each component's batch mass with a linear program, updates its statistics, and matches corresponding reference and generated components. This pairing is used for both Gaussian transport and KL-based updates.

Toy experiment comparing posterior-assigned global scores, LP-assigned global scores, and LP-paired updates over time. The LP-paired update aligns component positions and reference mass allocation. — Component assignment and the choice of score field both matter. The toy experiment compares soft-assigned global, LP-assigned global, and LP-paired updates.

03 / Results

Image quality. Prompt alignment.
One-step generation.

Selected comparisons from the paper.
↓ Lower is better. ↑ Higher is better.

ImageNet 256 × 256

JiT-H & pMF-H

100 epochs · SigLIP, Inception & MAE encoders

Method	NFE	FDr⁶ ↓	FDr³ ↓	FID ↓
JiT-HLi & He, 2026 ↗
Pretrained	50×2×2^†	7.66	9.07	1.97
FD-LossYang et al., 2026 ↗	1	2.65	4.44	0.75
AdvFDGao et al., 2026 ↗	1	1.80	2.93	0.72
AMFD-CLiu et al., 2026 ↗	1	2.15	3.10	0.85
AMFD-ULiu et al., 2026 ↗	1	1.79	2.78	0.83
MGFlow-W₂	1	2.55	4.42	0.99
MGFlow-KL	1	1.64	2.57	0.94
pMF-HLu et al., 2026 ↗
Pretrained	1	6.87	6.09	2.29
FD-LossYang et al., 2026 ↗	1	1.89	2.69	0.77
AdvFDGao et al., 2026 ↗	1	1.74	2.50	0.74
AMFD-CLiu et al., 2026 ↗	1	1.93	2.43	0.86
AMFD-ULiu et al., 2026 ↗	1	1.75	2.30	0.85
MGFlow-W₂	1	1.50	2.05	1.09
MGFlow-KL	1	1.45	1.88	1.07

FDr⁶ averages normalized Fréchet distances across six encoders; FDr³ uses three held-out encoders. ^† Full-CFG NFE upper bound for interval CFG. Bold marks the best value within each backbone among the methods shown.

Text-to-image 512 × 512

FLUX.2 [klein] 4B

MGFlow: 1,000 training steps · K = 1 + 4

Method	NFE	GenEval ↑	PickScore ↑
FLUX.2 [klein]Black Forest Labs, 2026 ↗	4	0.794	21.85
One-step initialization	1	0.535	20.10
DMD2Yin et al., 2024 ↗	1	0.804	—
FD-SIMYang et al., 2026 ↗	1	0.802	21.62
iRDMFeng et al., 2026 ↗	1	0.826	21.82
AMFD-U-SIMLiu et al., 2026 ↗	1	0.813	21.77
AMFD-C-SIMLiu et al., 2026 ↗	1	0.846	21.85
AMFD-C-10 enc.Liu et al., 2026 ↗	1	0.845	21.82
MGFlow · image-only	1	0.856	21.86
MGFlow · joint	1	0.900	21.98

The one-step initialization is evaluated before MGFlow post-training. Joint matching uses image and SigLIP2 text features. PickScore uses 499 Pick-a-Pic prompts. Published baseline values follow AMFD, Table 3; — denotes an unreported value.

04 / Models & data

Build on MGFlow.

Six ImageNet backbones, two objectives.
Two text-to-image variants, training data, and references.

ImageNet

Post-trained JiT and pMF models at B, L, and H scales. MGFlow-W₂ uses K = 1 + 4; MGFlow-KL uses K = 1 + 4 + 16.

Model checkpoints

BackboneObjective

Selected post-trained checkpointpMF-H · MGFlow-KL weights · 1 step Selected base checkpointpMF-H · pretrained weights

Training resources

Fitted Gaussian referencesSigLIP, Inception, MAE · K = 1, 4, 16

Text-to-image

FLUX.2 [klein] 4B post-trained with MGFlow-KL. Image-only and joint image–text distribution matching use the same reference construction.

Model checkpoints

Joint image–text checkpointFLUX.2 [klein] 4B weights · 1 step Image-only checkpointFLUX.2 [klein] 4B weights · 1 step

Training resources

T2I training imagesReconstructed COCO + GenEval reference images Fitted Gaussian referencesImage-only and joint features · K = 1 and 4

Model checkpoints contain trained weights and the training step. See the model repository for the current release and loading information.

Unifying Distributional Training for One-Step Visual Generation

From a single step.

One framework.Different distribution models.

Match the components,not just their local scores.

Image quality. Prompt alignment.One-step generation.

ImageNet 256 × 256

Text-to-image 512 × 512

Build on MGFlow.

ImageNet

Model checkpoints

Training resources

Text-to-image

Model checkpoints

Training resources

Unifying Distributional Training
for One-Step Visual Generation

One framework.
Different distribution models.

Match the components,
not just their local scores.

Image quality. Prompt alignment.
One-step generation.