The diffusion process

Let $p(x)$ be a target distribution (e.g. images, molecules, …) which we would like to approximate and sample from.

Notation

• ${\beta }_t$ is the variance schedule, ${\beta }_T\approx 1$ (can be clipped to 0.999 if some values are dependent on the inverse of ${\beta }_t$)

• ${\alpha }_t=1-{\beta }_t$

• ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$

• $q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$

• $q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-\overline{{\alpha }_t})I)$

  • $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,\text{where}{ϵ}_t\sim \mathcal{N}(0,1)$

• Using Bayes theorem, one can calculate the posterior $q(x_{t-1}|x_t,x_0)$ in terms of ${\tilde{\beta }}_t$ and $\tilde{\mu }(x_t,x_0)$ which are defined as follows:

  • ${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$ (posterior variance schedule)
  • $\tilde{\mu }(x_t,x_0)=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0+\frac{\sqrt{{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t$
  • $q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_{t}I)$

• $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$

• In pretty much all works, the mean is: ${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))$

• For the variance:

  • DDPM/DDIM: ${\Sigma }_{\theta }(x_t,t))={\sigma }_t^{2}I$ where ${\sigma }_t^2={\beta }_t$ or ${\sigma }_t^2={\tilde{\beta }}_t$
  • “Improved denoising diffusion probabilistic models” model it as ${\Sigma }_{\theta }(x_t,t)=exp(vlog{\beta }_t+(1-v)log{\tilde{\beta }}_t)$ where $v$ is a vector containing one component per dimension

Forward process. (simple)

A diffusion forward process gradually converts $p(x)$ into another known distribution $q(x)$ e.g. $q(x)\sim \mathcal{N}(0,1)$.

As defined by sohl-dickstein_deep_2015, the forward process is Markovian and discrete, and is defined as follows (in the variance-preserving formulation):

• Given a sample $x_0\sim p(x)$, and a variance schedule ${\beta }_1,\dots ,{\beta }_T$, we define: $x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}ϵ,\text{where}ϵ\sim \mathcal{N}(0,1)$

At the end of the forward diffusion process (around $T=1000$ for Stable Diffusion), we have ${\beta }_T=1$ and thus, $x_T\sim \mathcal{N}(0,1)$.

Closed form

Additionally, because $q(x)$ is Gaussian, the forward process allows sampling $x_t$ at an arbitrary timestep $t$ in closed form. Let ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$: $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,\text{where}{ϵ}_t\sim \mathcal{N}(0,1)$

Diffusion process (more involved)

• The forward process is a recursive application of VAEs (or variational diffusion model) where the encoder is a fixed deterministic function, just a Gaussian function i.e. $q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$
• The reverse process focuses on learning the decoder/ reverse distribution $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$
• ${\mu }_{\theta }(x_t,t)$ is usually learned by 1. score/noise matching 2. and is a function of the noise/score.

Variance fixing

• ${\Sigma }_{\theta }(x_t,t))$ is actually fixed in DDPM and DDIM. In the deterministic sampler of DDIM, it is fixed to zero. Otherwise, it is usually linked to the variance schedule of the forward process.
• Analytical-DDIM shows that there exists analytics forms w.r.t the score function for the optimal mean and variance. The optimal mean is the same us in usual works.
• However, at timestep n, the optimal variance is dependent on the L2 norm of the score of the marginal distribution $q_n(x_n)$ of the forward process at the timestep n.
• Thus, the optimal variance can be precomputed once for a pretrained model and reused during sampling.

Learning the variance schedule

• As described in Variational diffusion, it is also possible to learn the variance schedule by parametrizing it and learn it jointly along with the diffusion model. This requires optimizing the ELBO objective (explained in VAEs).
• Conditional Variational Diffusion Models is an extension of Variational diffusion that allows the variance schedule to be conditioned on class information.
  • This is useful as different class information content may have different robustness to noise. You can have class properties being dependent on precise local features, (e.g. to have the ability to differentiate between cheetah and leopard) vs. global features (e.g. no other animal has similar shape to an elephant).

Score-based Intepretation

• Optimizing such a variational diffusion model (VDM) boils down to learning a neural network to predict the original ground truth image from an arbitrarily noisified version of it.

• Equivalently, you can also show that you should learn to predict the noise that determines $x_t$ from $x_0$

• Score-based interpretation

  • Tweedie’s Formula states that the true mean of an exponential family distribution, given samples drawn from it, can be estimated by the maximum likelihood estimate of the samples (aka empirical mean) plus some correction term involving the score of the estimate. The score serves a correction in case of sample bias.

  • For a Gaussian variable $z\sim \mathcal{N}(z;{\mu }_z,{\Sigma }_z)$, the Tweedie’s Formula states that: $\mathbb{E}[{\mu }_z|z]=z+{\Sigma }_z{\nabla }_{z}logp(z)$

  • We know that our noisy samples $x_t\sim \mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$

  • Thus, by Tweedie’s Formula, $\mathbb{E}[{\mu }_{x_t}|x_t]=x_t+(1-{\bar{\alpha }}_t){\nabla }_{x_t}logp(x_t)$

  • Thus, according to the formula, our best estimate estimate for $x_0$ is : $x_0=\frac{x_t}{\sqrt{{\bar{\alpha }}_t}}+\frac{(1-{\bar{\alpha }}_t){\nabla }_{x_t}logp(x_t)}{\sqrt{{\bar{\alpha }}_t}}$

    • Tweedie’s formula allows us to jump directly to x_0 from a noisy sample.

  • Thus, we should have $s_{\theta }(x_t,t)$ is a neural network that learns to predict the score function ${\nabla }_{x_t}logp(x_t)$, which is the gradient of $x_t$ in data space, for any arbitrary noise level t. The score function measures how to move in data space to maximize the log probability.

  • (below, this is our best estimate for $x_{t-1}$, obtained by plugging our best estimate for $x_0$ in the original formulation of ${\mu }_{\theta }(x_t,t)$)

    • ${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}{x}_t+\frac{1-{\alpha }_t}{{\sqrt{\alpha }}_t}{s}_{\theta }(x_t,t)$

Equivalence between noise, sample and score prediction

• We can set

  • ${\mu }_{\theta }(x_t,t)=\frac{\sqrt{{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)}{1-{\bar{\alpha }}_t}{\hat{x}}_{\theta }(x_t,t)$
  • ${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}{x}_t-\frac{(1-{\alpha }_t)}{{\sqrt{\alpha }}_t\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)$
  • ${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}{x}_t+\frac{1-{\alpha }_t}{{\sqrt{\alpha }}_t}{s}_{\theta }(x_t,t)$

• Thus, we get that:

  • Source noise ${ϵ}_0$ and the score ${\nabla }_{x_t}logp(x_t)$ describe something very similar: ${\nabla }_{x_t}logp(x_t)=-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_0$
  • We have $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_0$

Variance schedule

• Cosine schedule from “Improved denoising diffusion probabilistic models” is
  • 
  •