Diffusion & Multimodal Models

Language models generate one token at a time. Images are generated differently: a diffusion model starts from pure noise and removes it step by step. And a model that reads images needs a way to turn pixels into tokens. This page covers both: the denoising diffusion model behind image generation, and the vision transformer and contrastive training behind vision-language models.

Diffusion: learning to remove noise

The idea of a diffusion model (Sohl-Dickstein et al., 2015; Ho, Jain and Abbeel, 2020) is easy to say: take a picture and add a little random noise, again and again, until only noise is left. Then train a network to undo one step of that. If it can undo every step, it can start from pure noise and end with a picture.

The forward process: adding noise (fixed)

Starting from a data point $x_0$ (an image), add a little Gaussian noise (random numbers from a bell-shaped distribution) at each of $T$ steps, with variances $\beta_1, \dots, \beta_T$. DDPM (Denoising Diffusion Probabilistic Models) uses $T = 1000$ and $\beta$ increasing linearly from $10^{-4}$ to $0.02$. With $\alpha_t = 1 - \beta_t$ and $\bar\alpha_t = \prod_{s \le t} \alpha_s$ ($\bar\alpha_t$ is "alpha bar": the product of all the $\alpha$ up to step $t$), any step can be sampled directly from $x_0$:

where $\epsilon \sim \mathcal{N}(0, I)$ means that $\epsilon$ is standard Gaussian noise. In words: $x_t$ is a mix of the clean image and noise; as $t$ grows the clean part shrinks and the noise part grows.

Example. A tiny schedule with $T = 4$ and $\beta = (0.1, 0.2, 0.3, 0.4)$.

• $\alpha = (0.9, 0.8, 0.7, 0.6)$ and $\bar\alpha = (0.9, 0.72, 0.504, 0.3024)$: each is the product of the $\alpha$ so far.
• A single number $x_0 = 2$ with noise $\epsilon = 1$, at step 2: $x_2 = \sqrt{0.72} \cdot 2 + \sqrt{0.28} \cdot 1 = 1.697 + 0.529 = 2.226$. At step 4 the clean part would be $\sqrt{0.3024} \cdot 2 = 1.10$ and the noise part $\sqrt{0.6976} = 0.84$: more noise than signal.

With the real DDPM schedule, $\bar\alpha_1 = 0.9999$, $\bar\alpha_{100} = 0.897$, $\bar\alpha_{500} = 0.079$ and $\bar\alpha_{1000} = 4 \times 10^{-5}$ (our computation): the signal part $\sqrt{\bar\alpha_T}$ is $0.006$, so $x_T$ is almost pure noise.

The reverse process: removing noise (learned)

A network $\epsilon_\theta(x_t, t)$ is trained to predict the noise that was added. The loss is a plain mean squared error, the "simple" objective of Ho et al.:

where $\mathbb{E}$ is the average over random choices of the step $t$ (drawn uniformly), the data point $x_0$ and the noise $\epsilon$, and the bars are the squared length of the difference between the true and the predicted noise. A training step is therefore: take an image, pick a random $t$, noise it with the formula above, ask the network what noise was added, and compare.

Sampling starts from $x_T \sim \mathcal{N}(0, I)$ (pure noise) and applies, for $t = T, \dots, 1$:

with $z \sim \mathcal{N}(0, I)$ (and $z = 0$ at the last step), and $\sigma_t^2 = \beta_t$. The first line removes the predicted noise from $x_t$; the second adds a small amount of fresh noise, which keeps the samples varied.

From a noise prediction one can also estimate the clean image directly: $\hat x_0 = (x_t - \sqrt{1 - \bar\alpha_t}\,\hat\epsilon) / \sqrt{\bar\alpha_t}$, where $\hat\epsilon$ is the predicted noise.

Example. The tiny schedule above, with $x_2 = 2.226$ made from $x_0 = 2$ and noise $1$.

• If the network predicts the noise exactly, $\hat\epsilon = 1$: $\hat x_0 = (2.226 - 0.529 \cdot 1) / 0.849 = 2$: the clean value is recovered exactly.
• One reverse step from $x_2 = 1.0$ with a prediction $\hat\epsilon = 0.5$ and a fresh noise $z = 0.7$: the first line gives $\frac{1}{\sqrt{0.8}}\left(1 - \frac{0.2 \cdot 0.5}{\sqrt{0.28}}\right) = 0.907$; the second adds $\sqrt{0.2} \cdot 0.7 = 0.313$. So $x_1 = 1.220$. At the last step ($t = 1$) the second line is dropped.

Later work changed the details: a cosine schedule that adds noise more gradually (Nichol and Dhariwal, 2021), and latent diffusion (Rombach et al., 2022), which runs the whole process in the compressed latent space of an autoencoder (a smaller representation of the image learned by a network) instead of on pixels, making high-resolution generation much cheaper.

Conditioning and classifier-free guidance

To generate "a photo of a cat", the network also receives the condition $c$ (a class or a text embedding): $\epsilon_\theta(x_t, t, c)$. Classifier-free guidance (Ho and Salimans, 2022) trains the same network to work with and without the condition, by replacing $c$ with a null token $\varnothing$ in a fraction $p_{\text{uncond}}$ of the training examples. At sampling time it extrapolates away from the unconditional prediction:

where $\tilde\epsilon$ is the guided prediction and $w$ the guidance weight. $w = 0$ is the plain conditional model; larger $w$ follows the condition more closely and gives less varied images.

Example. For one number, the conditional prediction is $0.8$ and the unconditional one is $0.5$. The difference, $0.3$, is "what the condition adds". The guided prediction is $0.5 + (1 + w) \cdot 0.3$:

$w$	guided prediction
0	$0.8$ (the conditional one)
1	$1.1$
2	$1.4$
5	$2.3$

The larger $w$, the further the prediction goes in the direction of the condition, past the conditional prediction itself.

Images as tokens: the vision transformer

The ViT (Vision Transformer; Dosovitskiy et al., 2021) applies the transformer of section 7 to images by cutting them into patches. An image of $H \times W$ pixels with $C$ channels (3 for colour: red, green, blue) becomes $N = HW/P^2$ patches of $P \times P$ pixels ("an image is worth 16×16 words"), each flattened into a vector of $P^2 C$ numbers and mapped to the model dimension by one linear layer. A learnable [class] embedding is prepended, as in BERT (Bidirectional Encoder Representations from Transformers), and learnable 1-D position embeddings are added. From there on, the image is a sequence of tokens like any other.

Example. An image of $224 \times 224$ pixels with 3 channels, $P = 16$: $N = (224/16)^2 = 14^2 = 196$ patches, each of $16 \cdot 16 \cdot 3 = 768$ numbers. A sequence of 197 tokens with the class token.

A toy image $4 \times 4$ with one channel, numbered $0, \dots, 15$ row by row, and $P = 2$: there are 4 patches, in row-major order (left to right, top to bottom): $[0, 1, 4, 5]$, $[2, 3, 6, 7]$, $[8, 9, 12, 13]$ and $[10, 11, 14, 15]$. The first patch is the $2 \times 2$ square in the top-left corner, flattened row by row.

Connecting images and text: CLIP

CLIP (Contrastive Language-Image Pre-Training; Radford et al., 2021) trains an image encoder and a text encoder together on image–caption pairs from the web. In a batch of $N$ pairs, each image embedding should match its own caption and none of the other $N - 1$. With both embeddings normalized to unit length, the $N \times N$ table of cosine similarities (the dot products of unit vectors: 1 if they point the same way, 0 if they are unrelated), divided by a learned temperature, is treated as logits, and the loss is the symmetric cross-entropy: once over rows (each image picks its caption) and once over columns (each caption picks its image), with the correct pairs on the diagonal. The temperature is initialized to the equivalent of 0.07.

Example. Two images, embeddings $(1, 0)$ and $(0, 1)$; two captions, embeddings $(1, 0.2)$ and $(0.1, 1)$.

• Cosine similarities: $\begin{pmatrix} 0.98 & 0.10 \\ 0.20 & 0.995 \end{pmatrix}$. The diagonal (the right pairs) is the largest in every row and column.
• With a temperature of $0.5$, the logits are the similarities times 2. Each row's cross-entropy is small, because the right caption wins by a wide margin: the loss is $0.17$. With a temperature of $1$ it is $0.36$. With the learned temperature, $0.07$, the logits are 14 times larger and the loss is almost 0.

A model that cannot tell the captions apart would give the loss $\log 2 = 0.69$ (two equally likely choices).

The result is a shared space where text and images can be compared, used for zero-shot classification (no training for the task: compare an image with the captions "a photo of a {class}", and pick the best) and as the vision part of multimodal models.

Vision-language models

The simplest way to give a language model eyes is to project the image tokens of a pre-trained vision encoder into its embedding space. LLaVA (Large Language and Vision Assistant; Liu et al., 2023) takes the visual features of CLIP's large vision transformer with 14-pixel patches (ViT-L/14), multiplies them by a trainable projection matrix $W$, and feeds the result to the language model as if they were word embeddings. A first stage trains only $W$ on image–caption pairs; a second fine-tunes on visual instructions (section 12). The language model never sees pixels: it sees a few hundred extra tokens.

References

• J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, S. Ganguli, Deep Unsupervised Learning using Nonequilibrium Thermodynamics, arXiv, 2015.
• J. Ho, A. Jain, P. Abbeel, Denoising Diffusion Probabilistic Models, arXiv, 2020. The schedule, $\mathcal{L}_{\text{simple}}$ and Algorithm 2 (sampling).
• A. Nichol, P. Dhariwal, Improved Denoising Diffusion Probabilistic Models, arXiv, 2021. The cosine schedule.
• J. Ho, T. Salimans, Classifier-Free Diffusion Guidance, arXiv, 2022.
• R. Rombach et al., High-Resolution Image Synthesis with Latent Diffusion Models, CVPR, 2022.
• A. Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR, 2021.
• A. Radford et al., Learning Transferable Visual Models From Natural Language Supervision, arXiv, 2021. CLIP; Figure 3 gives the symmetric loss in pseudocode.
• H. Liu, C. Li, Q. Wu, Y. J. Lee, Visual Instruction Tuning, NeurIPS, 2023. LLaVA.