Source-linked AI summary
Mean Flows for One-step Generative Modeling
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter, Kaiming He
TL;DR
One-step generative models need a principled training target beyond consistency constraints and instantaneous velocity fields. MeanFlow models average velocity through its intrinsic relation to instantaneous velocity, achieving 3.43 FID with 1-NFE ImageNet generation from scratch and narrowing the gap with multi-step models.
Problem
Consistency-based one-step modeling can be unstable and require a discretization curriculum, while coarse integration of curved marginal flow trajectories can be inaccurate.
Method
MeanFlow defines average velocity as interval displacement and trains a neural network to satisfy its intrinsic relation with instantaneous velocity.
Results
3.43 FID is achieved on ImageNet 256×256 with 1-NFE generation, trained from scratch without pre-training, distillation, or curriculum learning.
Takeaways & Limitations
MeanFlow largely closes the gap between one-step diffusion/flow models and their multi-step predecessors.
Takeaways & Limitations
The formulation includes a time-dependent r assumption requiring the Leibniz rule.
Abstract
from arXiv · showhide
We propose a principled and effective framework for one-step generative modeling. We introduce the notion of average velocity to characterize flow fields, in contrast to instantaneous velocity modeled by Flow Matching methods. A well-defined identity between average and instantaneous velocities is derived and used to guide neural network training. Our method, termed the MeanFlow model, is self-contained and requires no pre-training, distillation, or curriculum learning. MeanFlow demonstrates strong empirical performance: it achieves an FID of 3.43 with a single function evaluation (1-NFE) on ImageNet 256x256 trained from scratch, significantly outperforming previous state-of-the-art one-step diffusion/flow models. Our study substantially narrows the gap between one-step diffusion/flow models and their multi-step predecessors, and we hope it will motivate future research to revisit the foundations of these powerful models.
1 Introduction
MeanFlow addresses limitations of one-step generative modeling by defining and learning an average-velocity field with an intrinsic relation to instantaneous velocity. It achieves strong 1-NFE ImageNet performance from scratch without pre-training, distillation, or curriculum learning.
- Motivation: Flow Matching and diffusion models ordinarily use iterative sampling, while one-step consistency methods can require unstable training and a discretization curriculum.The underlying ground-truth field guiding consistency learning was previously unspecified.
- MeanFlow: MeanFlow introduces average velocity as a ground-truth field, defined by displacement over a time interval and related intrinsically to instantaneous velocity.Displacement is the time integral of instantaneous velocity.
- MeanFlow: The method trains a neural network to model average velocity using a loss enforcing the intrinsic average–instantaneous velocity relation, without an extra consistency heuristic.The target field can also incorporate classifier-free guidance without additional sampling cost.
- Results: 3.43 FID is achieved on ImageNet 256×256 with 1-NFE generation, outperforming previous methods by a relative margin of 50% to 70%.The model is trained entirely from scratch without pre-training, distillation, or curriculum learning.
2 Related Work
Related work develops diffusion- and flow-based approaches for reducing sampling steps, including distillation, consistency constraints, and methods modeling quantities across pairs of time steps. MeanFlow differs by targeting average velocity.
- Diffusion and Flow Matching: Diffusion models reverse progressively added noise, while Flow Matching models velocity fields defining paths between distributions.Flow Matching extends the diffusion-related framework through flow-path velocity modeling.
- Few-step Diffusion/Flow Models: Diffusion and flow research has pursued few-step generation through distillation, score distillation, and standalone consistency models.Consistency Models avoid distillation but impose consistency constraints across network outputs at different times.
- Time-dependent Quantities: Flow Map methods learn integrals of flow between time steps, whereas MeanFlow is based on average velocity corresponding to displacement over a time interval.Shortcut Models and Inductive Moment Matching impose additional cross-time self-consistency relationships.
3 Background: Flow Matching
Flow Matching constructs paths between prior and data distributions and trains on instantaneous velocity fields. Marginalizing conditional velocities produces the ground-truth field, whose curved trajectories can make coarse numerical integration inaccurate.
- Flow Path Construction: A flow path can be constructed as z_t = a_tx + b_tϵ, with conditional velocity defined as the time derivative z′_t.For the schedule a_t = 1 − t and b_t = t, the conditional velocity is v_t = ϵ − x.
- Marginal Velocity: A given z_t may arise from multiple data–prior pairs, so Flow Matching models the expectation over conditional velocities as the marginal velocity.The marginal velocity field is used as the underlying ground-truth field for network training.
- Training Objective: Flow Matching objectives can be equivalent when minimizing conditional-flow and marginal-flow losses.The supplied formulation states this equivalence as minimizing L_CFМ and L_FМ.
- Numerical Sampling: Euler integration advances each discrete step using z_ti+1 = z_ti + (t_i+1 − t_i)v(z_ti,t_i), while higher-order solvers are also possible.These solvers numerically follow the learned velocity field over discrete time steps.
- Trajectory Geometry: Even straight conditional flows typically yield curved marginal trajectories, and coarse discretization over them causes inaccurate numerical ODE solutions.The curvature arises from the ground-truth marginal field, not only neural-network approximation.
4 MeanFlow Models
MeanFlow models average velocity rather than instantaneous velocity, deriving an intrinsic identity between them to train a neural network for one-step generation. The formulation supports direct one-step sampling and classifier-free guidance while retaining single-NFE behavior.
- Mean Flows: Average velocity u is displacement over the interval [r, t], whereas Flow Matching models instantaneous velocity v.The displacement is (t − r)u(z, r, t), and u generally need not align with the instantaneous tangent direction.
- Mean Flows: A network approximating average velocity can represent the flow path without explicitly integrating at inference time, enabling one-step sampling.For one-step sampling, z0 = z1 − u(z1, 0, 1), with z1 drawn from the prior.
- Mean Flows: MeanFlow trains uθ with a regression target formed from instantaneous velocity and derivatives of the parameterized average-velocity field.Stop-gradient on the target avoids double backpropagation through the Jacobian-vector product, while JVP computation has comparable cost to backpropagation and under 20% overhead in the reported JAX implementation.
- Mean Flows: Restricting the time pair to t = r removes the average-velocity correction and recovers standard Flow Matching.The correction is −(t−r)(v_t∂_zuθ + ∂_tuθ), so the distinction arises from using two time points.
- Mean Flows: The MeanFlow Identity gives an intrinsic relation between average and instantaneous velocities that does not depend on neural-network parameterization.The identity is derived by differentiating the average-velocity definition and provides the target form used for training.
- Mean Flows with Guidance: MeanFlow incorporates classifier-free guidance into the learned target, avoiding the extra sampling evaluation required by naïve guidance.The guided average velocity is used directly for one-step sampling, preserving single-NFE behavior.
- Design Decisions: The design space includes alternative loss metrics, time-step sampling distributions, and conditioning schemes for the pair (r, t).The paper considers squared L2 and generalized loss forms, uniform or logit-normal time sampling, and positional embeddings for time conditioning.
5 Experiments
Experiments evaluate MeanFlow on ImageNet and CIFAR-10, including ablations of training choices, comparisons at one and two function evaluations, and scalability across model sizes. MeanFlow achieves strong one-step results while being trained from scratch, and its ablations identify effective configurations for 1-NFE generation.
- Experiment Setting: ImageNet experiments evaluate FID-50K on 256×256 generation, primarily using latent inputs from a pre-trained VAE tokenizer and models trained from scratch.The tokenizer produces a 32×32×4 latent space, and 1-NFE generation is the default evaluation setting.
- From Flow Matching to Mean Flows: A 0% r≠t sampling ratio fails for 1-NFE generation, whereas nonzero ratios enable meaningful MeanFlow results; the best FID occurs at 25%.A 100% ratio also produces a valid result, indicating that MeanFlow requires propagation beyond the r=t Flow Matching case.
- JVP Computation: Correct JVP computation is necessary for meaningful results, making the JVP the core relation connecting all (r, t) coordinates.The ablation intentionally corrupts JVP computation and finds that meaningful results occur only with the correct computation.
- Conditioning on (r, t): Embedding time and interval, (t, t−r), achieves the best result, while directly embedding (r, t) performs almost as well and interval-only embedding remains reasonable.All studied (r, t) embedding variants yield meaningful 1-NFE results.
- Time Samplers and Loss Metrics: The logit-normal sampler performs best for sampling (r, t), and adaptive loss weighting with p = 1 outperforms the tested alternatives.The standard squared L2 setting remains meaningful but underperforms; CFG also substantially improves 1-NFE generation while preserving 1-NFE sampling.
- Scalability: MeanFlow shows promising 1-NFE scalability across larger model sizes and training durations, while its ImageNet models are trained from scratch.The comparisons use CFG where applicable, and the reported scalability results preserve 1-NFE sampling behavior.
- Comparisons with Prior Work: MeanFlow achieves 3.43 FID at 1-NFE on ImageNet 256×256 and 2.20 FID at 2-NFE, with the latter on par with many-step DiT and SiT baselines.The 1-NFE result improves over prior methods by more than 50% relative to IMM and nearly 70% relative to Shortcut when comparing 1-NFE generation.
- Comparisons with Prior Work: The method is self-contained, requiring no pre-training, distillation, or curriculum learning, and is competitive with prior approaches on CIFAR-10.CIFAR-10 evaluation uses unconditional 1-NFE FID-50K with a roughly 55M-parameter U-net applied directly in pixel space.
6 Conclusion
The paper presents MeanFlow as a framework for one-step generation and connects its formulation to broader multi-scale simulation problems. It aims to bridge generative modeling, simulation, and dynamical systems research.
- MeanFlow is presented as a principled and effective framework for one-step generation.
- The formulation describes underlying quantities at coarsened levels of granularity, paralleling a common theme in physics applications.
- The authors hope the framework will bridge research in generative modeling, simulation, and dynamical systems.
A Implementation
The implementation uses DiT-based latent modeling for ImageNet and a roughly 55M-parameter pixel-space U-net for CIFAR-10. It also details an extended classifier-free guidance formulation with mixing scale κ.
- ImageNet 256×256: ImageNet models use a standard VAE tokenizer, 32×32×4 latent inputs, and DiT backbones with two-time-variable conditioning.
- CIFAR-10: CIFAR-10 models use a roughly 55M-parameter U-net directly in pixel space, condition on two time variables, and omit the EDM preconditioner.
- Improved CFG: The extended CFG target mixes class-conditional and class-unconditional predictions through κ while retaining the same loss function.
- Improved CFG: Mixing by κ can further improve generation quality when the effective guidance scale is fixed at 2.0.
B.2 Loss Metrics
The loss metrics generalize squared L2 regression through a powered loss and implement it as an adaptively weighted squared L2 objective. The weighting uses a small constant and stop-gradient operation.
- The powered L2 loss generalizes squared L2 using a user-specified exponent γ.
- The powered loss is equivalent to weighting squared L2 by a loss-adaptive factor determined by the regression error and γ.
- The practical weighting uses p = 1 − γ and a small positive constant c to avoid division by zero.
- The adaptively weighted objective applies stop-gradient to the weight, and p = 0.5 yields a loss similar to Pseudo-Huber.
B.3 On the Sufficiency of the MeanFlow Identity
The appendix proves that the MeanFlow Identity is not only necessary but also sufficient for the defining relation. The proof uses the displacement field’s boundary condition to cancel the integration constant.
- The MeanFlow Identity is shown to be a sufficient condition for the defining average-velocity relation, completing the equivalence.
- The proof addresses the usual integration-constant ambiguity that arises when equality of derivatives is used to establish equality of integrals.
- The boundary condition S|t=r = 0, together with the corresponding zero integral at t = r, cancels the constants.
- The sufficiency follows from modeling average velocity u rather than directly parameterizing displacement S, which automatically satisfies the boundary condition.
B.4 Analysis on Jacobian-Vector Product (JVP) Computation
MeanFlow’s JVP computation adds limited training overhead because the JVP product is stop-gradient during parameter backpropagation and only requires input-variable backpropagation.
- The JVP product is stop-gradient, so it is treated as constant during SGD backpropagation with respect to θ.
- JVP’s backward pass only propagates to input variables, making it less costly than standard parameter backpropagation.
- 0.052 sec/iter with MeanFlow versus 0.045 sec/iter with Flow Matching corresponds to a 16% wall-clock training overhead.
C Qualitative Results
Figure 5 presents curated 1-NFE generation examples from MeanFlow on ImageNet at 256×256 resolution.
- Figure 5 shows curated generation examples produced by the ImageNet 256×256 1-NFE model.
- The qualitative examples are drawn from ImageNet 256×256.
- The displayed model generates samples using 1-NFE inference.