---
title: 'Modular MeanFlow: Unified One-Step Modeling'
url: https://www.emergentmind.com/topics/modular-meanflow-mmf
type: topic
---

# Modular MeanFlow: Unified One-Step Modeling

Modular MeanFlow (MMF) is a unifying framework for stable and scalable one-step generative modeling, developed to efficiently generate high-quality data samples via direct mapping in a single function evaluation. MMF generalizes and interpolates between flow-matching and consistency-based models by introducing a principled family of regression losses built upon time-averaged velocity fields. Central to its design are a differential identity linking instantaneous and averaged velocities, a tunable gradient modulation mechanism, and a curriculum-style warmup schedule for training stability and expressiveness. Empirically, MMF achieves state-of-the-art performance across image synthesis, low-data, out-of-distribution (OOD), and trajectory modeling tasks, while circumventing the computational burden of higher-order derivatives [2508.17426].

## 1. Theoretical Framework

MMF builds on the continuous-time generative model defined by the ordinary differential equation (ODE):
$$
\frac{dx_t}{dt} = v(x_t, t), \qquad x_1 \sim p_\text{prior}, \;\; x_0 \sim p_\text{data}
$$
where $v(x_t, t)$ denotes the instantaneous velocity field parameterizing the mapping from $p_\text{prior}$ (usually a tractable distribution) to $p_\text{data}$. MMF introduces the time-averaged velocity field over the interval $[r, t]$:
$$
u(x_t, r, t) := \frac{1}{t - r} \int_{r}^{t} v(x_\tau, \tau)\,d\tau
$$
With Lipschitz assumptions on $v$, the averaged velocity recovers the instantaneous field as $t \to r$:
$$
\lim_{t \to r} u(x_t, r, t) = v(x_r, r)
$$
A key identity underpins MMF:
$$
v(x_t, t) = u(x_t, r, t) + (t - r) \frac{d}{dt} u(x_t, r, t)
$$
where $\frac{d}{dt}u = \partial_t u + (\nabla_x u)\,v(x_t, t)$. This relation enables the regression of averaged velocities and their time derivatives to approximate the model's functional path, decoupling expressiveness from the risk of instability intrinsic to higher-order supervision.

## 2. Modular Loss Construction and Gradient Modulation

MMF defines a spectrum of regression losses parametrized both by velocity averaging interval and a scalar $\lambda \in [0,1]$ controlling gradient flow. The “full” regression loss is given by:
$$
\mathcal{L}_\text{full} = \mathbb{E}_{x_0, x_1, r < t} \left\| u_\theta(x_t, r, t) + (t-r)\Big(\partial_t u_\theta + \nabla_x u_\theta \cdot u_\theta\Big) - \frac{x_1 - x_0}{t - r} \right\|^2
$$
To trade off training stability and functional expressiveness, a partial stop-gradient operator is introduced:
$$
\mathrm{SG}_\lambda[z] := \lambda z + (1-\lambda)\,\mathrm{stopgrad}(z)
$$
The MMF loss then generalizes as:
$$
\mathcal{L}_\lambda = \mathbb{E}_{x_0, x_1, r < t} \left\| u_\theta(x_t, r, t) + (t-r)\,\mathrm{SG}_\lambda\left[\partial_t u_\theta + \nabla_x u_\theta \cdot \frac{x_1 - x_0}{t-r}\right] - \frac{x_1 - x_0}{t - r} \right\|^2
$$
Here, $\lambda=1$ fully propagates gradients (maximum expressiveness but possible instability), $\lambda=0$ detaches Jacobian-vector products (maximum stability), and intermediate $\lambda$ governs a stability-expressiveness continuum. Explicitly blocking gradient flow through higher-order terms prevents gradient explosions and training oscillations.

## 3. Curriculum Warmup and Training Protocol

MMF adopts a curriculum for $\lambda$:
$$
\lambda(t_\mathrm{train}) = \min\left(1, \frac{t_\mathrm{train}}{T_\mathrm{warmup}}\right), \qquad T_\mathrm{warmup} \approx 10\%\;\text{of total steps}
$$
In early training ($\lambda \approx 0$), MMF behaves as a consistency or flow-matching model, yielding high stability by restricting second-order signal propagation. As $\lambda \to 1$, expressive gradients are introduced, allowing the model to capture richer curvature and achieve lower asymptotic loss. Empirically, this curriculum schedule yields both rapid convergence and low variance in training.

A typical MMF training protocol:

```python
Given: network u_θ, total steps N, warmup T_warmup
for step = 1 to N do
    sample x₀ ∼ p_data,  x₁ ∼ p_prior
    sample r∈[0,1), t∈(r,1]
    α ← (t - r)/(1 - r)
    xₜ ← (1-α) x₀ + α x₁
    λ ← min(1, step / T_warmup)
    # compute loss ℒ_λ using SG_λ on the JVP term
    L ← ‖ u_θ(xₜ,r,t)
           + (t-r)·SG_λ[∂ₜu + ∇ₓu·((x₁-x₀)/(t-r))]
           - (x₁ - x₀)/(t-r) ‖²
    θ ← θ - AdamStep(∇_θ L)
end for
```

Sampling proceeds in one step: $\hat{x}_0 = x_1 - u_\theta(x_1, r=0, t=1)$.

## 4. Connections to Prior Methods

MMF subsumes prior consistency and flow-matching models as special cases:

- **Consistency Models**: Fixing $(r, t) \equiv (0, 1)$ and $\lambda=0$ recovers the fixed-time consistency loss $\|u_\theta(x_t, 0, 1) - (x_1 - x_0)\|^2$.
- **Flow Matching**: The instantaneous limit $t \to r$, $u_\theta \to v_\theta$, and $(t-r)^{-1}(x_1 - x_0) \to v(x_r, r)$ recovers the flow-matching loss $\|v_\theta(x_r, r) - v_\mathrm{true}(x_r, r)\|^2$.
- **Gradient Efficiency**: Applying stop-gradient to the Jacobian-vector term ensures that backward computation never traverses $\nabla^2 u_\theta$, eliminating $\mathcal{O}(d^2)$ cost and Hessian-vector product overhead.

This unification allows MMF to inherit the interpretability and theoretical properties of both frameworks, while providing a tunable control for interpolation between them.

## 5. Empirical Results and Model Analysis

MMF's empirical evaluation focuses on image synthesis, robustness, and path modeling:

| Model                 | FID (↓) | 1-step MSE (↓) | LPIPS (↓) | Inference Time (s) (↓) |
|-----------------------|---------|---------------|-----------|------------------------|
| MeanFlow (full)       | 3.91    | 0.087         | 0.132     | 0.031                  |
| MeanFlow (stop-grad)  | 4.27    | 0.095         | 0.156     | 0.024                  |
| MMF (λ=0)             | 4.19    | 0.093         | 0.148     | 0.023                  |
| MMF (λ=0.5)           | 3.78    | 0.084         | 0.120     | 0.026                  |
| MMF (λ=1)             | 3.62    | 0.080         | 0.109     | 0.034                  |
| MMF (curriculum)      | 3.41    | 0.076         | 0.097     | 0.025                  |

On CIFAR-10 and ImageNet-64, curriculum MMF achieves the lowest FID, lowest 1-step MSE, and the highest diversity (LPIPS), matching or exceeding the efficiency of prior mean flow and consistency baselines. Few-shot and OOD experiments demonstrate that curriculum MMF retains low FID even with as little as 1% of CIFAR-10 data and achieves 10–20% lower FID in OOD settings (SVHN, STL-10, CIFAR-C) compared to baselines. In ODE-fitting and 2D control tasks, curriculum MMF yields smooth, accurate paths, outperforming noisy full-gradient and oversmoothed stop-grad alternatives.

Path deviation is formalized as:
$$
\mathcal{D}_\mathrm{path} = \mathbb{E}\big\| (s - r) u(x_s, r, s) + (t - s) u(x_t, s, t) - (t - r) u(x_t, r, t) \big\|
$$
Curriculum MMF achieves the lowest $\mathcal{D}_\mathrm{path}$, supporting latent interpolation smoothness.

## 6. Ablations and Practicalities

Extensive ablations reveal that:

- **Varying $\lambda$**: $\lambda=0$ yields maximum stability but higher FID (underfitting curvature). $\lambda=1$ is most expressive but unstable (loss oscillations). $\lambda=0.5$ provides some smoothing but with late-stage variance. Curriculum $\lambda(t)$ combines low early variance and best final performance.
- **Curriculum Horizon**: Short warmup (small $T_\mathrm{warmup}$) induces early instability; long warmup is too conservative with slower convergence. Optimal $T_\mathrm{warmup} \approx 10$–$15\%$ of total steps.
- **Compute**: Forward-mode autodiff for the Jacobian-vector-product yields ~15% overhead; with stopgrad, no backward is needed through this term.

The standard MMF implementation utilizes a UNet architecture with sinusoidal time embeddings, Adam optimizer (learning rate $1 \times 10^{-4}$, batch size 128, cosine decay), and a curriculum warmup over 100k steps.

## 7. Significance, Limitations, and Outlook

MMF provides a theoretically grounded, computationally efficient, and practically robust approach for one-step generative modeling. By enabling a tunable spectrum between expressiveness and stability—mediated by gradient modulation and curriculum scheduling—it addresses the instability and inefficiency intrinsic to prior higher-order methods. MMF’s empirical results demonstrate high generalization under low-data and out-of-distribution regimes and applicability beyond image synthesis to trajectory modeling. A plausible implication is that MMF may be extensible to other domains requiring stable, one-shot sampling of complex data distributions via learnable ODE flows [2508.17426].

Source: https://www.emergentmind.com/topics/modular-meanflow-mmf