This is a learning note of this series of videos.

Paper Link: https://arxiv.org/abs/2011.13456

1. Why we use SDE to describe the diffusion process?

We want to perturb the data with multiple noise scales(why??). The idea is to use SDE to provide an infinite number of noise scales(continuously changing).

The SDE (which is continuous) can be used for theoretical analysis. In practice, we discretize the SDE for numerical computation.

Goal: construct a diffusion process $\{\bold{x}(t)\}^T_{t=0}$ indexed by a continuous time variable $t \in [0, T]$, such that $\bold{x}(0) \sim p_0$, for which we have a dataset of i.i.d. samples, and $\bold{x}(T) \sim p_T$, for which we have a tractable form to generate samples efficiently. This diffusion process can be modeled as the solution to an Ito SDE:

$w$: a Brownian Motion whose increments follow the gaussian distribution, and variance increase with time:

$\bold{f} ( \cdot, t): \mathbb{R}^d \rightarrow \mathbb{R}^d$ is a vector valued function called the drift coefficient of $\bold{x}(t)$, and $g(\cdot): \mathbb{R} \rightarrow \mathbb{R}$ is a scalar function known as the diffusion coefficient of $\bold{x}(t)$.

Starting from samples of $\bold{x}(T) \sim p_T$ and reversing the process, we can obtain samples $\bold{x}(0) \sim p_0$. It is proved that the reverse of a diffusion process is also a diffusion process, running backward in time and given by the reverse-time SDE:

where $\bar{\bold{w}}$ is a standard Wiener process when time flows backwards from $T$ to 0, and $dt$ is a n infinitesimal negative timestep. Once the score of each marginal distribution, $\nabla_{\bold{x}} \log p_t(\bold{x})$, is known for all t, we can derive the above reverse diffusion process and simulate it to sample from $p_0$.

2. How to derive the reverse-time SDE?

Forward SDE: $d\bold{x} = \bold{f}(\bold{x}, t) dt + g(t) d \bold{w}$

Reverse SDE: $d\bold{x} = [\bold{f}(\bold{x}, t) - g(t)^2 \nabla_{\bold{x}} \log p_t(\bold{x})] dt + g(t) d \bar{\bold{w}}$

Proof 1:

Important Assumption: the diffusion coefficient is $g(t)$ rather than $g(\bold{x}, t)$.

Characteristics of the Brownian Motion $\{ w(t), t\geq 0\}$:

1. Gaussian Increments: $w(t) - w(s) \sim \mathcal{N}(0, t-s)$; $w(t) - w(0) \sim \mathcal{N}(0, t)$

2. Independent Increments: If $0 \leq u \leq s \leq t$, then $w(t) - w(s)$ and $w(s) - w(u)$ are independent.

3. Path Continuity: $w(t)$ is a continuous function of $t$.

Discretization of the forward SDE gives:

This indicates that:

According to Bayesian theorem:

Note: First order Taylor Expansion of $f(x)$ at $x_0$ is $f(x) \approx f(x_0) + f'(x_0) (x - x_0)$

Note: $\log p(x_t)$ is actually a function of both $x_t$ and $t$, therefore when taking derivatives, both of them should be considered.

By Taylor Expansion, we have

Plug in the above equation into $p(\bold{x}_t | \bold{x}_{t + \Delta t})$:

Here $C(\Delta t)$ is a polynomial of $\Delta t$ without a constant term. It goes to $0$ when $\Delta t \rightarrow 0$. Therefore,

Previously, we derived $p(\bold{x}_{t + \Delta t} | \bold{x}_t) \propto \exp \left( - \frac{||\bold{x}_{t + \Delta t} - \bold{x}_t - \bold{f}(\bold{x}_t, t) \Delta t||^2}{2 g(t)^2 \Delta t} \right)$ from $d\bold{x} = \bold{f}(\bold{x}, t) dt + g(t) d\bold{w}$. Therefore we can conclude that the corresponding reverse-time SDE is:

3. Core settings of two diffusion models: SMLD and DDPM

3.1 Denoising score matching with Langevin Dynamics (SMLD)

Let the perturbation kernel to be: $p_{\sigma} (\tilde{\bold{x}} | \bold{x}) \coloneqq \mathcal{N}(\tilde{\bold{x}}; \bold{x} , \sigma^2 \bold{I})$, and $p_{\sigma}(\tilde{\bold{x}}) \coloneqq \int p_{data}(\bold{x}) p_{\sigma}(\tilde{\bold{x}} | \bold{x}) d\bold{x}$, where $p_{data}(\bold{x})$ denotes the data distribution. Consider a sequence of positive noise scales $\sigma_{\min} = \sigma_1 < \sigma_2 < \dots < \sigma_N = \sigma_{\max}$. Usually $\sigma_{\min}$ is small enough such that $p_{\sigma_{\min}}(\bold{x}) \approx p_{data}(\bold{x})$ and $\sigma_{\max}$ is large enough such that $p_{\sigma_{\max}}(\bold{x}) \approx \mathcal{N}(\bold{x}; \bold{0}, \sigma_{\max}^2 \bold{I})$

Overall Step (derived from single step): $p(\bold{x}_t | \bold{x}_0) = \mathcal{N}(\bold{x}_0, \sigma_t^2 \bold{I})$

Single Step (by design): $p(\bold{x}_t | \bold{x}_{t-1}) = \mathcal{N}(\bold{x}_{t-1}, (\sigma_t^2 - \sigma_{t-1}^2) \bold{I})$

Previous work propose to train a Noise Conditional Score Network $\bold{s}_{\theta}(\bold{x}, \sigma)$ with a weighted sum of denoising score matching objectives:

Here, $\nabla_{\tilde{\bold{x}}} \log p_{\sigma_i}(\tilde{\bold{x}}|\bold{x}) = - \frac{\tilde{\bold{x}} - \bold{x}}{\sigma_i^2}$. The optimal score-based model $\bold{s}_{\theta^*}(\bold{x}, \sigma)$ can be obtained given sufficient data and model capacity. It matches $\nabla{\bold{x}} \log p_{\sigma}(\bold{x})$ almost everywhere for $\sigma \in \{ \sigma_i \}_{i=1}^N$.

For sampling, the work uses $M$ steps of Langevin MCMC to get a sample for each $p_{\sigma_i}(\bold{x})$ sequentially:

where $\epsilon_i > 0$ is the step size, and $\bold{z}_i^m$ is standard normal. The above process is repeated for $i = N, N-1, \dots, 1$ in turn with $\bold{x}_N^0 \sim \mathcal{N}(\bold{x} | \bold{0}, \sigma_{\max}^2 \bold{I})$ and $\bold{x}_i^0 = \bold{x}_{i+1}^M$ when $i < N$.

Note: $i$ means different noise levels; $M$ means denoise $M$ steps at each noise level; Therefore, to generate a sample, we need to run the score function $N \times M$ times:

3.2 Denoising Diffusion Probabilistic Models (DDPM)

Consider a sequence of positive noise scales $0 < \beta_1, \beta_2, \dots, \beta_N < 1$. For each training data point $\bold{x}_0 \sim p_{data}(\bold{x})$, construct a discrete Markov Chain $\{\bold{x}_0, \bold{x}_1, \dots, \bold{x}_N\}$. The single step and overall step of this chain is:

Single step (by design): $p(\bold{x}_i | \bold{x}_{i-1}) = \mathcal{N} (\bold{x}_i; \sqrt{1 - \beta_i} \bold{x}_{i-1}, \beta_i \bold{I})$

Overall step (derived): $p_{\alpha_i}(\bold{x}_i | \bold{x}_0) = \mathcal{N} (\bold{x}_i ; \sqrt{\alpha_i} \bold{x}_0, (1 - \alpha_i) \bold{I})$

Here $\alpha_i \coloneqq \Pi_{j=1}^i (1 - \beta_j)$

Proof of overall step:

Since $\bold{x}_i = \sqrt{1 - \beta_i}\bold{x}_{i-1} + \sqrt{\beta_i} \bold{z}_i$ and $\bold{x}_{i+1} = \sqrt{1 - \beta_{i+1}} \bold{x}_i + \sqrt{\beta_{i+1}} \bold{z}_{i+1}$, we have:

By induction, we can conclude that:

Therefore, $p_{\alpha_i}(\bold{x}_i | \bold{x}_0) = \mathcal{N} (\bold{x}_i ; \sqrt{\alpha_i} \bold{x}_0, (1 - \alpha_i) \bold{I})$.

End of proof.

The noise scales are prescribed such that $\bold{x}_N$ is approximately distributed according to $\mathcal{N} ( \bold{0}, \bold{I})$

The denoising distribution derived from Bayesian formula is:

and it can be trained with a re-weighted variant of the evidence lower bound (ELBO):

After solving the above equation we get the optimal model $\bold{s}_{\theta^*} (\bold{x}_i, i)$, then samples can be generated by starting from $\bold{x}_N \sim \mathcal{N}(\bold{0}, \bold{I})$ and following the estimated reverse Markov chain as below:

NOTE: we can notice that the weights of the i-th summand in both SMLD and DDPM loss functions, namely $\sigma_i^2$ and $(1 - \alpha_i)$, are related to the corresponding perturbation kernels in the same functional form: $\sigma_i^2 \propto 1 / \mathbb{E} [||\nabla_{\bold{x}} \log p_{\sigma_i}(\tilde{\bold{x}}| \bold{x})||^2_2]$ and $(1 - \alpha_i) \propto 1 / \mathbb{E} [||\nabla_{\bold{x}} \log p_{\alpha_i}(\tilde{\bold{x}}| \bold{x})||^2_2]$

4. VE, VP SDEs: Derive SDEs from Conditional Distributions

The noise perturbations used in SMLD and DDPM can be regarded as discretizations of two different SDEs.

VE Model

When using a total of N noise scales, each perturbation kernel $p_{\sigma_i}(\bold{x} | \bold{x}_0)$ of SMLD corresponds to the distribution of $\bold{x}_i$ in the following Markov chain:

where $\bold{z}_{i-1} \sim \mathcal{N}(\bold{0}, \bold{I})$, and we have introduced $\sigma_0 = 0$ to simplify the notation.

(Recall that single step perturbation of SMLD is: $p(\bold{x}_i | \bold{x}_{i-1}) = \mathcal{N}(\bold{x}_i; \bold{x}_{i-1}, (\sigma_i^2 - \sigma_{i-1}^2) \bold{I})$

When $N \rightarrow \infty$, $\{ \sigma_i \}_{i=1}^N$ becomes a function $\sigma(t)$, $\bold{z}_i$ becomes $\bold{z}(t)$, and the Markov chain $\{ \bold{x}_i \}_{i=1}^N$ becomes a continuous stochastic process $\{ \bold{x}(t) \}_{t=0}^1$, where we have used a continuous time variable $t \in [0, 1]$. The process $\{ \bold{x} (t) \}_{t=0}^1$ is given by the following SDE:

VP Model

Similarly for DDPM, the perturbation kernels are $p(\bold{x}_i | \bold{x}_{i-1}) = \mathcal{N} (\sqrt{1 - \beta_i} x_{i-1}, \beta_i \bold{I})$. Then the discrete Markov chain is:

As $N \rightarrow \infty$, the Markov chain converges to the following SDE:

Proof for SMLD(VE):

Define some notations first: $\bold{x}(\frac{i}{N}) = \bold{x}_i$, $\sigma(\frac{i}{N}) = \sigma_i$, and $\bold{z} (\frac{i}{N}) = \bold{z}_i$ for $i=1,\dots,N$. We can rewrite the Markov chain as follows with $\Delta t = \frac{1}{N}$ and $t \in \{ i, \frac{1}{N}, \dots, \frac{N-1}{N} \}$:

The approximation is given by the definition of derivative: $\frac{d[\sigma^2(t)]}{dt} = \lim_{\Delta t \rightarrow 0} \frac{\sigma^2(t + \Delta t) - \sigma^2 (t)}{\Delta t}$

As $\Delta t \rightarrow 0$, $\sqrt{\Delta t} z(t) = d \bold{w}$ where $\bold{w}$ is a Wiener process. This is because $\bold{w}(t + \Delta t) - \bold{w}(t) \sim \mathcal{N}(0, \Delta t)$.

End of Proof.

Proof for DDPM(VP):

Define an auxiliary set of noise scales $\{ \bar{\beta_i} = N \beta_i \}_{i=1}^N$ and rewrite the Markov chain as below:

In the limit of $N \rightarrow \infty$, $\{ \bar{\beta_i}\}_{i=1}^N$ becomes a function $\beta(t)$ indexed by $t \in [0, 1]$. Let $\beta(\frac{i}{N}) = \bar{\beta_i}$, $\bold{x}(\frac{i}{N}) = \bold{x}_i$, $\bold{z}(\frac{i}{N}) = \bold{z}_i$. We can rewrite the above Markov chain as the following with $\Delta t = \frac{1}{N}$ and $t \in \{ 0, 1, \dots, \frac{N-1}{N} \}$:

The first approximation comes from Taylor Expansion: $\sqrt{1 - x} = 1 - \frac{x}{2} - \frac{x^2}{8} - \dots \approx 1 - \frac{x}{2}$

Therefore, in the limit of $\Delta t \rightarrow 0$, the Markov chain converges to the following VP SDE:

5. How to solve the SDE?

Idea: Solve $\mathbb{E}(\bold{x}_t)$ and $\text{Var}(\bold{x}_t)$, then under the Gaussian assumption, we know that $p_{0t}(\bold{x}_t | \bold{x}_0) \sim \mathcal{N}(\cdot, \cdot)$

Theorem 5.1 (simplified from Equation (5.50), Equation (5.51) in Applied Stochastic Differential Equations)

If the SDE takes the form:

Then the expectation $\bold{m}(t)$ and covariance matrix $\bold{P}(t)$ of $\bold{x}(t)$ have:

Solution to VP SDE:

By Theorem 5.1, we have

Therefore,

For covariance matrix $P(t)$:

Since $\mathbb{E}[\bold{x}(t) - \bold{m}(t)] = 0$, we have:

Since $Cov(\bold{x}_0 | \bold{x}_0) = 0$, we have:

Therefore, the solution to VP model is:

Solution to VE SDE:

By Theorem 5.1:

Therefore, the solution to VE SDE is:

6. Derive mean and variance of the perturbation kernel from sub-VP SDE

Why sub-VP SDE:

1. Perform well on likelihood

2. Variance is bounded by VP SDE

Since VE, VP and sub-VP SDEs all have affine drift coefficients, their perturbation kernels $p_{0t} (\bold{x}(t) | \bold{x}(0))$ are all Gaussian and can be computed in closed-forms. This makes training with the score-matching loss:

Solution to sub-VP SDE:

Corollary 6.1:

Given an ODE of the form $y'(x) + p(x)y(x) = f(x)$, the solution is given by:

where $\mu(t) = \exp \left( \int^t p(\xi)d\xi \right)$.

Similar to VP SDE, we have:

By Theorem 5.1:

By Corollary 6.1:

Denote $\textcircled{1} = \int_0^t \beta(s) \exp \{ \int_0^s \beta(\xi) d\xi \} ds$ and $\textcircled{2} = \int_0^t \beta(s) \exp \{ -\int_0^s \beta(\xi) d\xi \} ds$. Solve them separately.