Mean Flows for One-step Generative Modeling

TL;DR

MeanFlow trains a one-step image generator from scratch by modelling the average velocity along a flow path — the displacement between two time steps divided by the interval — instead of Flow Matching’s instantaneous velocity. A simple identity (the “MeanFlow Identity”) rewrites the average-velocity target in a form that needs only the instantaneous velocity plus a Jacobian-vector product, so training reduces to Flow Matching with a modified target and one extra JVP per step. On ImageNet 256² it hits 3.43 FID at 1-NFE with a 676M DiT-XL/2, with no pre-training, distillation, or curriculum — roughly halving the prior 1-NFE FID and approaching the 2.06 FID of 250×2-NFE SiT-XL/2. CFG is folded into the ground-truth field rather than applied at inference, preserving single-NFE sampling.

Key claims

Average velocity u(z_t, t, r) = (1/(t-r)) ∫_r^t v(z_τ, τ) dτ is a well-defined ground-truth field induced by the instantaneous velocity v, jointly parameterised by two times, and satisfies built-in boundary (u → v as t → r) and additive consistency conditions without an extra consistency loss [§4.1].
The MeanFlow Identity u(z_t, r, t) = v(z_t, t) - (t-r) · d/dt u(z_t, r, t) follows directly from the definition of u and contains no integral, making it a viable training target [§4.1, Eq. 6].
The total time derivative on the RHS decomposes via JVP as d/dt u = ∂_t u + v · ∂_z u, computable in a single backward-equivalent pass with torch.func.jvp / jax.jvp, with stop-gradient on the target avoiding double-backprop [§4.1, Eq. 8].
Optimal loss configuration: 25% of training samples set r = t (collapsing to standard Flow Matching), logit-normal lognorm(-0.4, 1.0) time sampler, embed (t, t-r), adaptive loss weight with exponent p = 1.0 (matches Pseudo-Huber); destructive ablations of the JVP tangent confirm both 1-D time tangents are essential [Table 1, Fig. 4].
Classifier-Free Guidance is absorbed into the ground-truth field as v^cfg = ω·v_cond + (1-ω)·v_uncond and a corresponding u^cfg derived through the same identity, so a single network forward gives the guided sample — CFG without doubling NFE [§4.2, Eqs. 13–19].
On ImageNet 256² class-conditional, 1-NFE FID-50K scales smoothly: B/2 6.17, M/2 5.01, L/2 3.84, XL/2 3.43 [Table 2].
MeanFlow-XL/2 at 2-NFE reaches 2.20 FID, matching 250×2-NFE DiT-XL/2 (2.27) and SiT-XL/2 (2.06) under the same backbone [Table 2].
Reduces 1-NFE FID by 50–70% relative to prior one-step / 1-NFE methods (iCT-XL/2 34.24, Shortcut-XL/2 10.60, IMM 7.77) at matched DiT-XL/2 capacity [Fig. 1, Table 2].
JVP overhead is <20% of total training time in the JAX implementation, so the wall-clock cost relative to Flow Matching is small [§4.1].
Unlike Consistency Models, which fix r = 0 and condition on a single time variable, MeanFlow conditions on two times (t, r) and inherits the consistency property from the definition of u rather than from a neural-network behavioural constraint [§4.1 “Relation to Prior Work”].

Method

MeanFlow parameterises a network u_θ(z_t, t, r) to predict the average velocity over [r, t]. Training samples (t, r) from a logit-normal distribution (with 25% probability setting r = t so the loss reduces to Flow Matching on that subset), constructs z_t = (1-t)x + t·ε along the standard rectified-flow path, and minimises an L2-style adaptive loss against the MeanFlow Identity’s RHS. The target uses the conditional instantaneous velocity v = ε - x (following Flow Matching) and a JVP through u_θ itself to compute d/dt u, with stop-gradient on the target side. The result is mechanically nearly identical to Flow Matching code — same network, same noising — with the loss modified by (t - r) · jvp(u_θ, ...). For CFG, v is replaced with the guided velocity and the same identity is reused, so the network learns the guided average velocity directly and inference is a single forward pass z_r = z_t - (t-r) · u_θ(z_t, t, r). Backbone is DiT/SiT B/M/L/XL with patch size 2 on SD-VAE latents.

Results

ImageNet 256² 1-NFE class-conditional: 6.17 (B/2, 131M) → 5.01 (M/2, 308M) → 3.84 (L/2, 459M) → 3.43 (XL/2, 676M) FID-50K [Table 2].
2-NFE: XL/2 reaches 2.20 FID, on par with 250×2-NFE DiT-XL/2 (2.27) and SiT-XL/2 (2.06) [Table 2].
Comparison to prior 1-NFE from-scratch: 50%+ relative reduction vs IMM (7.77, one-step); ~70% vs Shortcut (10.60, 1-NFE) [Fig. 1, Table 2].
CFG guidance-scale sweep: best 1-NFE FID 15.53 at ω=3.0 (vs no-CFG 61.06 at the ViT-B/4 ablation scale); 80-epoch B/4 ablation shows CFG nearly 4× improvement on FID without doubling NFE [Table 1].
Ablations confirm: 25% r=t ratio optimal, logit-normal time sampler beats uniform, correct JVP tangents are essential (degrading them blows FID from 61 to 268–329), adaptive weighting with p = 1.0 best, (t, t-r) positional embedding marginally best [Table 1, Fig. 4].
CIFAR-10 unconditional pixel-space 1-NFE: competitive with EDM-style preconditioned consistency baselines on a 55M U-Net, no preconditioner [Table 3].

Why it’s interesting

MeanFlow is the foundational reference the wiki’s Diffusion Distillation thread keeps pointing back to — it is the from-scratch frontier that distillation methods are compared against. SoFlow: Solution Flow Models for One-Step Generative Modeling (SoFlow) is explicitly framed as a JVP-free, more PyTorch-friendly alternative to MeanFlow and beats it head-to-head at every DiT size (XL/2: 2.96 vs 3.43 FID at 1-NFE), but only by changing the parameterisation (solution function f(z_t, t, s) instead of average velocity u) and the consistency objective (Taylor expansion instead of MeanFlow Identity). Transition Matching Distillation for Fast Video Generation (TMD) reports that initialising its flow head with MeanFlow-style training outperforms vanilla Flow-Matching init, so the average-velocity formulation has migrated from one-step image to few-step video flow heads. AnyFlow: Any-Step Video Diffusion Model with On-Policy Flow Map Distillation (AnyFlow) generalises the same (z_t, t) → z_r flow-map framing to any-step distillation of a video teacher — i.e. MeanFlow’s bi-time map but distilled rather than trained from scratch. The pedagogical companion is Learning the integral of a diffusion model (flow maps), which argues flow maps (of which average velocity is one parameterisation) are the right primitive for the next generation of diffusion samplers. Sam asked specifically about this paper — the short read is: it is the baseline for one-step from-scratch, has been overtaken by SoFlow on FID but is conceptually cleaner (a single identity from the definition of average velocity), and its JVP-based training recipe remains a useful component in mixed pipelines.