Flow Matching for Gaussian Probability Paths

Probability Path

Conditional Gaussian Probability Path

The Gaussian conditional probability path forms the foundation of denoising diffusion models and flow matching.

Definition: Let ${\alpha }_t$, ${\beta }_t$ be noise schedulers: two continuously differentiable, monotonic functions with boundary conditions ${\alpha }_0={\beta }_1=0$ and ${\alpha }_1={\beta }_0=1$. We define the conditional probability path as a family of distribution $p_t(x|z)$ over ${\mathbb{R}}^d$:

$$p_t(\cdot |z)=\mathcal{N}({\alpha }_{t}z,{\beta }_t^2{I}_d)$$

Boundary Conditions: The imposed conditions on ${\alpha }_t$ and ${\beta }_t$ ensure:

$$p_0(\cdot |z)=\mathcal{N}(0,I_d),\text{ and }{p}_1(\cdot |z)=\mathcal{N}(z,0)={\delta }_z$$

where $z\in {\mathbb{R}}^d$ is a data point, ${\delta }_z$ is the Dirac delta “distribution” that sample from ${\delta }_z$ always returns $z$.

Marginal Gaussian Probability Path

Sampling Procedure: For $z\sim p_{data}$, $ϵ\sim p_{init}=\mathcal{N}(0,I_d)$, we can sample from marginal path:

$x_t={\alpha }_{t}z+{\beta }_tϵ\sim p_t$

This provides a tracable way to sample from the marginal distribution $p_t(x)=\int p_t(x|z)p_{data}(z)dz$.

Vector Field

Conditinoal Gaussian Vector Field

Definition: Let $\dot{{\alpha }_t}={\partial }_t{\alpha }_t$ and $\dot{{\beta }_t}={\partial }_t{\beta }_t$ denote respective time derivatives of ${\alpha }_t$ and ${\beta }_t$. The conditional Gaussian vector field given by:

$$u_t^{targe}(x|z)=(\dot{{\alpha }_t}-\frac{\dot{{\beta }_t}}{{\beta }_t}{\alpha }_t)z+\frac{\dot{{\beta }_t}}{{\beta }_t}x$$

is a valid conditional vector field model.

Property: This vector field generates ODE trajectories $X_t$ that satisfy $X_t\sim p_t(\cdot |z)=\mathcal{N}({\alpha }_{t}z,{\beta }_t^2{I}_d)$ if $X_0\sim \mathcal{N}(0,I_d)$.

Proof

Construct a conditional flow model ${\psi }_t^{targe}(x|z)$ by defining

$${\psi }_t^{targe}(x|z)={\alpha }_{t}z+{\beta }_{t}x$$

If $X_t$ is the ODE trajectory of ${\psi }_t^{target}(\cdot |z)$ with $X_0\sim p_{init}=\mathcal{N}(0,I_d)$, then

$$X_t={\psi }_t^{target}(X_0|z)={\alpha }_{t}z+{\beta }_t{X}_0\sim \mathcal{N}({\alpha }_{t}z,{\beta }_t^2{I}_d)=p_t(\cdot |z)$$

The conditional vector field is:

$$\begin{array}{rl}\frac{d}{dt}{\psi }_t^{\text{target}}(x|z)& =u_t^{\text{target}}({\psi }_t^{\text{target}}(x|z)|z)\text{for all }x,z\in {\mathbb{R}}^d\\ \stackrel{(i)}{\iff }{\dot{\alpha }}_{t}z+{\dot{\beta }}_{t}x& =u_t^{\text{target}}({\alpha }_{t}z+{\beta }_{t}x|z)\text{for all }x,z\in {\mathbb{R}}^d\\ \stackrel{(ii)}{\iff }{\dot{\alpha }}_{t}z+{\dot{\beta }}_t\left(\frac{x-{\alpha }_{t}z}{{\beta }_t}\right)& =u_t^{\text{target}}(x|z)\text{for all }x,z\in {\mathbb{R}}^d\\ \stackrel{(iii)}{\iff }\left({\dot{\alpha }}_t-\frac{{\dot{\beta }}_t}{{\beta }_t}{\alpha }_t\right)z+\frac{{\dot{\beta }}_t}{{\beta }_t}x& =u_t^{\text{target}}(x|z)\text{for all }x,z\in {\mathbb{R}}^d\end{array}$$

In (ii), we reparameterized $x\to (x-{\alpha }_{t}z)/{\beta }_t$.

Score Function for Conditional Gaussian Probability Paths

For the Gaussian path $p_t(x|z)=\mathcal{N}(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$, we can use the form of the Gaussian probability density to get the conditional Gaussian score function, which is the derivetive of $x$,

$$\nabla \log p_t(x|z)=-\frac{x-{\alpha }_{t}z}{{\beta }_t^2}$$

This linear relationship is a unique property of Gaussian distributions and is fundamental to efficient training.

Flow Matching for Gaussian Conditional Probability Paths

The conditional flow matching loss is

$$\begin{array}{rl}{\mathcal{L}}_{\text{CFM}}(\theta )& ={\mathbb{E}}_{t\sim \text{Unif}(0,1),z\sim p_{\text{data}},x\sim \mathcal{N}({\alpha }_{t}z,{\beta }_t^2{I}_d)}\left[{\Vert u_t^{\theta }(x)-\left({\dot{\alpha }}_t-\frac{{\dot{\beta }}_t}{{\beta }_t}{\alpha }_t\right)z-\frac{{\dot{\beta }}_t}{{\beta }_t}x\Vert }^2\right]\\ & \stackrel{(i)}{=}{\mathbb{E}}_{t\sim \text{Unif}(0,1),z\sim p_{\text{data}},ϵ\sim \mathcal{N}(0,I_d)}\left[{\Vert u_t^{\theta }({\alpha }_{t}z+{\beta }_tϵ)-({\dot{\alpha }}_{t}z+{\dot{\beta }}_tϵ)\Vert }^2\right]\end{array}$$

In (i) we replace $x$ by ${\alpha }_{t}z+{\beta }_tϵ$.

Let us make ${\mathcal{L}}_{\text{CFM}}$ even more concrete for the special case of ${\alpha }_t=t$, and ${\beta }_t=1-t$. The corresponding conditional probability path $p_t(x|z)=\mathcal{N}(tz,(1-t{)}^2)$ is referred to as the (Gaussian) CondOT probability path. Then we have $\dot{{\alpha }_t}=1,\dot{{\beta }_t}=-1$, so that

$${\mathcal{L}}_{\text{CFM}}(\theta )={\mathbb{E}}_{t\sim \text{Unif},z\sim p_{\text{data}},ϵ\sim \mathcal{N}(0,I_d)}\left[{\Vert u_t^{\theta }(tz+(1-t)ϵ)-(z-ϵ)\Vert }^2\right]$$

Models like Stable Diffusion 3, Meta’s Movie Gen Video are using this procedure.

Training Procedure

Given a dataset of samples $z\sim p_{data}$, vector field network $u_t^{\theta }$. For each batch of data:

1. Sample a data example $z$ from the dataset,

2. Sample a random time $t\sim {\text{Unif}}_{[}0,1]$,

3. Sample noise $ϵ\sim \mathcal{N}(0,I_d)$

4. Set $x=tz+(1-t)ϵ$

5. Compute loss ${\mathcal{L}}_{\theta }=||u_t^{\theta }(x)-(z-ϵ)|{|}^2$

6. Update the model.

Score Matching for Gaussian Probability Paths

The conditional score matching loss is

$$\begin{array}{rl}{\mathcal{L}}_{\text{CSM}}(\theta )& ={\mathbb{E}}_{t\sim \text{Unif},z\sim p_{data},x\sim p_t(\cdot |z)}[||s_t^{\theta }(x)+\frac{x-{\alpha }_{t}z}{{\beta }_t^2}|{|}^2]\\ & ={\mathbb{E}}_{t\sim \text{Unif},z\sim p_{data},ϵ\sim \mathcal{N}(0,I_d)}[||s_t^{\theta }({\alpha }_{t}z+{\beta }_tϵ)+\frac{ϵ}{{\beta }_t}|{|}^2]\\ & ={\mathbb{E}}_{t\sim \text{Unif},z\sim p_{data},ϵ\sim \mathcal{N}(0,I_d)}[\frac{1}{{\beta }_t^2}||{\beta }_t{s}_t^{\theta }({\alpha }_{t}z+{\beta }_tϵ)+ϵ|{|}^2]\end{array}$$

Note that $s_t^{\theta }$ learns to predict the noise that was used to corrupt a data sample $z$. Therefore, the above training loss is also called denoising score matching. It was soon realized that the above loss is numerically unstable for ${\beta }_t\approx 0$ close to zero (i.e. denoising score matching only works if you add a sufficient amount of noise).

In Denosing Diffusion Probabilitic Models (DDPM), it was proposed to drop the constant $\frac{1}{{\beta }_t^2}$ in the loss, and reparameterize $s_t^{\theta }$ into a noise predictor network ${ϵ}_t^{\theta }$ via:

$$-{\beta }_t{s}_t^{\theta }(x)={ϵ}_t^{\theta }(x)$$

thus,

$${\mathcal{L}}_{\text{DDPM}}(\theta )={\mathbb{E}}_{t\sim \text{Unif},z\sim p_{data},ϵ\sim \mathcal{N}(0,I_d)}[||{ϵ}_t^{\theta }({\alpha }_{t}z+{\beta }_tϵ)-ϵ|{|}^2]$$