Flow Matching

Let \( p_{\text{data}}(\mathbf{x}) \) be the distribution of real images on \( \mathbb{R}^{d} \). Flow matching models learn a time-dependent vector field

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

whose ODE transports a simple base distribution \( p_{0} \), usually Gaussian, into the complicated data distribution \( p_{1} = p_{\text{data}} \):

\[ \frac{d\mathbf{x}}{dt} = \mathbf{v}_{\theta}(t, \mathbf{x}_{t}) \]

So the training paradigm is simply:

1. Sample noise \( \mathbf{x}_{0}\sim\mathcal{N}(0,I) \)
2. Sample data \( \mathbf{x}_{1}\sim p_{\text{data}} \)
3. Uniformly sample time step \( t\sim \mathcal{U}(0,1) \)
4. Compute noisy sample \( \mathbf{x}_{t} = t \mathbf{x}_{1} + (1-t)\mathbf{x}_{0} \)
5. Compute velocity \( \mathbf{v} = \mathbf{x}_{1} - \mathbf{x}_{0} \)
6. Learn to predict the velocity \( L(\theta) = \mathbb{E}[ \| \mathbf{v}_{\theta}(t, \mathbf{x}_{t}) - \mathbf{v} \|^{2} ] \)

And the sampling paradigm is:

1. Using step size \( \Delta t \), starting from \( t=0 \).
2. Sample noise \( \mathbf{x}_{0} \sim \mathcal{N}(0,I) \)
3. While \( t\lt 1 \), repeat:
   1. \( \Delta \mathbf{x} = \mathbf{v}_{\theta}(\mathbf{x}_{t}, t) \Delta t \)
   2. \( \mathbf{x}_{t+\Delta t} \gets \mathbf{x}_{t} + \Delta \mathbf{x} \)
   3. \( t \gets t + \Delta t \)
4. Output \( \mathbf{x}_{1} \)

Suppose in the system, particles evolve by the ODE \( \frac{dX_{t}}{dt} = v(t,X_{t}) \). Let \( p_{t}(x) \) be the density of \( X_{t} \). Conversation of probability mass gives the continuity equation:

\begin{equation} \label{org650f0d4} \nabla_{t} p_{t}(x) + \nabla \cdot \left( p_{t}\left( x \right) v \left( t,x \right) \right) = 0 \end{equation}

This equation says, the density changes because probability mass flows through space with velocity field \( v \). A vector field satisfying the above continuity equation is called a probability velocity field for the path of distributions \( \{ p_{t} \}_{t\in[0,1]} \)

Now the core of problem is to learn a vector field \( v_{\theta} \) that matches the true velocity field \( v \) such that we can transform \( p_{0} \to p_{1} = p_{\text{data}} \)

Since we want to match the velocity field, a natural objective is the L2 norm between velocities. Suppose there exists a target velocity field \( u_{t}(x) \) at time \( t \) satisfies the continuity equation.

\begin{equation} \label{orgc9db1aa} \mathcal{L}_{FM}(\theta) = \mathbb{E}_{t\sim\mathcal{U}(0,1), x\sim p_{t}}\left[ \left\| v_{\theta}\left( t,x \right) - u_{t}\left( x \right) \right\|^{2} \right] \end{equation}

The issue is that the marginal velocity \( u_{t}(x) \) is usually not directly available. We will handle this with conditional probability paths.

Instead of defining \( p_{t}(x) \) directly, we define conditional paths \( p_{t}(x\mid z) \), then the marginal distribution is

\[ p_{t}(x) = \int p_{t}(x\mid z)q(z)\, dz \]

The endpoint \( t=1 \) is constructed to make \( p_{1}(x) = p_{data}(x) \).

For each conditional path, assume we know a conditional velocity field \( u_{t}(x\mid z) \) satisfying the continuity equation \( \nabla_{t}p_{t}(x\mid z) + \nabla\cdot \left( p_{t}\left( x\mid z \right) u_{t}\left( x\mid z \right) \right) = 0 \).

\begin{equation} \begin{split} p_{t}(x) &= \int p_{t}(x\mid z)q(z)\, dz\\ \implies \nabla_{t} p_{t}(x) &= \int \boxed{\nabla_{t}p_{t}(x\mid z)}q(z) \, dz\\ \nabla_{t} p_{t}(x) &= \int \nabla\cdot (p_{t}(x\mid z)u_{t}(x\mid z)) q(z)\, dz\\ \boxed{\nabla_{t} p_{t}(x)} &= \nabla\cdot \left[\int p_{t}(x\mid z)u_{t}(x\mid z) q(z)\, dz\right]\\ \nabla\cdot \boxed{p_{t}(x)u_{t}(x)} &= \nabla\cdot \left[\int p_{t}(x\mid z)u_{t}(x\mid z) q(z)\, dz\right]\\ u_{t}(x) &= \frac{\left[\int p_{t}(x\mid z)u_{t}(x\mid z) q(z)\, dz\right]}{p_{t}(x)}\\ u_{t}(x) &= \int u_{t}(x\mid z)q(z\mid x,t)\,dz\\ \implies u_{t} &= \mathbb{E}[u_{t}(x\mid z)] \end{split} \end{equation}