Improved Mean Flow (iMF) Modeling

• Improved Mean Flow (iMF) is a generative modeling framework that redefines regression targets to enhance training stability and reduce variance.
• It integrates explicit, tokenized guidance conditioning, allowing flexible adjustment of synthesis parameters during inference.
• Empirical results show that iMF achieves state-of-the-art 1-NFE performance in image, audio, and robot control domains compared to previous MF models.

Improved Mean Flow (iMF) is a generative modeling framework that addresses training stability, expressivity, and inference efficiency challenges inherent to one-step (“fastforward”) generative models, particularly those based on the MeanFlow (MF) paradigm. iMF reparameterizes the regression objective, disentangles guidance conditioning, and introduces architectural and curriculum refinements, thus achieving state-of-the-art performance in one-step generation across image, audio, and robot control domains. Central to all iMF variants is the use of average, rather than instantaneous, velocity fields—augmented with Jacobian-based corrections and explicit conditioning—which enables high-fidelity synthesis with a single function evaluation (1-NFE), approaching or surpassing the performance of multi-step diffusion and flow-based models (Geng et al., 1 Dec 2025).

1. MeanFlow Framework: Principles and Shortcomings

The MeanFlow framework replaces the many small ODE steps of traditional flow matching with a single large step by learning the average velocity $$u(z_t; r, t) = \frac{1}{t-r} \int_r^t v(z_\tau) d\tau$$ over a large interval [r, t]. A sample is generated as $x \approx z_1 - u(z_1; 0, 1)$, collapsing trajectory integration into one function evaluation (Geng et al., 19 May 2025). However, MF has two major issues:

1. Network-dependent targets: The MF training loss

$$L_{\mathrm{MF}} = \mathbb{E}\left\|u_\theta(z_t) - \mathrm{sg}[(e - x) - (t - r)\,\mathrm{JVP}(u_\theta; e-x)]\right\|^2$$

contains a Jacobian-vector product (JVP) evaluated with the network's own prediction, resulting in non-standard regression that complicates optimization and stability.

2. Fixed guidance scale: Classifier-free guidance (CFG) is supported by fixing the guidance weight ω at training, precluding flexibility at inference time and forcing suboptimal model behavior, e.g., inability to adjust guidance scale post-training.

2. Reformulation of the Training Objective

iMF addresses the network-dependent target by reformulating the MF identity as a regression in the instantaneous velocity (“v-space”) (Geng et al., 1 Dec 2025).

• The instantaneous velocity is $v(z_t) = e-x$ (for linear path), and the MF identity provides:

$$v(z_t) = u(z_t) + (t-r) \,[\partial_z u\, v + \partial_t u] = u(z_t) + (t-r) \, \mathrm{JVP}(u; v)$$

• iMF introduces a compound target

$$V_\theta(z_t) := u_\theta(z_t, r, t) + (t-r)\,\mathrm{JVP}(u_\theta; v_\theta)$$

where $v_\theta(z_t)$ is obtained either as $u_\theta(z_t, t, t)$ or via an auxiliary head.

• The new loss is standard regression, network-independent, with much lower variance:

$$L_{\mathrm{iMF}} = \mathbb{E}_{x,e,r,t} \big\|V_\theta(z_t) - (e-x)\big\|^2$$

This formulation brings training dynamics in line with classical supervised objectives, increasing robustness and convergence rate.

3. Explicit and Flexible Guidance Conditioning

Whereas MF locks CFG scale ω during training, iMF instead regards guidance hyperparameters as explicit conditioning variables (Geng et al., 1 Dec 2025). The guidance configuration inputs (including ω, [t_min, t_max], etc.) are incorporated as special tokens, jointly with class labels and time indices, into the transformer backbone. During training, ω and interval boundaries are sampled per-example; at inference, the same model can flexibly accept arbitrary guidance values.

• The guidance target is generalized as:

$v_g = (e-x) + (1 - 1/\omega)[v_c - v_u]$

where $x \approx z_1 - u(z_1; 0, 1)$0 are the conditional and unconditional instantaneous velocities.

This design allows exploitation of optimal guidance and interval settings post-training, enabling state-dependent guidance optimization and improving generative flexibility.

4. Model Architecture and Optimization

iMF employs a transformer backbone in the VAE-latent space (e.g., 32×32×4 for ImageNet 256²), adopting the following principal architecture and training design (Geng et al., 1 Dec 2025):

• Transformer depth 48, hidden size 1024, 16 attention heads—yielding 610M parameters
• Patch size 2×2, SwiGLU activations, RMSNorm, rotary positional embeddings (RoPE)
• In-context conditioning: class tokens (8), time tokens (4), guidance ω tokens (4), interval tokens (4)
• Auxiliary v-head attached to the last 8 layers, trained with a small flow-matching loss for boundary consistency
• Training from scratch (no distillation), batch size 1024, Adam optimizer (β₁=0.9, β₂=0.95), constant LR 1e−4, EMA 0.9999, 800 epochs.

This configuration, without reliance on AdaLN-zero or pretraining, achieves state-of-the-art 1-NFE performance while reducing the model footprint compared to prior MF systems.

5. Quantitative Results and Ablations

iMF sets new benchmarks in fastforward generative modeling. Notable results on ImageNet 256×256 latent space include (Geng et al., 1 Dec 2025):

Model	Params	1-NFE FID
MF-XL/2 (orig)	676M	3.43
α-Flow-XL/2	676M	2.58
iMF-XL/2	610M	1.72

• iMF-XL/2 halves the 1-NFE FID compared to MF-XL/2, and with 2-NFE sampling achieves a FID of 1.54, further narrowing the gap with multi-step models.
• Ablation studies highlight cumulative gains: introducing v-reparameterization with an auxiliary head, CFG ω-conditioning, full Ω-conditioning, and in-context conditioning yield FID improvements from 6.17 (MF-B/2 baseline) down to 3.39.
• In large-scale conditional sampling, iMF nearly matches StyleGAN-XL (FID 2.30) despite using only a single evaluation.

6. Theoretical and Algorithmic Advancements

iMF represents a broader shift in one-step generative modeling:

• The loss recasting removes network-dependent supervision, ensuring low-variance targets and reducing pathological gradient behaviors.
• Explicit CFG conditioning supports optimal scaling during synthesis and higher flexibility across tasks and datasets.
• In-context multi-token conditioning streamlines the model, reduces parameter count, and aggregates all conditional information equivalently to classic prompt-based approaches.

These features systematically address the “fastforward” pathologies outlined in original MF models—specifically, instability induced by self-dependency and guidance rigidity.

7. Significance and Implications

The improved MeanFlow method advances one-step (1-NFE) generative modeling to near parity with multi-step diffusion/flow models across FID, sample quality, and conditional control, facilitated by (i) standard regression targets via the MF identity, (ii) explicit, tokenized CFG conditioning, and (iii) compact in-context architectures (Geng et al., 1 Dec 2025). The framework allows for greater flexibility in conditional synthesis and post-training adaptation of generative parameters. iMF thus establishes a new state-of-the-art in highly efficient, scalable, and robust one-step generative modeling, and provides a blueprint for future architectural and algorithmic developments in fastforward generative paradigms.