This is a learning note from this series of videos.

Link to the paper: https://arxiv.org/abs/2210.02747

Supplementary papers:

Step-by-Step Diffusion: An Elementary Tutorial

Glow: Generative Flow with Invertible 1x1 Convolutions

1. What is Flow? Understand vector field and ODE.

Generative Modeling: With known samples $x_1 \sim q_1 \text{(unknown distribution)}$, we want to estimate this unknown distribution.

Method: $\underbrace{q_0}_{\text{a known distribution}} \xrightarrow{\phi} \underbrace{q_1}_{\text{unknown distribution to be estimated}}$. This mapping is denoted as $\phi$.

How to solve for $\phi$: 1. Normalizing Flow; 2. Flow Matching (ODE)

Keypoints of Flow Matching:

1. Let $q_0$ and $q_1$ be the initial and final points of ODE.

2. Use neural networks to fit the gradient term in ODE.

3. Solve the ODE.

Preliminaries

ODE → Flow → Normalizing Flow → Continuous Normalizing Flow

Flow

Definition 1. (Flow): A flow is a collection of time-indexed vector fields $v = \{ v_t \}_{t \in [0, 1]}$.

Any flow defines a trajectory taking initial points $x_1$ to final points $x_0$, by transporting the initial point along the velocity fields $\{ v_t \}$. It is equivalent to a transfer between two distributions.

Formally, for velocity field $v$ and initial point $x_1$, consider the ODE

with initial condition $x_1$ at time $t=1$. We write

to denote the solution to the flow ODE at time $t$, terminating at final point $x_0$. That is, RunFlow is the result of transporting point $x_1$ along the flow $v$ up to time $t$.

Flows also define transports between entire distributions by pushing forward points from the source distribution along their trajectories. If $p_1$ is a distribution on initial points, then applying the flow $v$ yields the distribution on final points:

This process is denoted as $p_1 \xrightarrow{v} p_0$, meaning the flow $v$ transports initial distribution $p_1$to final distribution $p_0$.

THE ULTIMATE GOAL OF FLOW MATCHING is to learn a velocity field $v^*$ which transport $p_1 \xrightarrow{v^*} p_0$, where $p_0$ is the target distribution and $p_1$ is some easy-to-sample base distribution (such as Gaussian).

Continuous Normalizing Flows

Let $\mathbb{R}^d$ denote the data space with data points $x = (x^1, \dots, x^d) \in \mathbb{R}^d$.

The probability density path $p: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}_{>0}$ is a time-dependent probability density function, i.e., $\int p_t(x) dx = 1$.

$v: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ is a time-dependent vector field.

A vector field $v_t$ can be used to construct a time-dependent diffeomorphic map, called a flow, $\phi: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$, defined via the ordinary differential equation (ODE):

Here $\phi_t(x)$ is a solution to the ODE, and we call it a flow. We can model the vector field $v_t$ with a neural network $v_t(x ; \theta)$, where $\theta \in \mathbb{R}^p$ are its learnable parameters.

2. The Continuity Equation and Fokker-Planck Equation

Continuity Equation

How to test if a vector field $v_t$ generates a probability path $p_t$?

Conditional VFs for Fokker-Planck probability paths

Consider a Stochastic Differential Equation (SDE) of the standard form:

with time parameter $t$, drift $f_t$, diffusion coefficient $g_t$, and $dw$ is the Wiener process.

The solution to the SDE is $y_t$, which is a stochastic process (a continuous time-dependent variable). Its probability density $p_t(y_t)$ is characterized by the Fokker-Planck equation:

where $\Delta$ represents the Laplace operator (in $y$), namely $\text{div} \nabla$, where $\nabla$ is the gradient operator. We can rewrite the above equation in the form of the continuity equation:

where the vector field

satisfies the continuity equation with the probability path $p_t$, and therefore generates $p_t$.

3. Continuous Normalizing Flow (CNF)

A CNF is used to reshape a simple prior density $p_0$ (e.g., pure noise) to a more complicated one, $p_1$, via the push-forward equation.

The push-forward (or change of variable) operator $*$ is defined by:

Relationships between Flow, Vector Field, and Probability Density

1. Equation (2)(3) describe the relationship between flow $\phi_t(x)$ and vector field $v_t(x)$

2. Equation (8)(9) describe how the flow $\phi_t(x)$ changes the density from $p_0$ to $p_t$.

3. Equation (4) (continuity equation) gives a necessary and sufficient condition to test whether a vector field $v_t$ generates a probability path $p_t$.

4. Flow Matching Objective

Notations:

1. $x_1$: a random variable distributed according to unknown distribution $q(x_1)$. We have access to samples from $q(x_1)$, but no access to the density function.

2. $p_t$: a probability path such that $p_0 = p$ is a simple distribution, e.g., the standard normal $p(x) = \mathcal{N}(x | 0, I)$

3. $p_1$: approximately equal in distribution to $q$.

The Flow Matching objective is designed to match this target probability path, which will allow us to flow from $p_0$ to $p_1$.

Given a target probability density path $p_t(x)$ and a corresponding vector field $u_t(x)$, which generates $p_t(x)$, the Flow Matching (FM) objective is defined as:

where $\theta$ denotes the learnable parameters of the CNF vector field $v_t$ (here is a neural network), $t \sim \mathcal{U}[0,1]$, and $x \sim p_t(x)$.

Problem: the above objective assumes that the density $p_t(x)$ and the ground truth vector field $u_t(x)$ are known, but we don’t.

Solution: Use conditional probability density for unknown items, and verify it’s valid.

5. From $p_t(x|x_1)$ and $u_t(x|x_1)$ to $p_t(x)$ and $u_t(x)$

Construct a target probability path via a mixture of simpler probability paths:

Given a sample $x_1$, let $p_t(x | x_1)$ denote a conditional probability path such that $p_0(x|x_1) = p(x)$ at time $t=0$ and $p_1(x|x_1)$ at $t=1$ to be a distribution concentrated around $x=x_1$, e.g., $p_1(x|x_1) = \mathcal{N}(x|x_1, \sigma^2 I)$ where $\sigma$ is sufficiently small.

The marginal probability path can be given by:

In particular at time $t=1$, the marginal probability $p_1$ is a mixture distribution that closely approximates the data distribution $q$:

Also, we can define the conditional vector fields in the following form:

where $u_t(\cdot|x_1):\mathbb{R}^d \rightarrow \mathbb{R}^d$ is a conditional vector field that generates $p_t(\cdot|x_1)$

Next, we want to prove that:

The vector field $u_t(x)$ in equation(13) generates the probability path $p_t(x)$ in equation(11).

6. The form of $u_t(x)$

Theorem 1. Given vector field $u_t(x|x_1)$ that generate conditional probability paths $p_t(x|x_1)$, for any distribution $q(x_1)$, the marginal vector field $u_t$ in equation (13) generates the marginal probability path $p_t$ in equation (11), i.e., $u_t$ and $p_t$ satisfy the continuity equation (equation (4)).

Proof:

To proof the theorem, we need to check that $p_t$ and $u_t$ satisfy the continuity equation.

Since $u_t(x|x_1)$ generates $p_t(x|x_1)$, by continuity equation we have:

Then we inspect $\frac{d}{dt} p_t(x)$:

It shows that $u_t(x)$ and $p_t(x)$ satisfy the continuity equation.

End of proof.

7. Conditional Flow Matching (Objectives are consistent)

There are still problems. The probability path $p_t(x)$ in Equation (11) and the vector field $u_t(x)$ in Equation (13) are still intractable due to the integration.

To solve this problem, we introduce a new objective $\mathcal{L}_{\text{CFM}}(\theta)$ defined as:

to avoid $u_t(x)$. And the following theorem states that the new objective $\mathcal{L}_{\text{CFM}}(\theta)$ is equal to $\mathcal{L}_{\text{FM}}(\theta)$ up to a constant difference w.r.t. the model parameters $\theta$.

The new objective is tractable as long as we can sample from $p_t(x|x_1)$ and compute $u_t(x|x_1)$. Consequently, this allows us to train a CNF to generate the marginal probability path $p_t$, which approximates the unknown data distribution $q$ at $t=1$. We don’t need to access to marginal probability path or marginal vector field. We simply need to design suitable conditional probability paths and vector fields.

Theorem 2. Assuming that $p_t(x) > 0$ for all $x \in \mathbb{R}^d$ and $t \in [0,1]$, then, up to a constant independent of $\theta$, $\mathcal{L}_{\text{CFM}}$ and $\mathcal{L}_{\text{FM}}$ are equal. Hence, $\nabla_{\theta} \mathcal{L}_{\text{CFM}} = \nabla_{\theta} \mathcal{L}_{\text{FM}}$.

Proof:

Some assumptions to ensure the existence of integrals and the changing of integration order (by Fubini’s Theorem):

1. $q(x)$ and $p_t(x|x_1)$ decrease to zero at sufficient speed as $||x|| \rightarrow \infty$.

2. $u_t, v_t, \nabla_{\theta} v_t$ are bounded.

First, we rewrite the two objectives:

Note that $||u_t(x)||^2$ and $||u_t(x|x_1)||^2$ are both constant w.r.t the parameters $\theta$. Now we only need to prove that the expectations of the first two terms are equal.

Therefore, the first term of two objectives are equal.

Therefore, the second term of two objectives are also equal. This two results lead to:

End of Proof.

8. Derive Conditional Vector Field from Conditional Probability Path (Gaussian)

The Conditional Flow Matching objective works with any choice of conditional probability path and conditional vector fields.

Here we consider a family of Gaussian conditional probability path:

where $\mu: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ is the time-dependent mean of the Gaussian distribution, while $\sigma: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}_{>0}$ describes a time-dependent scalar standard deviation (std).

When $t=0$, $\mu_0(x_1)=0$ and $\sigma_0(x_1)=1$, so that all conditional probability paths converge to the same standard Gaussian noise distribution.

When $t=1$, $\mu_1(x_1) = x_1$ and $\sigma_1(x_1) = \sigma_{\text{min}}$, so that $p_1(x|x_1)$ is a concentrated Gaussian distribution centered at $x_1$.

We decide to choose a simple form of flow (condition on $x_1$):

If we inspect Equation (9), we can verify that $\psi_t$ pushes the noise distribution $p_0(x|x_1) = p(x)$ to $p_t(x|x_1)$

This flow $\psi_t$ provides a vector field that generates the conditional probability path:

Reparameterizing $p_t(x|x_1)$ in terms of just $x_0$ and plugging Equation (16) in the CFM loss we get:

Theorem 3. Let $p_t(x|x_1)$ be a Gaussian probability path as in Equation (14), and $\psi_t$ its corresponding flow map as in Equation (15). Then, the unique vector field that defines $\psi_t$ has the form:

Consequently, $u_t(x|x_1)$ generates the Gaussian path $p_t(x|x_1)$.

9. Instances of Gaussian Conditional Probability Paths

The above formulation is general for arbitrary functions $\mu_t(x)$ and $\sigma_t(x)$. They can be set to any differentiable function satisfying desired boundary conditions.

Example 1. Diffusion Conditional VFs.

Diffusion models start with data points and gradually add noise until it approximates the pure noise. The process can be formulated as stochastic processes, resulting in Gaussian conditional probability paths $p_t(x|x_1)$, with specific choices of mean $u_t(x_1)$ and std $\sigma_t(x_1)$.

The Variance Exploding (VE) path has the form:

where $\sigma_t$ is an increasing function, $\sigma_0 = 0$ and $\sigma_1 \gg 1$. This this case, $\mu_t(x_1) = x_1$ and $\sigma_t(x_1) = \sigma_{1-t}$. By Theorem 3 we have:

The Variance Preserving diffusion path has the form:

where $\alpha_t = e^{-\frac{1}{2}T(t)}$, $T(t) = \int_0^t \beta(s) ds$, and $\beta$ is the noise scale function. It provides the choices of $\mu_t(x_1) = \alpha_{1-t}x_1$ and $\sigma_t(x_1) = \sqrt{1 - \alpha_{1-t}^2}$. Plug them into Theorem 3, we can get the corresponding vector field.