Diffusion model

Consider the problem of image generation. Let ${\displaystyle x}$ represent an image, and let ${\displaystyle p(x)}$ be the probability distribution over all possible images. If we have ${\displaystyle p(x)}$ itself, then we can say for certain how likely a certain image is. However, this is intractable in general.

Most often, we are uninterested in knowing the absolute probability that a certain image is -- when, if ever, are we interested in how likely an image is in the space of all possible images? Instead, we are usually only interested in knowing how likely a certain image is compared to its immediate neighbors -- how more likely is this image of cat, compared to some small variants of it? Is it more likely if the image contains two whiskers, or three, or with some gaussian noise added?

Consequently, we are actually quite uninterested in ${\displaystyle p(x)}$ itself, but rather, ${\displaystyle \nabla _{x}\ln p(x)}$. This performs two effects

• One, we no longer need to normalize ${\displaystyle p(x)}$, but can use any ${\displaystyle {\tilde {p}}(x)=Cp(x)}$, where ${\displaystyle C=\int {\tilde {p}}(x)dx>0}$ is any unknown constant that is of no concern to us.
• Two, we are comparing ${\displaystyle p(x)}$ neighbors ${\displaystyle p(x+dx)}$, by ${\displaystyle {\frac {p(x)}{p(x+dx)}}=e^{-\langle \nabla _{x}\ln p,dx\rangle }}$

Let the score function be ${\displaystyle s(x):=\nabla _{x}\ln p(x)}$, then consider what we can do with ${\displaystyle s(x)}$.

As it turns out, ${\displaystyle s(x)}$ allows us to sample from ${\displaystyle p(x)}$ using stochastic gradient Langevin dynamics, which is essentially an infinitesimal version of Markov chain Monte Carlo.[2]

The score function can be learned by noising-denoising.[1]

Suppose we wish to sample not from the entire distribution of images, but conditional on the image description. We don't want to sample a generic image, but an image that fits the description "black cat with red eyes". Generally, we want to sample from the distribution ${\displaystyle p(x|y)}$, where ${\displaystyle x}$ ranges over images, and ${\displaystyle y}$ ranges over classes of images (a description "black cat with red eyes" is just a very detailed class, and a class "cat" is just a very vague description).

Taking the perspective of the noisy channel model, we can understand the process as follows: To generate an image ${\displaystyle x}$ conditional on description ${\displaystyle y}$, we imagine that the requester really had in mind an image ${\displaystyle x}$, but the image is passed through a noisy channel and came out garbled, as ${\displaystyle y}$. Image generation is then nothing but inferring which ${\displaystyle x}$ the requester had in mind.

In other words, conditional image generation is simply "translating from a textual language into a pictorial language". Then, as in noisy-channel model, we use Bayes theorem to get

$${\displaystyle p(x|y)\propto p(y|x)p(x)}$$

in other words, if we have a good model of the space of all images, and a good image-to-class translator, we get a class-to-image translator "for free". The SGLD uses

$${\displaystyle \nabla _{x}\ln p(x|y)=\nabla _{x}\ln p(y|x)+\nabla _{x}\ln p(x)}$$

where ${\displaystyle \nabla _{x}\ln p(x)}$ is the score function, trained as previously described, and ${\displaystyle \nabla _{x}\ln p(y|x)}$ is found by using a differentiable image classifier.

The classifier-guided diffusion model samples from ${\displaystyle p(x|y)}$, which is concentrated around the maximum a posteriori estimate ${\displaystyle \arg \max _{x}p(x|y)}$. If we want to force the model to move towards the maximum likelihood estimate ${\displaystyle \arg \max _{x}p(y|x)}$, we can use

$${\displaystyle p_{\beta }(x|y)\propto p(y|x)^{\beta }p(x)}$$

where ${\displaystyle \beta >0}$ is interpretable as inverse temperature. In the context of diffusion models, it is usually called the guidance scale. A high ${\displaystyle \beta }$ would force the model to sample from a distribution concentrated around ${\displaystyle \arg \max _{x}p(y|x)}$. This often improves quality of generated images.[7] This can be done simply by SGLD with

$${\displaystyle \nabla _{x}\ln p_{\beta }(x|y)=\beta \nabla _{x}\ln p(y|x)+\nabla _{x}\ln p(x)}$$

If we do not have a classifier ${\displaystyle p(y|x)}$, we could still extract one out of the image model itself:[8]

$${\displaystyle \nabla _{x}\ln p_{\beta }(x|y)=(1-\beta )\nabla _{x}\ln p(x)+\beta \nabla _{x}\ln p(x|y)}$$

Such a model is usually trained by presenting it with both ${\displaystyle (x,y)}$ and ${\displaystyle (x,None)}$, allowing it to model both ${\displaystyle \nabla _{x}\ln p(x|y)}$ and ${\displaystyle \nabla _{x}\ln p(x)}$.

This is an integral part of systems like GLIDE,[9] DALL-E[10] and Google Imagen.[11]