Improved Mean Flows: On the Challenges of Fastforward Generative Models

TL;DR

iMF revisits the MeanFlow (MF) framework for one-step generative modeling and identifies two structural defects: (i) MF’s training target depends on the network’s own prediction, so it isn’t a standard regression problem, and (ii) MF freezes the CFG scale at training time, throwing away inference-time flexibility. The fix re-parameterizes the MF objective as a v-loss with u-pred — a network-independent target that takes only the noisy sample z as input — and treats CFG scale as a conditioning variable processed via multi-token in-context conditioning. Trained from scratch with no distillation or pretrained alignment, iMF-XL/2 reaches 1-NFE FID 1.72 on ImageNet 256², a ~50% relative improvement over MF-XL/2’s 3.43 and a new SOTA among 1-NFE-from-scratch diffusion/flow methods.

Key claims

The original MF objective is equivalent to a v-loss re-parameterized by u-pred, where the compound prediction function takes both the noisy sample z and the conditional velocity (e − x) as input — making it not a standard regression problem [§4.1, Fig. 1(a), Eq. 11].
Replacing the conditional velocity input with a network-predicted marginal (u_θ) yields a legitimate regression: the compound function takes only z as input and the training loss has both lower variance and is monotonically decreasing where original MF’s is non-decreasing [§4.1, Fig. 1(b), Fig. 3, Eq. 12].
Two realizations of the inner u_θ predictor work: a parameter-free boundary-condition reuse (u(z, t, t) ≡ v_θ) or an auxiliary u-head used only at training time. Boundary-condition gives the larger gain on MF-B/2 (FID 32.69 → 29.42 w/o CFG); the auxiliary head scales better with model size and is preferred with CFG (FID 6.17 → 5.68) [§4.1, Table 1(a)].
The original MF freezes the CFG scale ω at training time; the optimal ω depends on model size, training length, and inference NFE, so fixing it a priori is suboptimal [§4.2, Fig. 4].
Reformulating CFG scale as an explicit conditioning variable (analogous to time-step embedding) lets the model interpolate guidance at inference. Extending the conditioning to (ω, t_low, t_high) also supports CFG interval at inference, improving FID 6.17 → 4.57 on the B-size model [§4.2, Eq. 15, Table 1(b)].
Multi-token in-context conditioning (8 tokens per class label, 4 tokens per other condition) lets iMF drop adaLN-zero entirely, cutting the B-size model from 133M to 89M parameters while improving FID from 4.57 to 4.09 — the same recipe DiT dismissed works once multiple tokens per condition are used [§4.3, Fig. 5, Table 1(c)].
iMF-XL/2 (610M) achieves 1-NFE FID 1.72 / IS 282 on ImageNet 256² from scratch, substantially better than MF-XL/2 (676M, 3.43 / 247.5) and ahead of α-Flow-XL/2+ (676M, 2.58) and TiM-XL/2 (664M, 3.26) at the same NFE class [Table 2 — model comparison].
At 2-NFE, iMF-XL/2 reaches FID 1.54, closing most of the gap to multi-step distillation-based 1-NFE methods (FACM-XL/2 1.76, DMF-XL/2+ 2.16) without using any distillation or pretrained teacher [Table 2 — comparison block].

Method

iMF is a re-derivation, not a new architecture: it keeps MeanFlow’s average-velocity field u(z, t, r) but rewrites the training loss. Starting from the MeanFlow identity u = v − (t − r)·d/dt u, the authors note that v on the LHS is a network-independent target (the true marginal velocity, supervisable by the standard Flow-Matching loss), so MF can be written as a v-loss with the inner term re-parameterized through u_θ. The remaining issue is that the parameterized compound function f̂(z, e − x) still takes the high-variance conditional velocity (e − x) as input via the JVP tangent vector. iMF replaces (e − x) with u_θ(z, t, t) — either the boundary value of the same network (zero extra parameters) or a small auxiliary head — so the compound function takes only z.

For guidance, the fixed-ω CFG of MF is replaced with ω as a conditioning variable, identical in shape to time-step conditioning. Each CFG-related scalar gets 4 learnable tokens (8 for class labels) that are concatenated along the sequence axis with the image tokens and processed by Transformer blocks. This in-context multi-token conditioning removes the parameter-heavy adaLN-zero entirely, giving a 1/3 model-size reduction at fixed depth/width.

Training uses the same code, VAE tokenizer, and ImageNet-256² recipe as the original MF; iMF-XL/2 has 48 blocks at width 1024 (610M, 174.6 Gflops), trained from scratch with no distillation, no REPA-style alignment, and no pretrained model.

Results

ImageNet 256² 1-NFE FID (from scratch, with CFG): iMF-XL/2 1.72 vs MF-XL/2 3.43 (−50% relative), α-Flow-XL/2+ 2.58, TiM-XL/2 3.26, Shortcut-XL/2 10.60, iCT-XL/2 34.24 [Table 2].
Ablation on MF-B/2 baseline (240 epochs, w/ CFG): original MF 6.17 → +iMF u_θ via boundary 5.97 → +aux head 5.68 → +ω-conditioning 5.52 → +(ω, t_low, t_high)-conditioning 4.57 → +in-context conditioning (133M → 89M) 4.09 → +SwiGLU/RMSnorm/RoPE 3.82 → +longer training (640 ep) 3.39 [Table 1].
The same boundary-condition reformulation generalizes to XL: on MF-XL/2 with CFG, the boundary-condition variant alone drops FID 3.43 → 2.99 (−13% relative), and the authors hypothesize larger models benefit more from the legitimate prediction-function formulation [§5.1, “Comparison and Analysis”].
CFG-conditioning even improves the “w/o CFG” inference path: setting ω=1 at inference (no guidance) on a model trained with flexible ω improves FID from 30.76 → 20.95 on the B-size model — training across a range of guidance scales also acts as regularization [§5.1, Table 1(b)].
At 2-NFE, iMF-XL/2 reaches FID 1.54, beating MeanFlow-XL/2+ 2.20 and α-Flow-XL/2+ 1.95 [Table 2].

Why it’s interesting

iMF is the cleanest “from-scratch frontier” datapoint in the fastforward-generative-model line: it shows that two of MF’s limitations are formulation bugs, not capacity bugs, and that fixing them closes most of the gap to multi-step methods without the distillation crutch that Cross-Resolution Distribution Matching for Diffusion Distillation, Streaming Autoregressive Video Generation via Diagonal Distillation, and TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward depend on. It sharpens the Diffusion Distillation cluster’s open question of “teacher-free vs teacher-bound” by setting a higher from-scratch bar at 1-NFE FID 1.72 — close enough to teacher-bound 1-NFE numbers (FACM-XL/2 1.76) that the distillation premium may be smaller than assumed. The CFG-as-conditioning recipe also lands cleanly in the Classifier-Free Guidance variants cluster: it’s the training-time flexibility answer that complements SoFlow: Solution Flow Models for One-Step Generative Modeling‘s training-time-folded CFG and Normalized Attention Guidance: Universal Negative Guidance for Diffusion Models‘s inference-only attention-space patch, and it directly supersedes the fixed-ω constraint inherited from Mean Flows for One-step Generative Modeling.