This post introduces flow matching and mean flow for image generation, aiming to provide a thorough and detailed introduction to these two flow models. Much of Sections 1 and 3 follows the tutorial by Holderrieth and Erives [1].

Prerequisites are Linear Algebra, Probability Theory and Multivariable Calculus, plus a very little bit of Deep Learning.

1 Basics for Image Generation

Abstract:

We consider the task of generating images as generating objects that are represented as vectors \(z \in \mathbb{R}^d\). Generation is the task of sampling from a probability distribution \(p_{\text{data}}\), and during training we have access to a dataset of samples \(z_1, \dots, z_N\) from \(p_{\text{data}}\). Conditional generation assumes that we condition the distribution on a label \(y\): we want to sample from the conditional distribution \(p_{\text{data}}(\cdot \mid y)\), and during training we have access to a dataset of pairs \((z_1, y_1), \dots, (z_N, y_N)\). Our goal is to train a generative model that transforms samples from a simple distribution \(p_{\text{init}}\) (e.g., a Gaussian distribution) into samples from the target distribution \(p_{\text{data}}\).

1.1 Images as Vectors

Let’s begin with the kind of images we will encounter, and with how we represent them numerically.

Consider a gray image of \(h \times w\) pixels, where \(h\) and \(w\) denote the height and width of the image. A gray image can be represented as a matrix:

\[ \text{Gray-Image} = \begin{bmatrix} x_{11} & x_{12} & \ldots & x_{1w}\\ x_{21} & x_{22} & \ldots & x_{2w}\\ \vdots & \vdots & \ddots & \vdots\\ x_{h1} & x_{h2} & \ldots & x_{hw}\\ \end{bmatrix}_{h \times w} \]

where each pixel is assigned an intensity value \(x_{ij} \in \mathbb{R}\), and \(\text{Gray-Image} \in \mathbb{R}^{h \times w}\).

Next, extend gray images to \(h \times w\) RGB images with three color channels, where each channel is an intensity matrix, just like a gray image. That is:

\[ \text{RGB-Image} = \begin{bmatrix} \begin{bmatrix} r_{11} & \dots & r_{1w} \\ \vdots & \ddots & \vdots \\ r_{h1} & \dots & r_{hw} \end{bmatrix}\\ \begin{bmatrix} g_{11} & \dots & g_{1w} \\ \vdots & \ddots & \vdots \\ g_{h1} & \dots & g_{hw} \end{bmatrix}\\ \begin{bmatrix} b_{11} & \dots & b_{1w} \\ \vdots & \ddots & \vdots \\ b_{h1} & \dots & b_{hw} \end{bmatrix} \end{bmatrix}_{3 \times h \times w} \]

To cover both gray images and RGB images, we represent an image as:

\[ \text{Image} \in \mathbb{R}^{c \times h \times w}, \]

where \(c, h, w\) are the number of color channels, the image height and the image width, respectively. (\(\text{Image} \in \mathbb{R}^{h \times w \times c}\) is also acceptable, but with a different interpretation: a matrix whose elements are vectors in \(\mathbb{R}^c\).)

Any \(\text{Image} \in \mathbb{R}^{c \times h \times w}\) can be flattened into a vector \(z \in \mathbb{R}^d\), where \(d = c \cdot h \cdot w\).

For example:

\[ \text{RGB-Image} = \begin{bmatrix} \begin{bmatrix} r_{11} & \dots & r_{1w} \\ \vdots & \ddots & \vdots \\ r_{h1} & \dots & r_{hw} \end{bmatrix}\\ \begin{bmatrix} g_{11} & \dots & g_{1w} \\ \vdots & \ddots & \vdots \\ g_{h1} & \dots & g_{hw} \end{bmatrix}\\ \begin{bmatrix} b_{11} & \dots & b_{1w} \\ \vdots & \ddots & \vdots \\ b_{h1} & \dots & b_{hw} \end{bmatrix} \end{bmatrix}_{3 \times h \times w} \Rightarrow \begin{bmatrix} r_{11} \\ \vdots \\ r_{1w} \\ \vdots \\ r_{hw} \\ \vdots \\ g_{11} \\ \vdots \\ g_{hw} \\ \vdots \\ b_{11} \\ \vdots \\ b_{hw} \end{bmatrix}_{3 \cdot h \cdot w} \in \mathbb{R}^{3 \cdot h \cdot w}. \]

Therefore, throughout this post,

\[ \text{we identify the images to be generated as vectors: } \quad z \in \mathbb{R}^{d}. \]

1.2 Sampling as Generation

For simplicity, throughout this post, we don’t explicitly distinguish between a random variable \(Z\) and a sample \(z\). When we write \(z/Z \sim \text{Distribution}\), it means that \(z/Z\) is either a sample drawn from the distribution, or a random variable that follows the distribution.

Let’s define what it means to “generate an image”.

For example, let’s say we want to generate an image of a dog. Naturally, there are many possible images of dogs that we would be happy with. In particular, there’s no one single “best” image of a dog. Rather, there’s a spectrum of images that fit better or worse. In machine learning, it is common to think of this diversity of possible images as a probability distribution. We call it the data distribution and denote it as \(p_{\text{data}}\). In the example of dog images, this distribution would therefore give higher likelihood to images that look more like a dog. Therefore, how “good” an image fits - a rather subjective statement - is replaced by how “likely” it is under the data distribution \(p_{\text{data}}\). With this, we can mathematically express the task of generation as sampling from the (unknown) distribution \(p_{\text{data}}\).

\[ \text{Generating an object } z \text{ is modeled as sampling from the data distribution } p_{\text{data}}. \]

A generative model is a machine learning model that allows us to generate samples from \(p_{\text{data}}\). In machine learning, we require data to train models. In generative modeling, we usually assume access to a finite number of examples (train data) sampled independently from \(p_{\text{data}}\), which together serve as a proxy for the true distribution.

\[ \text{A dataset consists of a finite number of samples } \quad z_1, z_2 \dots z_N \sim p_{\text{data}}. \]

As the size of our dataset grows very large, it becomes an increasingly better representation of the underlying distribution \(p_{\text{data}}\).

1.3 Conditional Generation

The notation \(p(\cdot \mid y)\) for a conditional distribution describes a distribution \(p\) conditioned on \(y\). Here, “\(\cdot\)” means that you can insert any position \(x\) at “\(\cdot\)” to obtain the probability density \(p(x \mid y)\).

Throughout this post, we don’t explicitly distinguish between a probability density and a probability distribution. The notation \(p\) means either a probability distribution or a probability density.

In many cases, we want to generate an object conditioned on some data \(y\). For example, we might want to generate an image conditioned on \(y =\) “a dog running down a hill covered with snow with mountains in the background”. We can rephrase this as sampling from a conditional distribution:

\[ \text{Conditional generation involves sampling } z \sim p_{\text{data}}(\cdot \mid y), \text{ where } y \text{ is a conditioning variable}. \]

We call \(p_\text{data}(\cdot \mid y)\) the conditional data distribution. The conditional generative modeling task typically involves learning to condition on an arbitrary, rather than fixed, choice of \(y\). Using our previous example, we might alternatively want to condition on a different text prompt, such as \(y =\) “a photorealistic image of a cat blowing out birthday candles”. We therefore seek a single model which may be conditioned on any such choice of \(y\). It turns out that techniques for unconditional generation are readily generalized to the conditional case. We will first focus on unconditional generation, but keep in mind that conditional generation is what we’re building towards.

1.4 From Noise to Data

So far, we have discussed the what of generative modeling: generating samples from \(p_\text{data}\). Here, we briefly discuss the how.

For this, we assume that we have access to some initial distribution \(p_{\text{init}}\) that we can easily sample from, such as the Gaussian distribution \(p_{\text{init}} = N(0, I_d)\). The goal of generative modeling is then to transform samples \(x \sim p_{\text{init}}\) into samples from \(p_{\text{data}}\). We note that \(p_{\text{init}}\) does not have to be as simple as a Gaussian. As we shall see, there are interesting use cases for leveraging this flexibility. Despite this, in the majority of applications we take it to be a simple Gaussian, and it is important to keep that in mind.

2 Flow Basics: the Continuity Equation

This section is not necessary for understanding the flow matching model, but it offers some insight into a rigorous understanding of flow matching.

It presents a detailed walkthrough of the continuity equation, from physical intuition to mathematical formalism. We aim to derive the equation step by step and explain all quantities involved.

2.1 Motivation: Conservation in Physical Systems

Many physical quantities are conserved in nature. For instance:

  • Mass in a closed system
  • Charge in electromagnetic systems
  • Energy in isolated systems
  • Number of particles in fluid dynamics

Although these quantities are conserved globally, they vary locally in both space and time. To describe their local behavior, we define density functions that describe how much of a given conserved quantity exists per unit volume at any point and time:

  • Mass density: \(\rho(x, t)\)
  • Charge density: \(\rho_q(x, t)\)
  • Energy density: \(\epsilon(x, t)\)
  • Particle number density: \(n(x, t)\)

These are all functions \(\mathbb{R}^n \times \mathbb{R}^+ \to \mathbb{R}\) (typically \(n = 3\)).

To model how these densities evolve, we introduce the idea of an abstract fluid parcel at position \(x\) and time \(t\). This fluid parcel carries a certain amount of the conserved quantity. The idea of a fluid parcel is purely conceptual: it is like a massless particle that moves with the flow, carrying the density value \(\rho(x, t)\). Each parcel has a velocity \(\mathbf{v}(x, t)\).

The goal of this section is to derive a mathematical law that describes how the distribution of this quantity evolves over space and time. This law is known as the continuity equation.

2.2 Fundamental Definitions

Let’s define:

  • \(\rho(x, t): \mathbb{R}^n \times \mathbb{R}^+ \to \mathbb{R}\), the scalar density function of a conserved quantity
  • \(\mathbf{v}(x, t): \mathbb{R}^n \times \mathbb{R}^+ \to \mathbb{R}^n\), the velocity of the moving conserved quantity
  • \(\mathbf{J}(x, t): \mathbb{R}^n \times \mathbb{R}^+ \to \mathbb{R}^n\), the flux vector field, giving how much quantity passes through a unit area per unit time

For a conserved quantity transported by the velocity field \(\mathbf{v}(x, t)\), the flux is given by:

\[ \mathbf{J}(x, t) \triangleq \rho(x, t) \cdot \mathbf{v}(x, t) \]

Now consider a fixed spatial region \(V \subset \mathbb{R}^n\) with boundary \(\partial V\):

  • The total quantity inside \(V\) at time \(t\) is:
\[ Q_V(t) = \int_V \rho(x, t) \, dx \]
  • The time derivative of \(Q_V(t)\) represents how the total amount inside the region \(V\) is changing:
\[ \frac{d}{dt} Q_V(t) = \frac{d}{dt} \int_V \rho(x, t) \, dx \]

This change must result from flux across the boundary \(\partial V\).

\[ \begin{cases} \frac{d}{dt} \int_V \rho(x, t) \, dx > 0: & \text{Total quantity inside } V \text{ increases}\\ \\ \frac{d}{dt} \int_V \rho(x, t) \, dx < 0: & \text{Total quantity inside } V \text{ decreases} \end{cases} \]

2.3 (Optional) Flux Across the Boundary

For each point \(x \in \partial V\), we define the outward unit normal vector \(\mathbf{n}_x\). Let \(dS_x\) be an infinitesimal surface element at \(x\), and define the infinitesimal oriented surface vector element as \(d\mathbf{S}_x = \mathbf{n}_x \, dS_x\), pointing outward.

Then the net outflow of the quantity through \(\partial V\) is:

\[ \oint_{\partial V} \mathbf{J}(x, t) \cdot d\mathbf{S}_x \]

This measures how much of the quantity is leaving the region per unit time.

\[ \begin{cases} \oint_{\partial V} \mathbf{J}(x, t) \cdot d \mathbf{S}_x > 0: & \text{the quantity flows out of the region } V \\ \\ \oint_{\partial V} \mathbf{J}(x, t) \cdot d \mathbf{S}_x < 0: & \text{the quantity flows into the region } V \end{cases} \]

Naturally, by the conservation of the quantity, we obtain:

\[ \frac{d}{dt} \int_V \rho(x, t) \, dx = - \oint_{\partial V} \mathbf{J}(x, t) \cdot d\mathbf{S}_x \]

2.4 (Optional) Divergence

Additional mathematical background for the continuity equation. Feel free to skip this part if you are familiar with divergence and Gauss’s divergence theorem.

2.4.1 Definition and Geometric Interpretation

Distinguish the divergence \(\nabla \cdot \mathbf{J}\) from the gradient \(\nabla \mathbf{J}\).

The divergence of a vector field is defined as:

\[ \nabla_x \cdot \mathbf{J} (x, t) \triangleq \lim_{V \to 0} \frac{1}{|V|} \oint_{\partial V} \mathbf{J}(x, t) \cdot d \mathbf{S}_x \]

  • If \(\nabla_x \cdot \mathbf{J}(x, t) > 0\), more is flowing out than in \(\Rightarrow\) local decrease in \(\rho\), and \(x\) is a source.

  • If \(\nabla_x \cdot \mathbf{J}(x, t) < 0\), more is flowing in \(\Rightarrow\) local increase in \(\rho\), and \(x\) is a sink.

  • If \(\nabla_x \cdot \mathbf{J}(x, t) = 0\), net flow is balanced \(\Rightarrow\) \(\rho\) stays unchanged

This is a local measure of how much the vector field is “spreading out” at a point.

This geometric interpretation is essential: divergence at a point tells you whether the field is expanding (positive) or contracting (negative) at that location.

Analytically, this definition is equivalent to:

\[ \lim_{V \to 0} \frac{1}{|V|} \oint_{\partial V} \mathbf{J}(x, t) \cdot d \mathbf{S}_x = \sum_{i=1}^n \frac{\partial J_i(x, t)}{\partial x_i} \]

The proof is in Appendix I: Geometric and Analytical Definitions of Divergence.

2.4.2 Gauss’s Divergence Theorem

The rigorous proof is omitted. We only provide an intuitive interpretation of Gauss’s divergence theorem.

\[ \oint_{\partial V} \mathbf{J}(x, t) \cdot d\mathbf{S}_x = \int_V \nabla_x \cdot \mathbf{J}(x, t) \, dx \]

The theorem relates the total outward flux through the boundary to the divergence of \(\mathbf{J}\) inside the volume. Geometrically:

  • The left-hand side sums the quantity flowing out of every surface patch
  • The right-hand side sums the local expansion or compression of flow inside \(V\)

This formula transforms the boundary integral into a volume integral.

2.5 Continuity Equation

Using Gauss’s divergence theorem:

\[ \begin{align*} \frac{d}{dt} \int_V \rho(x, t) \, dx &=- \oint_{\partial V} \mathbf{J}(x, t) \cdot d\mathbf{S}_x \\[2em] &= - \int_V \nabla_x \cdot \mathbf{J}(x, t) \, dx \end{align*} \]

If \(\rho\) is smooth enough, we also have:

\[ \begin{aligned} \frac{d}{dt} \int_V \rho(x, t) \, dx &= \lim_{\Delta t \to 0} \frac{\int_V \rho(x, t + \Delta t) \, dx - \int_V \rho(x, t) \, dx}{\Delta t} \\[2ex] &= \lim_{\Delta t \to 0} \int_V \frac{\rho(x, t + \Delta t) - \rho(x, t)}{\Delta t} \, dx \\[2ex] &= \int_V \frac{\partial \rho(x, t)}{\partial t} \, dx \end{aligned} \]

Thus:

\[ \begin{align*} \int_V \frac{\partial \rho(x, t)}{\partial t} \, dx &= - \int_V \nabla_x \cdot \mathbf{J}(x, t) \, dx \\[2em] \int_V \left( \frac{\partial \rho(x, t)}{\partial t} + \nabla_x \cdot \mathbf{J}(x, t) \right) dx &= 0\\[2em] \frac{\partial \rho(x, t)}{\partial t} + \nabla_x \cdot \mathbf{J}(x, t) &= 0 \end{align*} \]

Finally, using \(\mathbf{J}(x, t) = \rho(x, t) \cdot \mathbf{v}(x, t)\), we obtain the continuity equation, which expresses the local conservation of the quantity carried by the flow.

\[ \text{Continuity Equation: } \frac{\partial \rho(x, t)}{\partial t} + \nabla_x \cdot (\rho(x, t) \cdot \mathbf{v}(x, t)) = 0 \]

The continuity equation arises naturally when we model a conserved quantity flowing through space. It reflects a deep idea:

Any local change in the quantity must be explained by its flow into or out of that region.

This continuity equation is fundamental in physics and mathematics, appearing in fluid mechanics, electromagnetism, quantum mechanics, and beyond.

3 Flow Matching Model

Abstracts:

We first introduce traditional flow models and their disadvantages. Then, we give two perspectives on a stochastic process, explaining why the velocity vector field \(v(x, t)\) can fully transform a sample from \(p_{\text{init}}\) into a sample from \(p_{\text{data}}\), and how the stochastic process is modeled with the continuity equation. Since \(v(x, t)\) can fully transform samples from \(p_{\text{init}}\) into samples from \(p_{\text{data}}\), all we need to do is model this velocity vector field with a neural network. In the third part, we construct a training target, the marginal velocity, and in the fourth part, we prove that in the loss the marginal velocity can be replaced by a much simpler conditional velocity.

3.1 Flow Models: Transform of Random Variables

For two distributions \(p_{\text{init}}\) and \(p_{\text{target}}\), and two random vectors \(x \sim p_{\text{init}}\) and \(y \sim p_{\text{target}}\), suppose we know the pdf \(p_{\text{init}}(x)\) and the invertible transformation \(y = g(x)\).

Then, any sample of \(p_{\text{init}}\) can be transformed into a sample of \(p_{\text{target}}\), and the density follows from the change of variables formula for random variables:

\[ \left| \int p_{\text{init}}(x) dx \right| = \left| \int p_{\text{target}}(y) dy \right| \Rightarrow p_{\text{target}} (y) = p_{\text{init}} (x) \left| \frac{\partial x}{\partial y} \right|, \]

where \(\left| \dfrac{\partial x}{\partial y} \right|\) denotes the absolute value of the determinant of the Jacobian matrix of the function \(x = g^{-1}(y)\).

For image generation, \(p_{\text{target}}\) is exactly \(p_{\text{data}}\), and thus we succeed in sampling objects from \(p_{\text{data}}\) (generating images).

However, such flow models are difficult in practice for the following reason:

  • Computing a determinant takes \(O(n^3)\) time. For images, \(n = 3 \times 256 \times 256 \approx 2 \times 10^5\), so this cubic cost becomes prohibitively expensive.

Then came the flow matching method [3]!

3.2 Stochastic Process and “Flows”

This part uses the concept of a stochastic process and Section 2, Flow Basics: the Continuity Equation to model the transformation of random variables.

Transforming directly from \(p_{\text{init}}\) to \(p_{\text{data}}\) is difficult. So, this time, we insert intermediate distributions. That is, instead of directly transforming \(p_{\text{init}} \Rightarrow p_{\text{data}}\), we do \(p_{\text{init}} \Rightarrow p_{\text{intermediate 1}} \Rightarrow p_{\text{intermediate 2}} \Rightarrow \ldots \Rightarrow p_{\text{data}}\). To model these intermediate steps, we introduce another variable \(t \in [0, 1]\) and define the distribution at time \(t\) as \(p_t\). Specifically, when \(t = 0\), \(p_0 = p_{\text{init}}\), and when \(t = 1\), \(p_1 = p_{\text{data}}\).

3.2.1 Model Transforms as Stochastic Process

A stochastic process is defined as \(X : \Omega \times [0,1] \to \mathbb{R}^d, (\omega, t) \mapsto X(\omega, t)\), where \(\omega\) is an outcome (an event) and \(t\) is time.

  • Perspective I: A Collection of Random Variables

    Here, \(X_t\) and \(X_t(\omega)\) both denote the random variable; the second explicitly writes the outcome \(\omega\).

    For any fixed time \(t\), the stochastic process reduces to a random variable, i.e., \(X_t: \Omega \to \mathbb{R}^d, \omega \mapsto X_t (\omega)\), and \(X_t \sim p_t\).

    For example, \(X_0 \sim p_0 = p_{\text{init}}\) and \(X_1 \sim p_1 = p_{\text{data}}\). As time \(t\) evolves from \(0\) to \(1\), we get a sequence of random variables, which describes how the random variable evolves over time.

    From this perspective, if we go through every \(t \in [0, 1]\), a stochastic process can be considered as a collection of random vectors. This perspective aligns with the flow model: transformations of random variables.

  • Perspective II: A Collection of Trajectories

    For any fixed outcome \(\omega\), the stochastic process reduces to a trajectory, i.e., \(X_{\omega}: [0, 1] \to \mathbb{R}^d, t \mapsto X_{\omega}(t)\).

    A single trajectory records where the outcome \(\omega\) is mapped by the random variables \(X_t\) as \(t\) evolves from \(0\) to \(1\).

    As we go through every \(\omega \in \Omega\), we obtain a collection of trajectories.

    For every \(\omega \in \Omega\), if we take the derivative of the trajectory \(X_{\omega}(t)\), the result \(\dfrac{d}{d t} X_{\omega}(t)\) is the corresponding velocity! Now, simply omit the outcome \(\omega\) and focus solely on trajectories. We can define a velocity vector field:

    \[ v: \mathbb{R}^d \times [0, 1] \to \mathbb{R}^d, (x, t) \mapsto v(x, t), \]

    by taking the derivative of the trajectories:

    \[ \text{For all } \omega \in \Omega, \ t \in [0, 1], \quad v\left(X_{\omega}(t), t\right) \triangleq \frac{d}{dt} X_{\omega}(t). \]

    Note that the vector field \(v\) is simply a mapping from \(\mathbb{R}^d \times [0, 1]\) to \(\mathbb{R}^d\). That means we don’t care which trajectory the position \(x\) in \(v(x, t)\) comes from: \(x\) is simply a point in \(\mathbb{R}^d\). For this definition to make sense, two different outcomes \(\omega_1\) and \(\omega_2\) must satisfy \(X_{\omega_1}(t) \ne X_{\omega_2}(t)\) for all \(t\); otherwise, if \(X_{\omega_1}(t_i) = X_{\omega_2}(t_i)\) at some time \(t_i\), then \(X_{\omega_1}(t) = X_{\omega_2}(t)\) for all \(t \in [t_i, 1]\).

    This doesn’t mean that two trajectories never interact with each other. It simply means that, for two different outcomes \(\omega_1\) and \(\omega_2\), the positions at the same time \(t\), i.e., their images under the same random variable \(X_t\), should be different.

    From this perspective, if we sample \(x_0\) from the initial distribution \(p_0\), we can equivalently obtain a sample \(x_t\) from the distribution \(p_t\) by following the velocity vector field:

    \[ x_t - x_0 = \int_0^t v(x_\tau, \tau) \, d\tau. \]

    When \(t = 1\), we have successfully sampled \(x_1\) from \(p_{\text{data}}\).

    Thus, just like the transformation of random variables in Perspective I, the velocity vector field in Perspective II can fully achieve the goal of sampling from the target distribution \(p_{\text{data}}\). If we push every point of \(p_{\text{init}}\) along the field, we equivalently obtain samples from \(p_{\text{data}}\). For simplicity, in the rest of this post we simply say that the velocity vector field transforms \(p_{\text{init}}\) into \(p_{\text{data}}\), instead of saying that it transforms a sample from \(p_{\text{init}}\) into a sample from \(p_{\text{data}}\). Compared with computing Jacobians and determinants, this ODE is much easier to simulate. So far, the key idea is to find a velocity vector field that fully describes this stochastic process.

    Once we have the velocity, we use the following procedure to “sample from \(p_\text{data}\)”:

    # Generation/Sample Process
Requires: 
- initial distribution: p_init
- velocity: v(x, t)

set t = 0
set time step n
set step size h = 1 / n
sample x_0 from p_init
set x = x_0

# Euler method to simulate the integral
for i in {0, 1, ..., n-1}, do
    x = x + h * v(x, i * h)
end for

return x

3.2.2 View Stochastic Process with “Flows”

In the rest of this post, when we say “flow conditions”, we mean that there is a quantity that is conserved globally but varies in both space and time. More specifically, at any time \(t\), the total quantity \(Q_t = \int \rho(x, t) \, dx\) is conserved.

From Section 2, Flow Basics, we learned that the concepts of a “flow” and the continuity equation apply to situations where a quantity is conserved globally at all times, but varies in both space and time.

In our case, the stochastic process, the probability mass \(\int p_t(x) \, dx\) is always equal to one at any time \(t \in [0, 1]\), while the probability density \(p_t(x)\) varies with the time \(t\) and the position \(x\). That means the conserved quantity, the probability mass, and the probability density \(p_t(x)\), varying with time \(t\) and position \(x\), exactly satisfy the “flow conditions” mentioned above!

Suppose there exists a stochastic process that transforms the initial distribution into the target distribution. Then the velocity vector field \(v(x, t)\) defined in Perspective II and the corresponding probability density \(p_t(x)\) must satisfy the continuity equation.

3.3 Constructing the Training Target

In the previous sections, we’ve shown how \(v(x, t)\) transforms \(p_{\text{init}}\) into \(p_{\text{data}}\). If we find \(v(x, t)\), we are done. In deep learning, we use a neural network \(v^{\theta}(x, t)\) to approximate \(v(x, t)\). We need a training target \(v^{\text{tgt}} (x, t)\) and a loss function \(\mathcal{L}(v^{\theta}, v^{\text{tgt}})\). In this section, we construct a training target \(v^{\text{tgt}}(x, t)\).

Note the difference between \(v(x, t)\) and \(v^{\text{tgt}}(x, t)\). Here, \(v(x, t)\) is the true but unknown velocity, and we need to construct an acceptable velocity as the training target, \(v^\text{tgt}(x, t)\). \(v(x, t)\) and \(v^\text{tgt}(x, t)\) are not necessarily the same, but \(v^\text{tgt}(x, t)\) must play the role of \(v(x, t)\); in other words, \(v^\text{tgt}(x, t)\) must transform the initial distribution \(p_\text{init}\) into the target distribution \(p_\text{data}\) while satisfying the “flow conditions”.

It’s difficult to find an analytical \(v(x, t)\) directly for arbitrary \((x, t)\) along the stochastic process from \(p_{\text{init}}\) to \(p_{\text{data}}\). Thus, we first look at a “conditional version”, and then use the conditional version to construct the marginal version.

3.3.1 Conditional and Marginal Probability Path

The concept of a probability path is similar to the concept of a stochastic process. A probability path describes how the distribution of a stochastic process evolves with time.

The first step of constructing the training target \(v^{\text{tgt}}\) is by specifying a probability path. Intuitively, a probability path specifies a gradual interpolation between the noise distribution \(p_{\text{init}}\) and the data distribution \(p_{\text{data}}\). You can think of a probability path as a trajectory in the space of distributions.

A deterministic evolution from \(p_\text{init}\) to \(p_\text{data}\) is considered a probability path: \(p_0 = p_\text{init}\), \(p_1 = p_\text{data}\), and the probability distribution at time \(t \in [0, 1]\) is \(p_t\). Since we will soon introduce another kind of probability path, we call the probability path from \(p_\text{init}\) to \(p_\text{data}\) the marginal probability path.

For a data point \(z \in \mathbb{R}^d\) from \(p_\text{data}\), let \(\delta_z\) denote the simplest distribution: sampling from \(\delta_z\) always returns \(z\) (i.e., \(\delta_z\) is deterministic). A conditional (interpolating) probability path is a set of distributions \(p_t(\cdot \mid z)\) over \(\mathbb{R}^d\) such that:

\[ p_0(\cdot \mid z) = p_\text{init}, \; \text{and } \; p_1(\cdot \mid z) = \delta_z, \quad \text{for all } z \in \mathbb{R}^d. \]

In other words, a conditional probability path gradually converts the initial noise distribution \(p_\text{init}\) into a single data point \(z\).

The relationship between the densities of the conditional and the marginal probability paths is:

\[ p_t (x) = \int p_t(x \mid z) p_\text{data}(z) dz \]

Here, \(p_t(x)\) is the probability density of sampling \(x\) from the distribution \(p_t\), \(p_t(x \mid z)\) is the probability density of sampling \(x\) from the distribution \(p_t(\cdot \mid z)\), and \(p_\text{data}(z)\) is the probability density of sampling \(z\) from the distribution \(p_\text{data}\). By the law of total probability, \(p_t(x) = \int p_t(x \mid z) p_\text{data}(z) \, dz\).

3.3.2 Conditional and Marginal Velocity

3.3.2.1 Conditional Velocity

We first define the conditional velocity.

For a conditional probability path,

\[ X_t (\cdot \mid z) = \alpha_t z + \beta_t X_0, \]

where \(X_t (\cdot \mid z) \sim p_t(\cdot \mid z)\), and \(X_0 \sim p_0 (\cdot \mid z) = p_{\text{init}}\), \(X_1 \sim p_1(\cdot \mid z) = \delta_z\) for all \(z \sim p_{\text{data}}\).

Here, \(\alpha_t\) and \(\beta_t\) are two simple monotonic functions \([0, 1] \to [0, 1]\) with:

\[ \begin{cases} \alpha_0 = 0 \\[1.5em] \beta_0 = 1 \end{cases} \quad \text{and} \quad \begin{cases} \alpha_1 = 1 \\[1.5em] \beta_1 = 0 \end{cases} \]

\(\alpha_t\) and \(\beta_t\) are simple functions such as:

\[ \begin{cases} \alpha_t = t \\[1.5em] \beta_t = 1 - t \end{cases} \quad \text{or} \quad \begin{cases} \alpha_t = \sqrt{t} \\[1.5em] \beta_t = \sqrt{1 - t} \end{cases} \]

The conditional velocity can be written as:

\[ \text{For all } \omega \in \Omega \text{ and all } z \in \mathbb{R}^d, \text{ at position } x = X_t(\omega \mid z), \quad v(x, t \mid z) \triangleq \frac{d}{d t} X_t (\omega \mid z) = \dot{\alpha}_t z + \dot{\beta}_t X_0. \]

We write \(x = X_t(\omega \mid z)\) because \(x\) simply denotes a position. Note that, for a given condition \(z\), the inputs of the conditional velocity are the position \(x\) and the time \(t\). On the right-hand side, \(\dfrac{d}{dt} X_t(\omega \mid z)\) treats \(X_t (\omega \mid z)\) as a function of the time \(t\). Since \(x = \alpha_t z + \beta_t X_0\), we can also eliminate \(X_0\) and write the conditional velocity as a function of the position alone: \(v(x, t \mid z) = \dot{\alpha}_t z + \dfrac{\dot{\beta}_t}{\beta_t} (x - \alpha_t z)\) for \(t < 1\) [1].

The velocity \(v(x, t \mid z)\) means the velocity, at position \(x\) and time \(t\), of a trajectory whose destination is \(z\). The word conditional means that this velocity is conditioned on the destination of the trajectory being \(z\).

3.3.2.2 Marginal Velocity

We’ve shown how to construct a conditional velocity: at position \(x\) and time \(t\), conditioned on the destination being \(z\), the velocity is \(v(x, t \mid z)\). But what we need is the marginal velocity \(v(x, t)\), the velocity at position \(x\) and time \(t\). A natural way is to build the marginal velocity from the conditional velocity: the marginal velocity at position \(x\) and time \(t\) is the weighted average of the conditional velocities. Here, \(v(x, t \mid z)\) is the velocity at \((x, t)\) conditioned on the destination being \(z\), and \(p_t(z \mid x)\) is the probability density that the final destination is \(z\), given that we are at position \(x\) at time \(t\). Using Bayes’ formula, we get:

\[ \begin{align*} v^\text{tgt}(x, t) &= \int v(x, t \mid z) p_t(z \mid x) \, dz \\ &= \int v(x, t \mid z) \frac{p_t(x \mid z) p_\text{data}(z)}{p_t(x)} \, dz \end{align*} \]

It may seem tricky that \(p_t(z \mid x)\), the density of the destination being \(z\) given that we are at \((x, t)\), equals \(\dfrac{p_t(x \mid z) p_\text{data}(z)}{p_t(x)}\). It becomes clear once we look at how \(x\) and \(z\) are generated together: first draw a destination \(z \sim p_\text{data}\), then draw a position \(x \sim p_t(\cdot \mid z)\). The pair \((z, x)\) then has the joint density \(p_t(x \mid z) \, p_\text{data}(z)\), and its marginal density in \(x\) is exactly \(p_t(x) = \int p_t(x \mid z) p_\text{data}(z) \, dz\). So \(p_t(x \mid z)\) is the density of being at \(x\) at time \(t\) given the destination \(z\), and Bayes’ formula flips the condition: \(p_t(z \mid x) = \dfrac{p_t(x \mid z) \, p_\text{data}(z)}{p_t(x)}\) is the density of the destination being \(z\) given that we are at \(x\) at time \(t\) [1].

This confusion comes from the notation \(p_t\). From the perspective of functions, we could simply write \(p_t(x \mid z)\) as \(p(x, t, z)\), since \(x\), \(t\) and \(z\) are just three input variables and the probability density is the output. We choose to write it as \(p_t(x \mid z)\), using the distribution at time \(t\) to model the variable \(t\) and a conditional probability to model the variable \(z\). With the joint distribution of \((z, x)\) described above, this choice is perfectly rigorous: it is just a convenient way to keep track of what each variable means.

3.3.3 (Optional) Checking “Flow Conditions”

In the conditional probability path, we simply scale and shift the random variable \(X_0 \sim p_0\), so it must satisfy the “flow condition” \(\int p_t(x \mid z) \, dx = \text{constant}\), and the conditional velocity \(v(x, t \mid z)\) successfully transforms the distribution \(p_0 (\cdot \mid z)\) into \(p_1 (\cdot \mid z)\). In short, the probability density and the corresponding velocity satisfy the continuity equation \(\dfrac{\partial p_t(x \mid z)}{\partial t} + \nabla_x \cdot (p_t(x \mid z) \cdot v(x, t \mid z)) = 0\).

But can the constructed \(v^\text{tgt}(x, t) = \int v(x, t \mid z) p_t(z \mid x) \, dz\) successfully drive the initial distribution \(p_\text{init}\) into the target distribution \(p_\text{data}\), and does this transformation satisfy the “flow conditions”? After all, \(v^\text{tgt}\) is just a velocity we constructed, and we don’t yet know whether the probability mass \(\int p_t (x) \, dx\) is conserved under it. To check that our construction is valid, we use the continuity equation:

\[ \begin{align*} \frac{\partial}{\partial t} p_t(x) &= \frac{\partial}{\partial t} \int p_t(x \mid z) p_\text{data}(z) dz \\ &= \int \frac{\partial}{\partial t} p_t(x \mid z) p_\text{data} (z) dz \\ \end{align*} \]

Then, we use the continuity equation of the conditional probability path:

\[ \frac{\partial}{\partial t} p_t(x \mid z) + \nabla_x \cdot (p_t(x \mid z) \cdot v(x, t \mid z)) = 0 \]

and substitute \(\dfrac{\partial}{\partial t} p_t(x \mid z)\):

\[ \begin{align*} \frac{\partial}{\partial t} p_t(x) &= \int - \nabla_x \cdot \left( p_t(x \mid z) \cdot v(x, t \mid z) \right) p_\text{data} (z) \, dz \\ &= - \int \nabla_x \cdot \left( p_t(x \mid z) \cdot v(x, t \mid z) p_\text{data}(z) \right) dz \\ &= - \nabla_x \cdot \int v(x, t \mid z) \cdot p_t(x \mid z) p_\text{data}(z) \, dz \\ &= - \nabla_x \cdot \int v(x, t \mid z) \frac{p_t(x \mid z) p_\text{data}(z)}{p_t (x)} p_t(x) dz \\ &= - \nabla_x \cdot \left( \left( \int v(x, t \mid z) \frac{p_t(x \mid z) p_\text{data}(z)}{p_t (x)} dz \right) \cdot p_t(x) \right) \\ \end{align*} \]

Then, using the definition \(v^\text{tgt}(x, t) \triangleq \int v(x, t \mid z) p_t(z \mid x) \, dz = \int v(x, t \mid z) \dfrac{p_t(x \mid z) p_\text{data}(z)}{p_t(x)} dz\):

\[ \begin{align*} \frac{\partial}{\partial t} p_t(x) &= - \nabla_x \cdot \left( v^\text{tgt}(x, t) \cdot p_t(x) \right) \\ \frac{\partial}{\partial t} p_t(x) + \nabla_x \cdot (v^\text{tgt}(x, t) \cdot p_t(x)) &= 0 \end{align*} \]

Amazing! The constructed \(v^\text{tgt}(x, t)\) satisfies the continuity equation, so the “flow conditions” are guaranteed.

3.4 Training the Model

3.4.1 Substitute \(v^\text{tgt} (x, t)\) with \(v(x, t \mid z)\)

In the previous section, we showed how to construct the training target \(v^\text{tgt}(x, t) = \int v(x, t \mid z) \dfrac{p_t(x \mid z) p_\text{data}(z)}{p_t(x)} dz\). However, there is no closed-form expression for \(v^\text{tgt}\), since we have no access to the distributions \(p_t(x)\) and \(p_\text{data}(z)\). All we have is the initial distribution \(p_\text{init}\) and finitely many samples from \(p_\text{data}\) (the dataset). Thus, we cannot compute our training target \(v^\text{tgt}(x, t)\). In this section, we show that, for a specific kind of loss function, the mean squared error (MSE), we can use the conditional velocity in place of the marginal velocity and get the same training objective.

The MSE loss of the flow matching model is:

\[ \begin{align*} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t, x \sim p_t} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right] \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right] \end{align*} \]

And the conditional MSE loss is:

\[ \mathcal{L}_{\text{CFM}} (\theta) = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) - v(x, t \mid z) \right\|_2^2 \right] \]

Since the loss is a function of \(\theta\), \(x\) and \(t\) are treated as constants independent of \(\theta\). In this section, we aim to show that:

\[ \mathcal{L}_\text{FM} (\theta)= \mathcal{L}_\text{CFM} (\theta) + C \]

Or, equivalently, in gradient form:

\[ \nabla_\theta \mathcal{L}_\text{FM}(\theta) = \nabla_\theta \mathcal{L}_\text{CFM} (\theta) \]

Here is the proof:

\[ \begin{align*} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t, x \sim p_t} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right] \\ &= \iint \|v^{\theta}(x, t) - v^\text{tgt}(x, t)\|_2^2 \;p_t(x) \, d x \, dt \\ &= \iint \|v^{\theta}(x, t) - v^\text{tgt}(x, t)\|_2^2 \; \left(\int p_t (x \mid z) \, p_{\text{data}}(z) \, dz \right) \, d x \, dt \\ &= \iiint \|v^{\theta}(x, t) - v^\text{tgt}(x, t)\|_2^2 \ \, p_t(x \mid z) \, p_{\text{data}}(z) \, dz \, d x \, dt \\ &= \iiint \|v^{\theta}(x, t) - v^\text{tgt}(x, t)\|_2^2 \ \, p_t(x \mid z) \, p_{\text{data}}(z) \, dt \, d z \, dx \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right] \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t)\right\|_2^2 \right] - 2 \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] + \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\|v^\text{tgt}(x, t) \right\|_2^2 \right] \\ \end{align*} \]

And we have:

\[ \begin{align*} \mathcal{L}_{\text{CFM}}(\theta) &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) - v(x, t \mid z) \right\|_2^2 \right] \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) \right\|_2^2 \right] - 2 \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v(x, t \mid z) \right] + \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\|v(x, t \mid z) \right\|_2^2 \right] \\ \end{align*} \]

Since \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\|v(x, t \mid z) \right\|_2^2 \right]\) and \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\|v^\text{tgt}(x, t) \right\|_2^2 \right]\) are both independent of \(\theta\), if we can show that \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v(x, t \mid z) \right] + \text{constant}\), then the proof is done.

Just as we showed that the two expectations of the squared error agree, \(\mathbb{E}_{t, x \sim p_t} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right] = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| v^\theta(x, t) - v^\text{tgt}(x, t) \right\|_2^2 \right]\), we have \(\mathbb{E}_{t, x \sim p_t} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right]\). So:

\[ \begin{align*} \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] &= \mathbb{E}_{t, x \sim p_t} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] \\ &= \iint {v^\theta(x, t)}^\top v^\text{tgt}(x, t) p_t(x) \, dt \, dx \\ &= \iint {v^\theta(x, t)}^\top \left( \int v(x, t \mid z) \frac{p_t(x \mid z) p_\text{data}(z)}{p_t(x)} \, dz \right) p_t(x) \, dt \, dx \\ &= \iiint {v^\theta(x, t)}^\top v(x, t \mid z) p_t(x \mid z) p_\text{data}(z) \, dz \, dt \, dx \\ &= \iiint {v^\theta(x, t)}^\top v(x, t \mid z) p_\text{data}(z) p_t(x \mid z) \, dt \, dz \, dx \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v(x, t \mid z) \right] \end{align*} \]

Amazing! \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v^\text{tgt}(x, t) \right] = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ {v^\theta(x, t)}^\top v(x, t \mid z) \right]\), and we’ve finished the proof!

Note that the conclusion \(\mathcal{L}_\text{FM} (\theta) = \mathcal{L}_\text{CFM} (\theta) + C\), or \(\nabla_\theta \mathcal{L}_\text{FM}(\theta) = \nabla_\theta \mathcal{L}_\text{CFM} (\theta)\), holds in a more general situation: for an MSE loss, as long as the error is a linear function of \(v^\text{tgt}\), the conclusion still holds. For details, see Appendix II: Linear Error and MSE Loss.

3.4.2 An Overall Training Procedure

We’ve shown how to find the conditional velocity, construct the training target, and replace \(\mathcal{L}_\text{FM}\) with \(\mathcal{L}_\text{CFM}\). Here is the overall procedure for training a neural network \(v^\theta(x, t)\) to approximate \(v(x, t)\). My implementation, with conditional image generation on MNIST and CIFAR-10, is available at github.com/hanbinzheng/flow-matching.

# Training Process
Requires:
- neural network: v_theta(x, t)
- training dataset: data
- initial distribution: p_init
- two simple functions: alpha(t) and beta(t)

for each mini-batch in data, do
    get an image z from the dataset                  # z ~ p_data
    sample x_0 from p_init                           # x_0 ~ p_init
    sample time t
    x_t = alpha(t) z + beta(t) x_0                   # equivalent to sampling x_t from p_t(. | z)
    v_cond = d x_t / dt                              # = alpha'(t) z + beta'(t) x_0
    L(θ) = ||v_theta(x_t, t) - v_cond||_2^2          # compute the loss
    update the model parameters θ via gradient descent on L(θ)
end for

4 Mean Flow Model

4.1 Motivation

Flow matching successfully transforms a sample from the initial distribution \(p_\text{init}\) into a sample from the target distribution \(p_\text{data}\). However, simulating the ODE takes several steps to produce the final sample. So, can we simulate it in a single step and achieve one-step generation? Yes, we can. That’s the role of the mean flow (MeanFlow) model [2]. Building on flow matching, it uses the average velocity instead of the instantaneous velocity to achieve one-step generation.

4.2 Average Velocity

4.2.1 Definition of Average Velocity

For times \(t\) and \(r\) with \(0 \le t < r \le 1\), and the position \(x_t\) at time \(t\) (sampled from \(p_t\), i.e., \(x_t \sim p_t\)), the average velocity between times \(t\) and \(r\) is defined as:

\[ u(x_t, t, r) \triangleq \frac{1}{r - t} \int_t^r v(x_\tau, \tau) \, d \tau \]

where \(x_\tau\) is the position at time \(\tau\) on the trajectory that passes through \(x_t\).

To emphasize the difference between the instantaneous velocity and the average velocity, throughout this post the instantaneous velocity is denoted by \(v\) and the average velocity by \(u\).

Note that the average velocity \(u\) is a function of \(v\) only, and has nothing to do with the neural network \(v^\theta\). That is, \(u = f(v) \triangleq \dfrac{1}{r - t} \int_t^r v \, d \tau\). Conceptually, just as the instantaneous velocity \(v\) serves as the ground-truth field in flow matching, the average velocity in this formulation provides an underlying ground-truth field for learning. Thus, we train a model \(u^\theta\) from scratch, rather than building \(u^\theta\) on top of \(v^\theta\).

By definition, the average velocity must satisfy the following “consistency” conditions for any \(t, s, r \in [0, 1]\):

\[ \begin{cases} \lim_{r \to t} u(x_t, t, r) &= \quad v(x_t, t) \\[2em] (r-t) u(x_t, t, r) &= \quad (s - t) u(x_t, t, s) + (r - s) u(x_s, s, r) \end{cases} \]

Our ultimate goal is to approximate the average velocity with a neural network \(u^\theta (x_t, t, r)\). As long as the neural network is accurate enough, we can sample \(\epsilon\) from \(p_\text{init}\) and use the average velocity \(u^\theta(\epsilon, 0, 1)\) to sample from \(p_\text{data}\) in a single step: \(x_1 = \epsilon + u^\theta(\epsilon, 0, 1)\). However, directly using \(u (x_t, t, r) = \dfrac{1}{r - t} \int_t^r v(x_\tau, \tau) \, d \tau\) to train the neural network is intractable, since we have no access to \(v(x_t, t)\), and even if we had it, computing \(u(x_t, t, r)\) would require an integral. In the next section, we show how to construct an optimization target that is amenable to training, using the defining equation of the average velocity.

4.2.2 Mean Flow Identity

To obtain a formulation amenable to training, we rewrite the definition of the average velocity as:

\[ (r - t) u(x_t, t, r) = \int_t^r v(x_\tau, \tau) \, d \tau \]

Now, differentiate both sides with respect to \(t\). Since \(r\) is independent of \(t\):

\[ \begin{align*} \frac{d}{d t} \left[ (r - t) u(x_t, t, r) \right] &= \frac{d}{d t} \int_t^r v(x_\tau, \tau) \, d \tau \\ - u(x_t, t, r) + (r - t) \frac{d}{d t} u(x_t, t, r) &= - v(x_t, t) \\ u(x_t, t, r) &= v(x_t, t) + (r - t) \frac{d}{d t}u(x_t, t, r) \end{align*} \]

We call the equation

\[ u(x_t, t, r) = v(x_t, t) + (r - t) \dfrac{d}{d t}u(x_t, t, r) \]

the MeanFlow identity. It describes the relationship between the instantaneous velocity \(v\) and the average velocity \(u\). In fact, the MeanFlow identity is equivalent to the definition of the average velocity (see Appendix III: Sufficiency of the MeanFlow Identity for details).

4.2.3 JVP to Compute the Derivative

The right-hand side of the MeanFlow identity provides a training target for \(u(x_t, t, r)\), which we will use to construct a loss function for training a neural network. To serve as a suitable target, the time-derivative term must be further decomposed:

\[ \frac{d}{d t} u(x_t, t, r) = \frac{\partial u(x_t, t, r)}{\partial x_t} \frac{d x_t}{d t} + \frac{\partial u(x_t, t, r)}{\partial t} \frac{d t}{d t} + \frac{\partial u(x_t, t, r)}{\partial r} \frac{d r}{d t} \]

This equation shows that the total time derivative can be computed with a Jacobian-vector product (JVP).

More specifically:

\[ u(x_t, t, r): \mathbb{R}^{d} \times [0, 1] \times [0, 1] \to \mathbb{R}^d, \]

or, viewed as a map \(\mathbb{R}^{d+2} \to \mathbb{R}^d\):

\[ u(x_t, t, r) = \begin{bmatrix} u_1 (x_t, t, r) \\ u_2 (x_t, t, r) \\ u_3 (x_t, t, r) \\ \vdots \\ u_d (x_t, t, r) \end{bmatrix}_d \]

\(\dfrac{\partial u(x_t, t, r)}{\partial x_t}\), \(\dfrac{\partial u(x_t, t, r)}{\partial t}\) and \(\dfrac{\partial u(x_t, t, r)}{\partial r}\) are three Jacobian matrices:

\[ \frac{\partial u(x_t, t, r)}{\partial x_t} = \begin{bmatrix} \frac{\partial u_1}{\partial x_{t 1}} & \dots & \frac{\partial u_1}{\partial x_{t d}} \\ \vdots & \ddots & \vdots \\ \frac{\partial u_d}{\partial x_{t 1}} & \dots & \frac{\partial u_d}{\partial x_{t d}} \\ \end{bmatrix}_{d \times d}, \frac{\partial u(x_t, t, r)}{\partial t} = \begin{bmatrix} \frac{\partial u_1}{\partial t} \\ \vdots \\ \frac{\partial u_d}{\partial t} \end{bmatrix}_{d \times 1}, \frac{\partial u(x_t, t, r)}{\partial r} = \begin{bmatrix} \frac{\partial u_1}{\partial r} \\ \vdots \\ \frac{\partial u_d}{\partial r} \end{bmatrix}_{d \times 1} \]

Here,

\[ \frac{d x_t}{dt} = v(x_t, t) = \begin{bmatrix} v_1 (x_t, t) \\ \vdots \\ v_d (x_t, t) \end{bmatrix} , \frac{d t}{d t} = 1, \frac{d r}{d t} = 0 \]

and \(\dfrac{d}{d t} u(x_t, t, r) = \dfrac{\partial u(x_t, t, r)}{\partial x_t} \dfrac{d x_t}{d t} + \dfrac{\partial u(x_t, t, r)}{\partial t} \dfrac{d t}{d t} + \dfrac{\partial u(x_t, t, r)}{\partial r} \dfrac{d r}{d t}\) can be written as:

\[ \begin{align*} \frac{d}{d t} u(x_t, t, r) &= \begin{bmatrix} \frac{\partial u_1}{\partial x_{t 1}} & \dots & \frac{\partial u_1}{\partial x_{t d}} \\ \vdots & \ddots & \vdots \\ \frac{\partial u_d}{\partial x_{t 1}} & \dots & \frac{\partial u_d}{\partial x_{t d}} \\ \end{bmatrix} \begin{bmatrix} v_1 (x_t, t) \\ \vdots \\ v_d (x_t, t) \end{bmatrix} + \begin{bmatrix} \frac{\partial u_1}{\partial t} \\ \vdots \\ \frac{\partial u_d}{\partial t} \end{bmatrix} 1 + \begin{bmatrix} \frac{\partial u_1}{\partial r} \\ \vdots \\ \frac{\partial u_d}{\partial r} \end{bmatrix} 0 \\[2em] &= \overbrace { \underbrace { \begin{bmatrix} \frac{\partial u_1}{\partial x_{t 1}} & \dots & \frac{\partial u_1}{\partial x_{t d}} & \frac{\partial u_1}{\partial t} & \frac{\partial u_1}{\partial r} \\ \vdots & \ddots & \vdots & \vdots & \vdots \\ \frac{\partial u_d}{\partial x_{t 1}} & \dots & \frac{\partial u_d}{\partial x_{t d}} & \frac{\partial u_d}{\partial t} & \frac{\partial u_d}{\partial r}\\ \end{bmatrix} }_{\text{Jacobian matrix of } u: \ \mathcal{J} \in \mathbb{R}^{d \times (d+2)}} \underbrace { \begin{bmatrix} v_1 (x_t, t) \\ \vdots \\ v_d (x_t, t) \\ 1 \\ 0 \end{bmatrix} }_{\text{tangent vector: } \tilde{v} \in \mathbb{R}^{d+2}} }^\text{Jacobian Vector Product (JVP)} \end{align*} \]

In modern libraries, the JVP can be computed efficiently with a JVP interface, such as torch.func.jvp in PyTorch. Rather than constructing the Jacobian matrix, these libraries compute the JVP with forward-mode automatic differentiation.

For simplicity, we write \(\dfrac{d}{d t} u(x_t, t, r)\) as \(\dfrac{\partial u(x_t, t, r)}{\partial x_t} v(x_t, t) + \dfrac{\partial u(x_t, t, r)}{\partial t}\), but keep in mind that it is computed by a single JVP, not by computing three Jacobian matrices and three products.

Using the JVP result, the MeanFlow identity can be written as:

\[ \begin{align*} u(x_t, t, r) &= v(x_t, t) + (r - t) \dfrac{d}{d t}u(x_t, t, r) \\ &= v(x_t, t) + (r - t) \left( \frac{\partial u(x_t, t, r)}{\partial x_t} v(x_t, t) + \frac{\partial u(x_t, t, r)}{\partial t} \right) \end{align*} \]

From the definition of the average velocity, through the MeanFlow identity, to the JVP, we have used nothing but basic calculus. The equation \(u(x_t, t, r) = v(x_t, t) + (r - t) \left( \dfrac{\partial u(x_t, t, r)}{\partial x_t} v(x_t, t) + \dfrac{\partial u(x_t, t, r)}{\partial t} \right)\) is exactly the analytical relation that the average velocity must satisfy. That’s the prototype of the training target. In the next section, we show how to turn it into a target that is amenable to training.

4.3 Training the Model

4.3.1 Constructing the Training Target

Up to now, the formulations above are independent of any neural network. We now introduce a model to learn \(u(x_t, t, r)\). Formally, we parameterize a neural network \(u^\theta(x_t, t, r)\) and encourage it to satisfy the MeanFlow identity. More specifically, we minimize the objective:

\[ \begin{align*} \mathcal{L}_\text{MF} (\theta) &= \mathbb{E}_{t, r, x_t \sim p_t} \left[ \left\| u^\theta(x_t, t, r) - \text{sg} \left[ u^\text{tgt}(x_t, t, r) \right] \right\|_2^2 \right]\\[1em] \text{where } u^\text{tgt} (x_t, t, r) &\triangleq v(x_t, t) + (r - t) \left( \frac{\partial u^\theta(x_t, t, r)}{\partial x_t} v(x_t, t) + \frac{\partial u^\theta (x_t, t, r)}{\partial t} \right) \end{align*} \]

Here, \(v(x_t, t)\) is the \(v^\text{tgt}(x_t, t)\) of flow matching, i.e., the underlying instantaneous velocity field. You may notice that \(u^\text{tgt}\) differs from the MeanFlow identity, and that we apply \(\text{sg} \left[ \cdot \right]\) to \(u^\text{tgt}\).

  • On the difference between \(u^\text{tgt}(x_t, t, r)\) and the MeanFlow identity

    The MeanFlow identity \(u(x_t, t, r) = v(x_t, t) + (r - t) \left( \dfrac{\partial u(x_t, t, r)}{\partial x_t} v(x_t, t) + \dfrac{\partial u(x_t, t, r)}{\partial t} \right)\) only reveals the relationship between the instantaneous velocity \(v(x_t, t)\) and the average velocity \(u(x_t, t, r)\). It does not give a closed-form expression for \(u(x_t, t, r)\), so we cannot use \(v(x_t, t) + (r - t) \left( \dfrac{\partial u(x_t, t, r)}{\partial x_t} v(x_t, t) + \dfrac{\partial u(x_t, t, r)}{\partial t} \right)\) to compute the training target \(u^\text{tgt}(x_t, t, r)\) directly. Instead, we use \(\dfrac{\partial u^\theta(x_t, t, r)}{\partial x_t}\) and \(\dfrac{\partial u^\theta (x_t, t, r)}{\partial t}\) to approximate \(\dfrac{\partial u(x_t, t, r)}{\partial x_t}\) and \(\dfrac{\partial u(x_t, t, r)}{\partial t}\), even though using the model \(u^\theta\) itself to construct its own training target is a little tricky. This approximation is clearly imprecise, but it is what we do.

  • On the \(\text{sg} \left[ \cdot \right]\) operation

    In the loss function, the stop-gradient operation \(\text{sg} \left[ \cdot \right]\) is applied to the training target \(u^\text{tgt}\). This choice is based on experience rather than derivation. It eliminates the need for “double backpropagation” through the JVP, thereby avoiding higher-order optimization. According to Jianlin Su, this operation also has something to do with label leakage and with the ability to achieve one-step generation. And, in my view, the stop-gradient is required to replace the marginal velocity with the conditional velocity (see 4.3.2). Anyway, this operation mainly concerns model training and has little to do with the mathematical derivation, so we simply follow the MeanFlow paper [2] rather than probing into why.

4.3.2 From Loss to “Conditional Loss”

One question may come up here. In practice, we feed \(x_t \sim p_t(\cdot \mid z)\), instead of \(x_t \sim p_t\), into the average velocity \(u(x_t, t, r)\). At first sight this looks wrong, but it is allowed: drawing \(z \sim p_\text{data}\) and then \(x_t \sim p_t(\cdot \mid z)\) produces \(x_t\) with exactly the marginal distribution \(p_t\), which is the same fact we used in 3.4.1.

We’ve constructed the training target and the loss function. However, \(v(x_t, t)\) is still unavailable. Here, we use the conclusion of 3.4.1 (Substitute \(v^\text{tgt}(x, t)\) with \(v(x, t \mid z)\)): for an MSE loss, if the error is a linear function of \(v^\text{tgt}(x, t)\), we can replace the marginal velocity with the conditional velocity and obtain the same training result.

Here:

\[ \begin{align*} \text{error} &= u^\theta(x_t, t, r) - v(x_t, t) - (r - t) \left( \frac{\partial u^\theta(x_t, t, r)}{\partial x_t} v(x_t, t) + \frac{\partial u^\theta (x_t, t, r)}{\partial t} \right) \\[1em] &= \left( u^\theta(x_t, t, r) - (r - t) \frac{\partial u^\theta (x_t, t, r)}{\partial t} \right) - \left( I_d + (r - t) \frac{\partial u^\theta(x_t, t, r)}{\partial x_t}\right) v(x_t, t) \end{align*} \]

To make it clearer, we define:

\[ \begin{cases} b_{\theta, t, r, x_t} &= \quad u^\theta(x_t, t, r) - (r - t) \frac{\partial u^\theta (x_t, t, r)}{\partial t} \\[2em] A_{\theta, t, r, x_t} &= \quad I_d + (r - t) \frac{\partial u^\theta(x_t, t, r)}{\partial x_t} \end{cases} \]

Then the error and the loss function can be rewritten as:

\[ \begin{align*} \text{error} &= b_{\theta, t, r, x_t} - A_{\theta, t, r, x_t} v(x_t, t) \\[2em] \text{and the loss function: } \mathcal{L}_\text{MF} (\theta) &= \mathbb{E}_{t, r, x_t \sim p_t} \left[ \left\| b_{\theta, t, r, x_t} - A_{\theta, t, r, x_t} v(x_t, t) \right\|_2^2 \right] \end{align*} \]

The error is a linear function of the marginal velocity \(v(x_t, t)\). By the conclusion of Appendix II: Linear Error and MSE Loss, minimizing \(\mathcal{L}_\text{MF} (\theta) = \mathbb{E}_{t, r, x_t \sim p_t} \left[ \left\| b_{\theta, t, r, x_t} - A_{\theta, t, r, x_t} v(x_t, t) \right\|_2^2 \right]\) is equivalent to minimizing \(\mathcal{L}_\text{CMF} (\theta) = \mathbb{E}_{t, r, z \sim p_\text{data}, x_t \sim p_t(\cdot \mid z)} \left[ \left\| b_{\theta, t, r, x_t} - A_{\theta, t, r, x_t} v(x_t, t \mid z) \right\|_2^2 \right]\). This replacement makes the loss computable.

Note that Appendix II requires the matrix \(A\) to be independent of the parameters \(\theta\). Here, the matrix \(A_{\theta, t, r, x_t}\) does depend on \(\theta\). But since we apply \(\text{sg} \left[\cdot \right]\) to \(u^\text{tgt}\), \(\dfrac{\partial u^\theta(x_t, t, r)}{\partial x_t}\) is treated as a constant independent of \(\theta\), and so is the matrix \(A_{\theta, t, r, x_t}\). Thus, in my view, the \(\text{sg} \left[\cdot \right]\) operation is necessary to justify replacing the marginal velocity here.

4.3.3 More Training Tricks

  • How to sample \(t\) and \(r\)

    The MeanFlow paper [2] provides two ways to sample \((t, r)\). The first is to sample two points directly from the uniform distribution \(U(0, 1)\), then assign the larger one to \(r\) and the smaller one to \(t\). The second is to sample from a normal distribution \(N(\mu, \sigma)\) and map the samples into \([0, 1]\) with the logistic function \(f(x) = \dfrac{1}{1 + e^{-x}}\); again, the larger one is assigned to \(r\) and the smaller one to \(t\). The second method is called “lognorm” (logit-normal sampling). By adjusting \(\mu\) and \(\sigma\), you can adjust the distribution of \(t\) and \(r\).

  • Adaptive L2 loss

    Replacing the marginal velocity with the conditional velocity requires an MSE loss. In practice, MeanFlow uses the MSE loss with adaptive loss weights. More specifically, the MSE loss is \(\mathcal{L} = \| \Delta \|_2^2\); if we instead use the loss \(\mathcal{L}_{\text{adaptive L2}} = \text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right] \cdot \|\Delta\|_2^2\), it is equivalent to the loss \(\mathcal{L}_{2 \gamma} = \|\Delta\|_2^{2 \gamma}\) while keeping the MSE form (which guarantees that we can replace the marginal velocity with the conditional velocity). Here is the proof:

    \[ \nabla_\theta \, \text{sg} \left[ \frac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right] \cdot \|\Delta\|_2^2 = \text{sg} \left[ \frac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right] \cdot \nabla_\theta \|\Delta\|_2^2 \]

    This shows that, for each sample, minimizing \(\mathcal{L}_{\text{adaptive L2}} = \text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right] \cdot \|\Delta\|_2^2\) moves the parameters in the same direction as minimizing \(\mathcal{L} = \| \Delta \|_2^2\); the only difference is the step size. Here, \(\text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right]\) serves only as a scaling factor: it changes the scale, but not the direction, of the gradient. That’s why we call it an adaptive loss weight. When the error \(\| \Delta \|_2^2\) is large, the adaptive loss weight becomes small, and vice versa. This keeps the gradient at a moderate scale, which helps training.

    We’ve shown that the loss \(\mathcal{L}_{\text{adaptive L2}} = \text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right] \cdot \|\Delta\|_2^2\) plays the role of \(\mathcal{L} = \| \Delta \|_2^2\), but why is it equivalent to the loss \(\mathcal{L}_{2 \gamma} = \|\Delta\|_2^{2 \gamma}\)?

    \[ \begin{align*} \mathcal{L}_{2 \gamma} &= \|\Delta\|_2^{2 \gamma} \\[1em] &= \left( \|\Delta\|_2^2 \right)^\gamma \\[1em] \nabla_\theta \mathcal{L}_{2 \gamma} &= \gamma \left( \|\Delta\|_2^2 \right)^{\gamma - 1} \cdot \nabla_\theta \|\Delta\|_2^2 \\[1em] &= \frac{\gamma}{\left( \|\Delta\|_2^2 \right)^{1 - \gamma}} \cdot \nabla_\theta \|\Delta\|_2^2 \end{align*} \]

    You see, \(\dfrac{\gamma}{\left( \|\Delta\|_2^2 \right)^{1 - \gamma}}\) and \(\text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2)^{1 - \gamma}} \right]\) differ only by the constant factor \(\gamma\), so \(\nabla_\theta \mathcal{L}_{2 \gamma} = \gamma \cdot \nabla_\theta \mathcal{L}_{\text{adaptive L2}}\). Neglecting this factor \(\gamma\), they are the same.

    In practice, we use \(\mathcal{L}_{\text{adaptive L2}} = \text{sg} \left[ \dfrac{1}{(\| \Delta \|_2^2 + c)^{1 - \gamma}} \right] \cdot \|\Delta\|_2^2\) with a small constant \(c\) to avoid division by zero.

4.3.4 Overall Training/Sampling Procedure

Training procedure:

# Training Process
# fn: neural network u_theta(x_t, t, r) that approximates u(x_t, t, r)
for each mini-batch in data, do
    z = get_image()
    t, r = sample_t_r()
    x_0 = sample_from_p_init()

    x_t = alpha(t) z + beta(t) x_0
    v_cond = d x_t / dt                              # = alpha'(t) z + beta'(t) x_0
    u, dudt = jvp(fn, (x_t, t, r), (v_cond, 1, 0))

    u_tgt = stop_gradient(v_cond + (r - t) * dudt)
    error = u - u_tgt

    loss = adaptive_l2_loss(error)
    update the model parameters θ via gradient descent
end for

Sampling procedure:

# Generation/Sample Process
Requires: 
- initial distribution: p_init
- average velocity: u(x, t, r)

sample x_0 from p_init
return x_0 + u(x_0, 0, 1)

Appendix

I. Geometric and Analytical Definitions of Divergence

\[ \lim_{V \to 0} \frac{1}{|V|} \oint_{\partial V} \mathbf{J}(x, t) \cdot d \mathbf{S}_x = \sum_{i=1}^n \frac{\partial J_i(x, t)}{\partial x_i} \]

Consider the three-dimensional case, where \(\mathbf{x} = (x, y, z)\). Take the region \(V\) to be a cube centered at \(\mathbf{x}\) with edge length \(2 \epsilon\), so its volume is \(|V| = 8 \epsilon^3\). Call the six faces of the cube \(S_{\pm x}, S_{\pm y}, S_{\pm z}\), with outward unit normal vectors \(\pm \hat{x}, \pm \hat{y}, \pm \hat{z}\). Write \(\mathbf{J}_{t, x}\) for the \(x\)-component of \(\mathbf{J}(\cdot, t)\), and similarly for \(y\) and \(z\).

The outward flux through the face \(S_{+x}\) is:

\[ \Phi_{+x} = \int_{y, z \in [- \epsilon, \epsilon]} \mathbf{J}((x + \epsilon, y, z), t) \cdot \hat{x} \, dy \, dz \]

or, more simply:

\[ \Phi_{+x} = \int_{y, z} \mathbf{J}_{t, x}(x + \epsilon, y, z) \, dy \, dz \]

Similarly,

\[ \Phi_{-x} = - \int_{y, z} \mathbf{J}_{t, x}(x - \epsilon, y, z) \, dy \, dz \]

The outward flux per unit volume in the \(x\) direction is:

\[ \frac{\Phi_{+x} + \Phi_{-x}}{|V|} = \frac{1}{8 \epsilon^3} \int_{y, z} \left( \mathbf{J}_{t, x} (x + \epsilon, y, z) - \mathbf{J}_{t, x} (x - \epsilon, y, z)\right) \, dy \, dz \]

Since \(\mathbf{J}\) is differentiable, we have:

\[ \mathbf{J}_{t, x} (x \pm \epsilon, y, z) = \mathbf{J}_{t, x} (x, y, z) \pm \epsilon \frac{\partial \mathbf{J}_{t, x}}{\partial x} + O(\epsilon^2) \]

so, as \(\epsilon \to 0\),

\[ \begin{align*} \frac{\Phi_x}{|V|} & = \frac{\Phi_{+x} + \Phi_{-x}}{|V|} \\[1.5em] &= \frac{1}{4 \epsilon^2} \int_{y, z} \frac{\partial \mathbf{J}_{t, x}}{\partial x} \, dy \, dz \\[1.5em] &= \frac{1}{4 \epsilon^2} \frac{\partial \mathbf{J}_{t, x}}{\partial x} \int_{y, z \in [- \epsilon, \epsilon]} dy dz \\[1.5em] &= \frac{\partial \mathbf{J}_{t, x}}{\partial x} \end{align*} \]

Thus, we have:

\[ \begin{align*} \frac{\Phi_{\text{total}}}{|V|} &= \frac{\Phi_x + \Phi_y + \Phi_z}{|V|} \\ &= \frac{\partial \mathbf{J}_{t, x}}{\partial x} + \frac{\partial \mathbf{J}_{t, y}}{\partial y} + \frac{\partial \mathbf{J}_{t, z}}{\partial z} \\ &= \sum_{i=1}^3 \frac{\partial J_i (\mathbf{x}, t)}{\partial x_i} \end{align*} \]

II. Linear Error and MSE Loss

For an MSE loss, if the error is a linear function of \(v^\text{tgt}(x, t)\), the conclusion derived in 3.4.1 still holds.

More specifically, let the matrix \(A\) be independent of \(\theta\), and consider the linear error \(A v^\text{tgt}(x, t) + b\), the MSE loss \(\mathcal{L}(\theta) = \mathbb{E}_{t, x \sim p_t} \left[ \left\| A v^\text{tgt}(x, t) + b \right\|_2^2 \right]\) and the conditional MSE loss \(\mathcal{L}_\text{C} (\theta) = \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| A v(x, t \mid z) + b \right\|_2^2 \right]\). Then \(\mathcal{L}(\theta) = \mathcal{L}_\text{C}(\theta) + C\), or equivalently \(\nabla_\theta \mathcal{L}(\theta) = \nabla_\theta \mathcal{L}_\text{C}(\theta)\), always holds.

Note that we don’t need the vector \(b\) to be independent of \(\theta\). If \(b\) is independent of \(\theta\), then the errors \(A v^\text{tgt}(x, t) + b\) and \(A v(x, t \mid z) + b\) are clearly independent of \(\theta\), so \(\nabla_\theta \mathcal{L}(\theta) = \nabla_\theta \mathcal{L}_\text{C}(\theta) = 0\). The interesting case is when \(b\) is a function of \(\theta\), \(b = f(\theta)\), and then we want to prove that \(\nabla_\theta \mathcal{L}(\theta) = \nabla_\theta \mathcal{L}_\text{C}(\theta)\).

Here is the proof:

\[ \begin{align*} \mathcal{L}(\theta) &= \mathbb{E}_{t, x \sim p_t} \left[ \left\| A v^\text{tgt}(x, t) + b \right\|_2^2 \right] \\ &= \mathbb{E}_{t, x \sim p_t} \left[ \left\| A v^\text{tgt}(x, t) \right\|_2^2 \right] + 2 \mathbb{E}_{t, x \sim p_t} \left[b^\top {A v^\text{tgt}(x, t)} \right] + \mathbb{E}_{t, x \sim p_t} \left[ \left\| b \right\|_2^2 \right] \\[2em] \text{and} \\[2em] \mathcal{L}_\text{C}(\theta) &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| A v(x, t \mid z) + b \right\|_2^2 \right] \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| A v(x, t \mid z) \right\|_2^2 \right] + 2 \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ b^\top A v(x, t \mid z) \right] + \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| b \right\|_2^2 \right] \\ \end{align*} \]

Since \(A\) is a matrix independent of \(\theta\), \(\mathbb{E}_{t, x \sim p_t} \left[ \left\| A v^\text{tgt}(x, t) \right\|_2^2 \right]\) and \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| A v(x, t \mid z) \right\|_2^2 \right]\) are both independent of \(\theta\), so:

\[ \nabla_\theta \mathbb{E}_{t, x \sim p_t} \left[ \left\| A v^\text{tgt}(x, t) \right\|_2^2 \right] = \nabla_\theta \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| A v(x, t \mid z) \right\|_2^2 \right] = 0 \]

For \( \mathbb{E}_{t, x \sim p_t} \left[ \left\| b \right\|_2^2 \right] \) and \(\mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| b \right\|_2^2 \right]\):

\[ \begin{align*} \mathbb{E}_{t, x \sim p_t} \left[ \left\| b \right\|_2^2 \right] &= \iint \left\| b \right\|_2^2 \, p_t(x) \, dt \, dx \\ &= \iint \left\| b \right\|_2^2 \, \left( \int p_t(x \mid z) p_\text{data}(z) \, dz \right) \, dt \, dx \\ &= \iiint \left\| b \right\|_2^2 \, p_t(x \mid z) p_\text{data}(z) \, dz \, dt \, dx \\ &= \iiint \left\| b \right\|_2^2 \, p_t(x \mid z) p_\text{data}(z) \, dt \, dz \, dx \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ \left\| b \right\|_2^2 \right] \end{align*} \]

Thus, what’s left is to prove:

\[ \nabla_\theta \mathbb{E}_{t, x \sim p_t} \left[b^\top {A v^\text{tgt}(x, t)} \right] = \nabla_\theta \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ b^\top A v(x, t \mid z) \right]. \]

Here is the proof:

\[ \begin{align*} \mathbb{E}_{t, x \sim p_t} \left[b^\top {A v^\text{tgt}(x, t)} \right] &= \iint b^\top A v^\text{tgt}(x, t) \, p_t(x) \, dt \, dx \\ &= \iint b^\top A \left( \int v(x, t \mid z) \frac{p_t(x \mid z) p_\text{data} (z)}{p_t(x)} \, dz \right) \, p_t(x) \, dt \, dx \\ &= \iiint b^\top A v(x, t \mid z)p_t(x \mid z) p_\text{data} (z) \,dz \, dt \, dx \\ &= \iiint b^\top A v(x, t \mid z)p_t(x \mid z) p_\text{data} (z) \,dt \, dz \, dx \\ &= \mathbb{E}_{t, z \sim p_\text{data}, x \sim p_t(\cdot \mid z)} \left[ b^\top A v(x, t \mid z) \right] \end{align*} \]

Thus, \(\mathcal{L}(\theta) = \mathcal{L}_\text{C}(\theta) + C\) and \(\nabla_\theta \mathcal{L}(\theta) = \nabla_\theta \mathcal{L}_\text{C}(\theta)\) holds.

III. Sufficiency of the MeanFlow Identity

In Section 4.2.2, we showed that:

\[ u(x_t, t, r) = \frac{1}{r - t} \int_t^r v(x_\tau, \tau) \, d \tau \quad \Rightarrow \quad u(x_t, t, r) = v(x_t, t) + (r - t) \frac{d}{d t}u(x_t, t, r) \]

Here, we show that

\[ u(x_t, t, r) = v(x_t, t) + (r - t) \frac{d}{d t}u(x_t, t, r) \quad \Rightarrow \quad u(x_t, t, r) = \frac{1}{r - t} \int_t^r v(x_\tau, \tau) \, d \tau \]

Treat \(u(x_t, t, r)\) and \(v(x_t, t)\) as functions of \(t\) alone (with \(r\) fixed and \(x_t\) moving along its trajectory). That is, \(u(t) = u(x_t, t, r)\) and \(v(t) = v(x_t, t)\). Then:

\[ u(x_t, t, r) = v(x_t, t) + (r - t) \frac{d}{d t}u(x_t, t, r) \Rightarrow u(t) = v(t) + (r - t) \frac{d}{d t}u(t) \]

Write this ODE in standard form:

\[ \begin{aligned} \frac{d}{d t} u(t) - \frac{1}{r - t} u(t) &= - \frac{1}{r - t} v(t) \\ \frac{d}{d t} u(t) + \frac{1}{t - r} u(t) &= \frac{1}{t - r} v(t) \end{aligned} \]

Then, multiply both sides by a function \(g(t)\) (an integrating factor):

\[ g(t) \cdot \frac{d}{d t} u(t) + \frac{g(t)}{t - r} u(t) = \frac{g(t)}{t - r} v(t) \]

We want the left-hand side to be the derivative of a product:

\[ \begin{align*} \frac{d}{dt} (u(t) \cdot g(t)) &= g(t) \cdot \frac{d}{dt} u(t) + u(t) \cdot \frac{d}{d t} g(t) \\ &= g(t) \cdot \frac{d}{d t} u(t) + \frac{g(t)}{t - r} u(t) \end{align*} \]

Then:

\[ \frac{d}{dt} (u(t) \cdot g(t)) = \frac{g(t)}{t - r} v(t) \]

Integrate both sides from \(t\) to \(r\):

\[ \begin{align*} \frac{d}{d \tau} (u(\tau) \cdot g(\tau)) &= \frac{g(\tau)}{\tau - r} v(\tau) \\[1em] d (u(\tau) \cdot g(\tau)) &= \frac{g(\tau)}{\tau - r} v(\tau) \, d \tau \\[1em] \int_t^r d (u(\tau) \cdot g(\tau)) &= \int_t^r \frac{g(\tau)}{\tau - r} v(\tau) \, d \tau \\[1em] u(r) \cdot g(r) - u(t) \cdot g(t) &= \int_t^r \frac{g(\tau)}{\tau - r} v(\tau) \, d \tau \end{align*} \]

Next, we need to find \(g(t)\). Comparing the two expressions for \(\dfrac{d}{dt} (u(t) \cdot g(t))\):

\[ \begin{align*} \frac{d}{dt} (u(t) \cdot g(t)) &= g(t) \cdot \frac{d}{dt} u(t) + u(t) \cdot \frac{d}{d t} g(t) \\[1em] &= g(t) \cdot \frac{d}{d t} u(t) + \frac{g(t)}{t - r} u(t) \\[1em] & \Downarrow \\ \frac{d g (t)}{d t} &= \frac{g(t)}{t - r} \\[1em] \end{align*} \]

Thus:

\[ \begin{align*} \frac{1}{g(t)} d g(t) &= \frac{1}{t - r} dt \\[1em] \ln |g(t)| &= \ln |t - r| + C \\[1em] g(t) &= A \, |t - r| \\[1em] &= A (r - t), \quad \text{where } A = e^C > 0 \text{ (since } t < r). \end{align*} \]

Substituting \(g(t) = A (r-t)\) into the equation \(u(r) \cdot g(r) - u(t) \cdot g(t) = \int_t^r \frac{g(\tau)}{\tau - r} v(\tau) \, d \tau\), we obtain:

\[ \begin{align*} u(r) \cdot A(r - r) - u(t) \cdot A(r - t) &= \int_t^r \frac{A (r - \tau)}{\tau - r} v(\tau) \, d \tau \\[1em] - A (r - t) \, u(t) &= - A \int_t^r v(\tau) \, d \tau \\[1em] u(t) &= \frac{1}{r - t} \int_t^r v(\tau) \, d \tau \end{align*} \]

Thus:

\[ u(x_t, t, r) = v(x_t, t) + (r - t) \frac{d}{d t}u(x_t, t, r) \quad \Rightarrow \quad u(x_t, t, r) = \frac{1}{r - t} \int_t^r v(x_\tau, \tau) \, d \tau \]

References

[1] Peter Holderrieth and Ezra Erives. “An Introduction to Flow Matching and Diffusion Models.” arXiv:2506.02070, 2025.

[2] Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter and Kaiming He. “Mean Flows for One-step Generative Modeling.” arXiv:2505.13447, 2025.

[3] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel and Matt Le. “Flow Matching for Generative Modeling.” arXiv:2210.02747, 2022.