01-Overview of Flow Matching

The series of tutorial is based on Flow Matching Guide and Code

• arXiv: 2412.06264
• Thank you, META

1 The Definition of the Velocity ODE

Let’s consider a particle moving in the space, we use function $\psi_t(x)$ to describe its current position at time $t\in[0,1]$ starting from $x$. That says, given the initial position $x$ and the time $t$, the function $\psi:[0,1]\times \mathbb{R}^d\rightarrow \mathbb{R}^d$ tells us the current position.

Now we define another function $u_t(x)$, which is a vector field $u:[0,1]\times \mathbb{R}^d \rightarrow \mathbb{R}^d$, this vector field describes the velocity of the particle. Given the current position $x$ and time $t$, it tells us the velocity of the particle.

We now can define an ODE describing the movement of the particle:

$$\frac{\mathrm{d}}{\mathrm{d}t} \psi_t(x) = u_t(\psi_t(x))\tag{1}$$

Basically, this equation says, the derivative of the position function $\psi_t$ is the velocity of the particle given by the vector field function $u_t$. We define that $\psi_0(x)=x$, since at $t=0$ the particle is at the initial position.

Now, if we know the velocity field $u_t$, solving the ODE will generate a path flowing from $\psi_0$ to $\psi_1$, we say the velocity field $u_t$ generates the probability path $p_t$ if its flow $\psi_t$ satisfies

$$X_t := \psi_t(X_0)\sim p_t \ \ \text{for}\ \ X_0\sim p_0$$

Therefore, if we have the velocity field, solving the ODE provides us $X_1=\psi_1(X_0)$, let’s say $X_1$ is from target distribution $q$ and $X_0$ is from the source distribution $p$, then $p_0=p$ and $p_1=q$, Flow Matching is trying to learn the field $u_t^\theta$ such that its flow $\psi_t$ generates such a probability path.

2 Flow Matching Objective

Let’s start from a standard example, let the source distribution be a standard Gaussian $p:=p_0=\mathcal{N}(x|0,I)$, and the probability path $p_t$ is the aggregation of the conditioned probability path $p_{t|1}(x|x_1)$, where $x_1$ is from the training dataset, let’s take the aggregation:

$$p_t(x) = \int p_{t|1}(x|x_1)q(x_1)dx_1\quad p_{t|1}(x|x_1)=\mathcal{N}(x|tx_1,(1-t)^2I)$$

This is just the marginal probability calculation, this path is also known as the conditional optimal-transport or linear path. Using this path, we may say $X_t\sim p_t$ is the linear combination of $X_0\sim p$ and $X_1\sim q$, note that here the $X_0$ and $X_1$ are random variables drawn from distributions.

$$X_t = [tX_1 + (1-t)X_0]\sim p_t\tag{2}$$

The flow matching loss is straightforward

$$\mathcal{L}_{FM} = \mathbb{E}_{t,X_t}\|u_t^\theta(X_t)-u_t(X_t)\|^2\quad t\sim \mathcal{U}[0,1]\ \ \text{and}\ \ X_t\sim p_t$$

The above loss function is intractable, but this objective simplifies drastically if we condition the loss on a single target example $X_1=x_1$ picked randomly from training set.

From Equation (2) let’s write

$$X_{t|1} = tx_1 + (1-t)X_0\ \ \sim\ \ p_{t|1}(\cdot|x_1) = \mathcal{N}(\cdot|tx_1,(1-t)^2I)\tag{3}$$

Now $x_1$ is a specific sample from training dataset, let’s go back to Equation (1) and solve

$$\frac{\mathrm{d}}{\mathrm{d} t} X_{t|1}=u_t(X_{t|1}|x_1)\Rightarrow x_1-X_0 = u_t(x|x_1) \tag{4}$$

Rearrange Equation (3) gives us

$$X_0 = \frac{X_{t|1}-tx_1}{1-t}\tag{5}$$

Take (5) back to (4) yields

$$u_t(x|x_1) = \frac{x_1-x}{1-t} \tag{6}$$

this is the conditional vector field.

Now let’s take (6) into our objective, we obtain the conditional Flow Matching loss

$$\begin{aligned} \mathcal{L}_{CFM} &= \mathbb{E}_{t,X_t,X_1}\|u_t^\theta(X_t)-u_t(X_t|X_1)\|^2 \\ &=\mathbb{E}_{t,X_t,X_1}\left\|u_t^\theta(X_t) - \frac{X_1-X_t}{1-t}\right\|^2 \\ &\overset{\text{Equ.(2)}}{=} \mathbb{E}_{t,X_t,X_1}\left\|u_t^\theta(X_t) - \frac{X_1-tX_1-(1-t)X_0}{1-t}\right\|^2\\ &=\mathbb{E}_{t,X_t,X_1}\|u_t^\theta(X_t) - (X_1-X_0)\|^2 \end{aligned}$$

where $t\sim U[0,1]$, $X_0\sim \mathcal{N}(0,I)$ and $X_1\sim q$.

This is the simplest implementation of Flow Matching using optimal transportation and Gaussian source.