Generative Diffusion Models Chat (30): From Instantaneous Velocity to Average Velocity
As is well known, slow generation speed has always been a pain point of diffusion models, and to solve this problem, everyone has been showing off their unique tricks, proposing all kinds of solutions. However, for a long time no single piece of work has managed to stand out and become the de facto standard. What kind of work would meet that bar? In my view, it should satisfy at least a few conditions:
1. Clear mathematical principles that reveal the essence of fast generation;
2. A single objective trainable from scratch, without needing extra techniques such as adversarial training or distillation;
3. Single-step generation close to SOTA, with the ability to improve results by increasing the number of steps.
Based on my reading experience, almost no work satisfies all three criteria at once. However, just a few days ago, a paper appeared on arXiv titled Mean Flows for One-step Generative Modeling (abbreviated "MeanFlow"), which looks very promising. Taking this as an opportunity, we will discuss the related ideas and progress in this area. more
Existing Approaches
There has already been a great deal of work on accelerating diffusion model generation, and this blog has briefly covered some of it before. Broadly speaking, acceleration approaches can be divided into three categories.
First, converting the diffusion model into an SDE/ODE and then studying more efficient solvers, with DPM-Solver and its series of follow-up improvements being representative. However, this approach typically can only bring the NFE (Number of Function Evaluations) down to around 10; going lower noticeably degrades generation quality. This is because the convergence rate of solvers is usually proportional to some power of the step size, and when the NFE is very small the step size cannot be small enough, so convergence isn't fast enough to be usable.
Second, distilling a trained diffusion model into a generator with fewer steps. This has spawned a great many works and schemes; we previously covered one such scheme called SiD. Distillation is a fairly conventional and general approach, but it shares a common drawback: it requires additional training cost and is not a from-scratch training scheme. Some works, in order to distill down to a single-step generator, add multiple optimization strategies such as adversarial training on top, making the overall scheme often overly complex.
Third, approaches based on the Consistency Model (CM), including the CM we introduced simply in Generative Diffusion Models Chat (28): Understanding the Consistency Model Step by Step, its continuous-time version sCM, and CTM, among others. CM represents its own distinct line of thinking—it can be trained from scratch to get a model with very small NFE, and can also be used for distillation. But CM's objective depends on EMA or the stop_gradient operation, meaning it is coupled to the optimizer's dynamics, which always gives one a somewhat ambiguous, hard-to-pin-down feeling.
Instantaneous Velocity
So far, the diffusion models with the smallest generation NFE have basically all been ODEs, because deterministic models tend to be easier to analyze and solve. This post likewise focuses only on ODE-based diffusion, using the framework of ReFlow introduced in Generative Diffusion Models Chat (17): The General Procedure for Constructing ODEs (Part 2), which is essentially equivalent to Flow Matching but more intuitive.
ODE-based diffusion aims to learn an ODE
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)\label{eq:ode}\end{equation}
in order to construct a transform $\boldsymbol{x}_1\to \boldsymbol{x}_0$. Specifically, let $\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)$ be some easily sampled random noise, and let $\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)$ be a real sample from the target distribution. We hope that, through the above ODE, we can achieve a transformation from random noise to a target sample; that is, we randomly sample $\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)$ as the initial value, and solving the above ODE to get $\boldsymbol{x}_0$ gives us a sample of $p_0(\boldsymbol{x}_0)$.
If we regard $t$ as time and $\boldsymbol{x}_t$ as displacement, then $d\boldsymbol{x}_t/dt$ is instantaneous velocity, so ODE-based diffusion is essentially modeling instantaneous velocity. So how do we train $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$? ReFlow proposes a very intuitive method: first construct an arbitrary interpolation between $\boldsymbol{x}_0$ and $\boldsymbol{x}_1$, such as the simplest linear interpolation $\boldsymbol{x}_t=(1-t)\boldsymbol{x}_0 + t \boldsymbol{x}_1$. Then taking the derivative of $t$ with respect to time gives
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{x}_1 - \boldsymbol{x}_0\end{equation}
This is an extremely simple ODE, but it doesn't meet our requirements, because $\boldsymbol{x}_0$ is our target—it shouldn't appear inside the ODE. To address this, ReFlow proposes a very intuitive idea: use $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ to approximate $\boldsymbol{x}_1 - \boldsymbol{x}_0$:
\begin{equation}\mathbb{E}_{t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\Vert\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\Vert^2\right]\label{eq:loss-reflow}\end{equation}
This is ReFlow's objective function. It's worth pointing out that: 1) in theory ReFlow permits an arbitrary interpolation scheme between $\boldsymbol{x}_0$ and $\boldsymbol{x}_1$; 2) although ReFlow is intuitive, it is also theoretically rigorous—one can prove that its optimal solution really is the ODE we're looking for. For the relevant details, please refer to Generative Diffusion Models Chat (17): The General Procedure for Constructing ODEs (Part 2) and the original paper.
Average Velocity
However, an ODE is merely a pure mathematical form; actually solving it still requires discretization, such as the simplest Euler scheme:
\begin{equation}\boldsymbol{x}_{t - \Delta t} = \boldsymbol{x}_t - \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) \Delta t\end{equation}
The NFE from $1$ to $0$ is $1/\Delta t$, and wanting a small NFE is equivalent to wanting a large $\Delta t$. However, ReFlow's theoretical foundation is the exact ODE—that is, only when the ODE is solved exactly can it produce the target sample—which means smaller $\Delta t$ is better, running counter to our goal. Although ReFlow claims that using straight-line interpolation makes the ODE trajectory straighter, thereby allowing larger $\Delta t$, in practice the trajectory is still curved, and $\Delta t$ can hardly get close to 1, so ReFlow struggles to achieve one-step generation.
At the end of the day, the ODE is inherently a $\Delta t\to 0$-type object, and forcing it to be used for $\Delta t \to 1$ while also demanding it perform well is, in itself, "asking too much of the model." So changing the modeling target—rather than continuing to "torture" the model—is the essential path toward faster generation. To this end, let us consider integrating both sides of equation $\eqref{eq:ode}$:
\begin{equation}\boldsymbol{x}_t - \boldsymbol{x}_r = \int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau = (t-r)\times \frac{1}{t-r}\int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau\end{equation}
If we can model
\begin{equation} \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) \triangleq \frac{1}{t-r}\int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau\end{equation}
then we have $\boldsymbol{x}_0 = \boldsymbol{x}_1 - \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_1, 0, 1)$, meaning that in theory we can achieve exact one-step generation, without having to resort to approximate relations. If $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ is the instantaneous velocity at time $t$, then clearly $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$ is the average velocity over the time interval $[r,t]$. In other words, to accelerate generation, or even achieve one-step generation, our modeling target should be the average velocity, rather than the ODE's instantaneous velocity.
Identity Transformations
Of course, the shift from instantaneous velocity to average velocity is not hard to think of; the truly difficult part is how to construct a loss function for it. ReFlow only tells us how to build a loss function for instantaneous velocity—about training average velocity, we know nothing.
The natural next thought is to "turn the unknown into the known," i.e., to start from the average velocity $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$ and construct the instantaneous velocity $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ from it, then substitute this into ReFlow's objective function. This requires us to derive the identity relation connecting the two. From the definition of $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$, we get
\begin{equation} \int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau = (t-r)\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) \end{equation}
Taking the derivative of both sides with respect to $t$ gives
\begin{equation}\begin{aligned} \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) =&\, \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\frac{d}{dt}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) \\ =&\, \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\left[\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right] \end{aligned}\label{eq:id1}\end{equation}
This is the first identity relation between $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ and $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$. Where there's a first, naturally there's a second; the second identity relation comes from the definition of average velocity:
\begin{equation}\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) = \lim_{r\to t}\frac{1}{t-r}\int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau = \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)\label{eq:id2}\end{equation}
In plain terms, the average velocity over an infinitesimally small interval equals the instantaneous velocity.
First Objective
Based on $d\boldsymbol{x}_t/dt = \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ and the identity $\eqref{eq:id2}$, we can replace $d\boldsymbol{x}_t/dt$ in the identity $\eqref{eq:id1}$ with either $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ or $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)$. The former is an implicit relation, which we'll discuss later; let's first look at the latter, in which case we have:
\begin{equation}\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) = \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\left[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right]\end{equation}
Substituting this into ReFlow, we obtain the first objective function that can be used to train $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$:
\begin{equation}\mathbb{E}_{r,t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\left[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right] - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\label{eq:loss-1}\end{equation}
This is an extremely satisfying result—it meets all our expectations for the objective function of a generative model:
1. A single, explicit minimization objective;
2. No EMA, stop_gradient, or similar operations;
3. Theoretical guarantees (from ReFlow).
These properties mean that, regardless of what optimization algorithm we use, as long as we can find the minimum point of the above expression, it will be the average velocity model we want—that is, a generative model that, in theory, can achieve one-step generation. In other words, it has the training simplicity and theoretical guarantees of diffusion models, yet can generate in one step like a GAN, without having to pray that the model doesn't "give up" and collapse during training.
The JVP Operation
However, for some readers, implementing the objective function $\eqref{eq:loss-1}$ might still be a bit tricky, since it involves the "Jacobian-Vector Product" (JVP), which is relatively unfamiliar to ordinary users. Specifically, we can write the term inside the square brackets of the objective function as:
\begin{equation}\underbrace{\left[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t),0,1\right] \\[10pt]}_{\text{vector}}\cdot\underbrace{\left[\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t), \frac{\partial}{\partial r}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t), \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right] \\[10pt]}_{\text{Jacobian matrix}}\end{equation}
that is, the product of the Jacobian matrix of $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$ with a given vector $[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t),0,1]$, yielding a vector the same size as $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$. This operation is called a JVP, and it is already implemented in both Jax and Torch. For example, the reference code in Jax is:
u = lambda xt, r, t: diffusion_model(weights, [xt, r, t])
urt, durt = jax.jvp(u, (xt, r, t), (u(xt, t, t), r * 0, t * 0 + 1))
Here urt is $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$, and durt is the corresponding JVP result; usage in Torch is similar. Once you understand the JVP operation, implementing the objective function $\eqref{eq:loss-1}$ becomes basically trivial.
Second Objective
If we had to point out a drawback of the objective function $\eqref{eq:loss-1}$, in my view there's only one: its computational cost is relatively high. This is because it requires two separate forward passes, $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$ and $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)$, followed by one gradient computation for the JVP, and then—when optimizing via gradient descent—another gradient computation on top of that. So it essentially requires second-order gradients, similar to the earlier WGAN-GP.
To reduce the computational cost, we can consider adding a stop_gradient operation ($\newcommand{\sg}[1]{\color{skyblue}{\mathop{\text{sg}}\left[\color{blue}{#1}\right]}}\sg{\cdot}$) to the JVP part:
\begin{equation}\mathbb{E}_{r,t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\sg{\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)} - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\label{eq:loss-2}\end{equation}
This avoids taking a gradient through the JVP a second time (though it still requires two forward passes). Empirical results show that, compared with the first objective $\eqref{eq:loss-1}$, the above objective trains nearly twice as fast under gradient-based optimizers, and visually the results show no loss in quality.
Note that the stop_gradient here is used purely to reduce computational cost; the actual optimization direction still aims for the smallest possible loss value. This differs from CM-series models, especially sCM, whose loss functions are merely "equivalent losses" with equivalent gradients—not necessarily objectives that should be minimized as low as possible—and whose stop_gradient is usually mandatory: removing it will almost certainly cause training to collapse.
Third Objective
Earlier we mentioned that another way to handle $d\boldsymbol{x}_t/dt$ in the identity $\eqref{eq:id1}$ is to replace it with $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$, which leads to
\begin{equation}\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) = \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\left[\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right]\end{equation}
If we were to solve for $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ from this, the result would be
\begin{equation}\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t) = \left[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right]\cdot\left[\boldsymbol{I} - (t-r)\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right]^{-1}\end{equation}
which requires inverting an enormous matrix and is not practical. MeanFlow uses a compromise: since the regression target of $d\boldsymbol{x}_t/dt = \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$ is $\boldsymbol{x}_1-\boldsymbol{x}_0$, why not simply replace $d\boldsymbol{x}_t/dt$ with $\boldsymbol{x}_1-\boldsymbol{x}_0$? The objective function then becomes
\begin{equation}\mathbb{E}_{r,t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\left[(\boldsymbol{x}_1-\boldsymbol{x}_0)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)\right] - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\end{equation}
However, here $\boldsymbol{x}_1-\boldsymbol{x}_0$ is both the regression target and appears in the definition of the model $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t,t)$, which inevitably gives a sense of "label leakage." To avoid this problem, MeanFlow's approach is once again to add stop_gradient to the JVP part:
\begin{equation}\mathbb{E}_{r,t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + (t-r)\sg{(\boldsymbol{x}_1-\boldsymbol{x}_0)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)} - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\label{eq:loss-3}\end{equation}
This is the loss function ultimately used by MeanFlow, which we call the "third objective" here. Compared with the second objective $\eqref{eq:loss-2}$, it eliminates one forward pass, $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t)$, so training is somewhat faster. But the introduction of "label leakage" here and the countermeasure of stop_gradient mean that the training of the third objective is coupled to the gradient optimizer, just like CM, adding a bit of that hard-to-pin-down mystique.
The paper's experimental results show that the objective $\eqref{eq:loss-3}$ with $\sg{\cdot}$ included can train to reasonable results—so what happens if we remove it? I asked the authors about this, and they indicated that after removing $\sg{\cdot}$, training still converges and can do multi-step generation, but loses the ability to do one-step generation. This isn't hard to understand actually, because when $r=t$, regardless of whether $\sg{\cdot}$ is present or not, the objective function degenerates to ReFlow:
\begin{equation}\mathbb{E}_{t,\boldsymbol{x}_0,\boldsymbol{x}_1}\left[\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t, t) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\Vert^2\right]\label{eq:loss-reflow-2}\end{equation}
In other words, MeanFlow always has ReFlow as a "safety net" underneath, so it can never get too bad. Removing $\sg{\cdot}$ intensifies the negative effect of "label leakage," so it's better to keep it.
Let's Prove It
Can we, as with ReFlow, prove theoretically that the optimal solution of the third objective $\eqref{eq:loss-3}$ really is the average velocity model we're after? Let's try. First let's recall the two key lemmas used in proving ReFlow:
1. $\mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}[\Vert\boldsymbol{\mu} - \boldsymbol{x}\Vert^2] = \mathbb{E}[\boldsymbol{x}]$, i.e., minimizing the squared error between $\boldsymbol{\mu}$ and $\boldsymbol{x}$, has the optimal solution being the mean of $\boldsymbol{x}$;
2. The ODE form solution that transforms $\boldsymbol{x}_1$ into $\boldsymbol{x}_0$ following the distributional trajectory $\boldsymbol{x}_t=(1-t)\boldsymbol{x}_0 + t \boldsymbol{x}_1$ is $d\boldsymbol{x}_t/dt = \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}[\boldsymbol{x}_1-\boldsymbol{x}_0]$.
Among these, the proof of Lemma 1 is fairly simple—taking the gradient of $\boldsymbol{\mu}$ directly gives $\mathbb{E}[\boldsymbol{\mu} - \boldsymbol{x}] = \boldsymbol{\mu} - \mathbb{E}[\boldsymbol{x}]$, and setting it to zero suffices. The proof details of Lemma 2 need to be found in Generative Diffusion Models Chat (17): The General Procedure for Constructing ODEs (Part 2), where $\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}[\boldsymbol{x}_1-\boldsymbol{x}_0]$ is obtained by first using $\boldsymbol{x}_t=(1-t)\boldsymbol{x}_0 + t \boldsymbol{x}_1$ to eliminate $\boldsymbol{x}_1$, obtaining a function of $\boldsymbol{x}_0,\boldsymbol{x}_t$, and then taking the expectation over the distribution $p_t(\boldsymbol{x}_0|\boldsymbol{x}_t)$, resulting in a function of $t,\boldsymbol{x}_t$.
Using Lemma 1, we can prove that the theoretical optimal solution of ReFlow's objective function $\eqref{eq:loss-reflow}$ is exactly $\boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t,t) = \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}[\boldsymbol{x}_1-\boldsymbol{x}_0]$; combined with Lemma 2, this gives us that $d\boldsymbol{x}_t/dt=\boldsymbol{v}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t,t)$ is the ODE we're after. The proof of the third objective $\eqref{eq:loss-3}$ proceeds similarly: since it contains $\sg{\cdot}$, taking the gradient with respect to $\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t)$ and setting it to zero yields
\begin{equation}\begin{aligned} \boldsymbol{0} =&\, \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[(t-r)\left[(\boldsymbol{x}_1-\boldsymbol{x}_0)\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t)\right] - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right] \\ =&\, \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + (t-r)\left[\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}[\boldsymbol{x}_1-\boldsymbol{x}_0]\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t)\right] - \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}[\boldsymbol{x}_1 - \boldsymbol{x}_0] \\ =&\, \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + (t-r)\left[\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t)\right] - \frac{d\boldsymbol{x}_t}{dt} \\ =&\, \frac{d}{dt}\left[(t - r) \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) - (\boldsymbol{x}_t - \boldsymbol{x}_r)\right] \\ \end{aligned}\end{equation}
So under appropriate boundary conditions we have $\boldsymbol{x}_t - \boldsymbol{x}_r = (t - r) \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t)$, which is precisely the average velocity model we expected.
The key to this process is that the introduction of $\sg{\cdot}$ avoids taking the gradient of the JVP part, thereby simplifying the gradient expression and yielding the correct result. If we removed $\sg{\cdot}$, the right-hand side of the above expression would need an extra multiplicative term—the Jacobian matrix of the JVP part with respect to $\boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t)$—and the result would be that we could no longer isolate the term $\frac{d}{dt}\left[(t - r) \boldsymbol{u}_{\boldsymbol{\theta}^*}(\boldsymbol{x}_t, r, t) - (\boldsymbol{x}_t - \boldsymbol{x}_r)\right]$. The mathematical purpose of introducing $\sg{\cdot}$ is precisely to resolve this issue.
Of course, as I've said before, the introduction of $\sg{\cdot}$ also couples the entire model's training to the gradient optimizer, adding a touch of ambiguity. The point where the gradient equals zero can at best be called a stationary point rather than a (local) minimum, so its stability is not entirely clear either—this is, in fact, a common feature of all models coupled to $\sg{\cdot}$.
Related Work
Interestingly, two earlier posts we covered on accelerating generation, Generative Diffusion Models Chat (21): The Mean Value Theorem for Accelerated ODE Sampling and Generative Diffusion Models Chat (27): Using the Step Size as a Conditioning Input, are both centered on average velocity, and can be said to share a common lineage of ideas. Even though the authors may not necessarily have been connected to each other, their work indeed gives a sense of continuity—of one building upon the other.
In the mean value theorem post, the author had already recognized the importance of the average velocity
\begin{equation}\frac{1}{t-r}\int_r^t \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_{\tau},\tau) d\tau\end{equation}
but the approach taken there was to draw an analogy with the integral mean value theorem for one-dimensional functions, attempting to find $s\in[r,t]$ such that $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_s,s)$ equals the average velocity. This is essentially still the idea of a higher-order solver, but it's no longer training-free—it requires a small number of distillation steps, which counts as a modest breakthrough for solvers.
The Shortcut model proposed in the step-size-as-input post, meanwhile, comes very close to touching MeanFlow, because using the step size as an extra input is essentially equivalent to MeanFlow's dual time parameter $r,t$. The difference is that it trains the model directly using the properties of average velocity as an extra regularization term. In this post's notation, the property that average velocity should satisfy is
\begin{equation}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) = \frac{1}{2}\left[\boldsymbol{u}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_t, s, t\right) + \boldsymbol{u}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_s, r, s\right)\right]\end{equation}
where $s = (r+t)/2$. So Shortcut simply uses this to construct a regularization term
\begin{equation}\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, r, t) - \frac{1}{2}\sg{\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, s, t) + \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_s, r, s)}\right\Vert^2\end{equation}
which is mixed with ReFlow's objective $\eqref{eq:loss-reflow-2}$ during training; in actual training, $\boldsymbol{x}_s = \boldsymbol{x}_t - (t-s)\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, s, t)$. The introduction of $\sg{\cdot}$ is, in my view, mainly to save computational cost. The Shortcut model is actually more intuitive than MeanFlow, but since it lacks the identity-transformation and ReFlow-based rigorous theoretical support, it comes across more as a product of a transitional stage.
Consistency Models
Finally, let's discuss consistency models once more. Given that CM and sCM came before it, MeanFlow's success actually also draws on their experience, especially the operation of adding $\sg{\cdot}$ to the JVP, which is also mentioned in the original paper. Of course, one of MeanFlow's authors, Kaiming He, is himself a master of manipulating gradients (e.g., SimSiam), so the emergence of MeanFlow looks like a very natural culmination of prior work.
We analyzed the discrete CM carefully in Generative Diffusion Models Chat (28): Understanding the Consistency Model Step by Step. If we replace CM's EMA operator with stop_gradient, take the gradient, and let $\Delta t\to 0$ approach its limit, we obtain the objective function of sCM from Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models:
\begin{equation}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\cdot \frac{d}{dt}\boldsymbol{f}_{\sg{\boldsymbol{\theta}}}(\boldsymbol{x}_t, t) = \boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\cdot\sg{\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) + \frac{\partial}{\partial t}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)}\label{eq:loss-scm}\end{equation}
If we replace $\frac{d\boldsymbol{x}_t}{dt}$ with $\boldsymbol{x}_1 - \boldsymbol{x}_0$ and denote $\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) = \boldsymbol{x}_t - t\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t , 0, t)$, then its gradient is equivalent to MeanFlow's third objective $\eqref{eq:loss-3}$ when $r=0$:
\begin{equation}\begin{aligned} \nabla_{\boldsymbol{\theta}}\eqref{eq:loss-scm} =&\, \nabla_{\boldsymbol{\theta}}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\cdot \left[\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) + \frac{\partial}{\partial t}\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right] \\[10pt] =&\, -t\nabla_{\boldsymbol{\theta}}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t)\cdot \left[\frac{d\boldsymbol{x}_t}{dt} - t\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) - \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) - t\frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t)\right] \\[10pt] =&\, t\nabla_{\boldsymbol{\theta}}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t)\cdot \left[\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) + t\left[\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t)\right]- \frac{d\boldsymbol{x}_t}{dt}\right] \\[10pt] =&\, \frac{t}{2}\nabla_{\boldsymbol{\theta}}\left\Vert\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) + t\sg{\frac{d\boldsymbol{x}_t}{dt}\cdot\frac{\partial}{\partial \boldsymbol{x}_t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t) + \frac{\partial}{\partial t}\boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, 0, t)}- \frac{d\boldsymbol{x}_t}{dt}\right\Vert^2 \\[10pt] \sim &\, \left.\nabla_{\boldsymbol{\theta}}\eqref{eq:loss-3}\right|_{r=0} \end{aligned}\end{equation}
So, from this angle, sCM is a special case of MeanFlow at $r=0$. As mentioned earlier, introducing the additional time parameter $r$ allows ReFlow to serve as MeanFlow's "safety net" (when $r=t$), thus better avoiding training collapse—this is one of its advantages. Of course, one could also introduce a dual time parameter starting from sCM and obtain a result identical to the third objective, but from a personal aesthetic standpoint, the physical meaning of CM/sCM is ultimately not as intuitive as MeanFlow's interpretation via average velocity.
Furthermore, starting from the combination of average velocity and ReFlow, we can also obtain another first objective $\eqref{eq:loss-1}$ and second objective $\eqref{eq:loss-2}$, which is an extremely comfortable and elegant result for someone like me who has a stop_gradient purity fetish. In my view, from the standpoint of computational cost, it's fine to consider adding stop_gradient to the loss function, but the first-principles derivation and the fundamental result shouldn't be coupled to stop_gradient—otherwise it means the result is strongly coupled to the optimizer and its dynamics, which is not how an essential result ought to behave.
Summary
Centered on the recently released MeanFlow, this post discussed the idea of accelerating diffusion model generation through the lens of "average velocity."
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.