Mean Flow

Diffusion models (Song et al., 2021) and flow matching (Lipman et al., 2023) have become the dominant paradigm for image and video generation. Although iterative refinement eases the optimization challenge, it results in slow sampling. A common acceleration approach is progressive distillation (Salimans & Ho, 2022) (Luo et al., 2023) (Yin et al., 2024), but it requires a pretrained teacher and offers limited flexibility.

Consisteny Models

Consistency models (Song et al., 2023) (Song & Dhariwal, 2024) address this by learning a direct mapping from noisy samples to clean data, enforcing self-consistency across noise levels. This enables training from scratch and supports one-step generation without iterative refinement or teacher distillation.

However, consistency models remain heuristic, constraining model outputs rather than a ground-truth field, which often leads to training instability and requires curriculum learning over time discretization.

Mean Flows

MeanFlow (Geng et al., 2025) avoids both distillation and ad-hoc curriculum learning by modeling the average velocity field directly, providing a principled and stable one-step training framework from scratch.

Denote $\mathrm{x \sim p_{\text{data}}}$ and $\mathrm{\epsilon \sim p_{\text{prior}}}$. Consider the linear interpolation $\mathrm{z_t = (1 - t)x + t\epsilon}$. The probability flow ODE yields $\mathrm{\nu_t \equiv z_t’ = \epsilon - x}$ (link), where \(\mathrm{\nu_t \equiv \nu_t(z_t \mid x)}\) is the conditional velocity. The authors introduce the elegant concept of the average velocity:

The above equality reduces to flow matching when $\mathrm{r=t}$ and the visualization is shown below.

Taking the derivative w.r.t. $\mathrm{t}$ on both sides of \eqref{def_mf}, we have

The derivatives $\mathrm{[\partial_z u, \partial_r u, \partial_t u]}$ on the LHS can be computed via a Jacobian–vector product (JVP) along the tangent vector $\mathrm{[\nu, 0, 1]}$. Let $\mathrm{u_{\theta}}$ parameterize the mean velocity $\mathrm{u}$. The training loss becomes:

where $\mathrm{sg\big[\cdot\big]}$ denotes the stop-gradient operator.

After optimizing $\mathrm{u_{\theta}}$, using the probability flow ODE $\mathrm{\nu_t \equiv z_t’}$, sampling is simply:

Core Code Snippet

# -----------------------------------------------------------------
# Instantaneous velocity
t, r = self.sample_tr(bz)

e   = jax.random.normal(self.make_rng('gen'), x.shape, dtype=self.dtype)
z_t = (1 - t) * x + t * e
v   = e - x

# Guided velocity
v_g = self.guidance_fn(v, z_t, t, labels, train=False) if self.guidance_eq else v

# Cond dropout (dropout class labels)
y_inp, v_g = self.cond_drop(v, v_g, labels)

# -----------------------------------------------------------------
# Compute u_tr (average velocity) and du_dt using jvp
def u_fn(z_t, t, r):
    return self.u_fn(z_t, t, t - r, y=y_inp, train=train)

dt_dt = jnp.ones_like(t)
dr_dt = jnp.zeros_like(t)
u, du_dt = jax.jvp(u_fn, (z_t, t, r), (v_g, dt_dt, dr_dt))

# -----------------------------------------------------------------
# Compute loss
u_tgt = v_g - jnp.clip(t - r, a_min=0.0, a_max=1.0) * du_dt
u_tgt = jax.lax.stop_gradient(u_tgt)

loss = (u - u_tgt) ** 2
loss = jnp.sum(loss, axis=(1, 2, 3)) # sum over pixels

# Adaptive weighting
adp_wt = (loss + self.norm_eps) ** self.norm_p
loss = loss / jax.lax.stop_gradient(adp_wt)

# -----------------------------------------------------------------
loss = loss.mean()  # mean over batch