COMP90086 Computer Vision · Workshop 10

From noise to digits

A primer on diffusion models and a walk through the worksheet. Every code cell from the notebook is here, next to a live version you can poke at. The U-Net runs in this page.

clean data \(x_0\) noise \(\epsilon\) network predictions \(\hat\epsilon,\hat x_0\)
16 samples from the worksheet U-Net
A0 · Prologue

From classifiers to diffusion

Most of the models in this course so far answer a question about an image. Generative models go the other way and produce the image. Here is the short path from one to the other.

1

Discriminative models: CNNs on ImageNet 2012

A CNN such as AlexNet2 trained on ImageNet1 maps an image to a label. The softmax head models the conditional distribution of the label given the image, trained with cross-entropy:

$$p_\theta(y\mid \cx{x}) = \operatorname{softmax}\big(f_\theta(\cx{x})\big)_y, \qquad \mathcal L = -\log p_\theta(y\mid \cx{x})$$

LimitThe output is a label \(y\). The model never learns what images look like, so it cannot produce a new \(x\).

2

Autoencoders 2006

An encoder squeezes the image into a small code and a decoder rebuilds it3. For the first time the training target is the data \(x\) itself, not a label:

$$z = E(\cx{x}), \qquad \hat x = D(z), \qquad \mathcal L = \lVert \cx{x} - D(E(\cx{x}))\rVert^2$$

LimitIt is an information compressor, not a generator. Codes of real images land in scattered pockets of latent space, so decoding a random \(z\) usually gives garbage.

3

Variational autoencoders (VAE) 2013

The encoder now outputs a distribution \(q_\phi(z\mid x)=\mathcal N(\mu,\sigma^2)\), and a KL term pulls it towards a standard Gaussian4. Training maximises the evidence lower bound (ELBO):

$$\log p_\theta(\cx{x}) \;\ge\; \underbrace{\mathbb E_{q_\phi(z\mid x)}\big[\log p_\theta(\cx{x}\mid z)\big]}_{\text{reconstruct}} - \underbrace{\mathrm{KL}\big(q_\phi(z\mid \cx{x})\,\|\,\ce{\mathcal N(0,I)}\big)}_{\text{keep the latent Gaussian}}$$

GainSince the latent is forced to look like \(\mathcal N(0,I)\), we can sample it. At inference time the encoder is thrown away: draw \(\ce{z}\sim\mathcal N(0,I)\) and decode \(x = D(\ce{z})\).

LimitSamples tend to be blurry. A pixel-wise Gaussian likelihood rewards averaging over plausible images.

4

Generative adversarial networks (GAN) 2014

A generator \(G\) turns noise into images. A discriminator \(D\) learns to tell real from fake, and its gradients are the generator's training signal5:

$$\min_G \max_D \;\; \mathbb E_{\cx{x}}\big[\log D(\cx{x})\big] + \mathbb E_{\ce{z}\sim\mathcal N(0,I)}\big[\log\big(1 - D(G(\ce{z}))\big)\big]$$

GainMuch sharper, more realistic samples than VAEs.

LimitThe minimax game is hard to optimise: training oscillates or diverges. Mode collapse6 lets the generator cover a few modes of the data and ignore the rest.

5

Many small steps: iterative denoising and diffusion 2015–2020

Autoencoders, VAEs and GANs all generate in one step: one network pass maps \(\ce{z}\) straight to \(\cx{x}\). The map from a Gaussian to the distribution of natural images is enormously complex to learn in a single jump.

So split the journey into many small steps. Each one only has to remove a little noise, which is an easy, well-posed regression problem. This is the idea behind denoising autoencoders7, diffusion probabilistic models8, score-based models9 and DDPM10, the model this workshop implements:

$$\ce{x_T}\sim\mathcal N(0,I) \;\to\; x_{T-1} \;\to\; \cdots \;\to\; x_1 \;\to\; \cx{x_0}$$
FigureOne big jump versus many small steps
AE / VAE / GAN: \(x = G(z)\), one pass
Diffusion: \(x_{t-1} = \text{denoise}(x_t, t)\), many passes
Left: a single learned map has to throw each Gaussian sample straight onto a thin, curved data manifold. Right: the same trip as a random walk of small denoising steps. Each step is a small, nearly Gaussian move (see A4), which is why it is easy to learn.
A1 · Background

Why generative models are hard for images

The goal of an (unconditional) generative model is a process for sampling from a distribution. For images that is hard: we cannot write down the probability that a given set of pixel values is a photo. A random point in pixel space is static, and real images occupy a thin, curved region of that space.

What we do have is lots of samples of real images. Diffusion uses them in a roundabout way: it learns to walk from random noise to the image region in many small steps.

FigureRandom images vs. real images (lecture slides 6, 16, 20, 26)
Training (slide 20) takes a real image and pushes it into the noise ball. Inference (slides 16, 26) starts at a random point in the ball and walks back to the region of real images, one small denoising step at a time.

Where diffusion sits among generative models

GANadversarial training
\(z\)→Generator \(G(z)\)→\(x^\prime\)→Discriminator \(D(x)\)→\(0/1\)
VAEmaximise the variational lower bound
\(x\)→Encoder \(q_\phi(z\mid x)\)→\(z\)→Decoder \(p_\theta(x\mid z)\)→\(x^\prime\)
Flowinvertible transform
\(x\)→\(f(x)\)→\(z\)→\(f^{-1}(z)\)→\(x^\prime\)
Diffusionadd Gaussian noise gradually, then learn to reverse
\(x_0\)⇄\(x_1\)⇄\(x_2\)⇄ … ⇄\(x_T = z\)
Lecture slide 8, adapted from Weng18. GAN5, VAE4, flow11 and diffusion10. All four map images to a latent \(z\) you can sample. In diffusion the latent has the same size as the image and the map is a long chain of small noising steps.
A2 · The forward process

Diffusion destroys structure, one small step at a time

Diffusion models the destruction of a structured distribution as it moves towards Gaussian noise. Each sample takes a random walk. Every step is a mixture of the current sample and fresh noise:

$$x_t = \underbrace{\sqrt{1-\beta_t}\;x_{t-1}}_{\mu} + \underbrace{\sqrt{\beta_t}}_{\sigma}\;\ce{\epsilon}, \qquad \ce{\epsilon}\sim\mathcal N(0, 1)$$

The \(\beta_t\) values are small (the worksheet uses \(10^{-4}\) up to \(0.02\)), so each step blurs the sample only slightly. Shrinking by \(\sqrt{1-\beta_t}\) while adding variance \(\beta_t\) keeps the overall variance at about 1, so samples never explode. After enough steps every sample looks like a draw from \(\mathcal N(0, I)\).

InteractiveA 2-D spiral diffusing into a Gaussian (lecture slide 10)
Each dot is one sample \(x_t\). It uses the worksheet's schedule, 1000 steps of the update above. Colour shifts from data to noise as \(\sqrt{1-\bar\alpha_t}\) grows. The traced path is a single random walk.

Images live in a much higher-dimensional space (MNIST: 784 dimensions), but the same update applies to every pixel independently. We can watch it on one sample, as on slide 11. The Notebook §2 section below does exactly this with a real digit.

A3 · The closed form

Many noising steps are one big noising step

Because a sum of independent Gaussians is Gaussian, we can collapse the chain. This is lecture slide 13; step through it.

$$x_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\ce{\epsilon_t}$$
substitute \(x_{t-1} = \sqrt{1-\beta_{t-1}}\,x_{t-2} + \sqrt{\beta_{t-1}}\,\epsilon_{t-1}\)
$$x_t = \sqrt{1-\beta_t}\sqrt{1-\beta_{t-1}}\,x_{t-2} + \sqrt{1-\beta_t}\sqrt{\beta_{t-1}}\,\ce{\epsilon_{t-1}} + \sqrt{\beta_t}\,\ce{\epsilon_t}$$
two independent zero-mean Gaussians add: their variances add, \((1-\beta_t)\beta_{t-1} + \beta_t\)
$$x_t = \sqrt{1-\beta_t}\sqrt{1-\beta_{t-1}}\,x_{t-2} + \sqrt{\beta_t + \beta_{t-1} - \beta_t\beta_{t-1}}\;\ce{\epsilon}$$
tidy up both square roots
$$x_t = \sqrt{(1-\beta_t)(1-\beta_{t-1})}\,x_{t-2} + \sqrt{1-(1-\beta_t)(1-\beta_{t-1})}\;\ce{\epsilon}$$
repeat all the way down to \(x_0\)
$$\boxed{\,x_t = \sqrt{\bar\alpha_t}\;\cx{x_0} + \sqrt{1-\bar\alpha_t}\;\ce{\epsilon}\,}\qquad \bar\alpha_t = \prod_{i=1}^{t}(1-\beta_i)$$

The two coefficients always satisfy \((\sqrt{\bar\alpha_t})^2 + (\sqrt{1-\bar\alpha_t})^2 = 1\). The pair is a point on the unit circle that slides from "all image" to "all noise" as \(t\) grows. That is the same rule you find in Exercise 1.1.

InteractiveHow much image is left at step t?
\(\sqrt{\bar\alpha_t}\) image\(\sqrt{1-\bar\alpha_t}\) noise
\(x_0\)
\(\epsilon\)
\(x_t\)
One draw of \(\epsilon\), one formula, any \(t\). There is no need to simulate the 1000 intermediate steps. That is why training is cheap: pick a random \(t\) and jump straight there.
A4 · Reversing it

Why denoise in many small steps?

Erasing an image with noise is easy. Reversing a big jump is very hard, because many clean images could have produced the same static. If the added noise is very small, noising and denoising are about equally easy (slide 15).

We can write the forward step as a distribution, \(P(x_t\mid x_{t-1}) = \mathcal N\big(\sqrt{1-\beta_t}\,x_{t-1},\ \beta_t\big)\). The generative process has to model the reverse, \(P(x_{t-1}\mid x_t)\). In general this is not Gaussian (slide 21). The demo shows why small steps rescue us: for a tiny step the reverse distribution is almost exactly Gaussian, so a network only has to predict its mean.

InteractiveThe exact reverse distribution for 1-D data with three clusters
\(x_{t-\Delta}\)
true \(P(x_{t-\Delta}\mid x_t)\)best Gaussian fit\(P(x_{t-\Delta})\), rescaled
The solid curve is the true \(P(x_{t-\Delta}\mid x_t)\), computed exactly from a 3-cluster data distribution. The dashed curve is the best Gaussian with the same mean and variance. With \(\Delta = 1\) they coincide. Drag \(\Delta\) up and the reverse step becomes multi-modal, so a single Gaussian (one prediction) cannot capture it.
A5 · Using deep learning

A U-Net that predicts the clean image or the noise

A U-Net12 takes the noisy image \(x_t\) and the timestep \(t\) and outputs an estimate. There are two equivalent choices, linked by the closed form \(x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon\):

The worksheet network has both heads. It is trained on both targets, and only the noise head is used to generate. Intuitively, the model learns the direction in which to move a noisy image to reduce its noise.

A6 · Training and inference

The two loops

Training (slides 19–25)

  1. Take a real image \(\cx{x_0}\) from the dataset.
  2. Draw \(t \sim \mathrm{Uniform}\{0,\dots,T-1\}\) and \(\ce{\epsilon}\sim\mathcal N(0,I)\).
  3. Make \(x_t = \sqrt{\bar\alpha_t}\,\cx{x_0} + \sqrt{1-\bar\alpha_t}\,\ce{\epsilon}\) in one shot.
  4. Predict \(\cm{\hat\epsilon}, \cm{\hat x_0} = \text{UNet}(x_t, t)\).
  5. Loss \(= \lVert\cm{\hat\epsilon}-\ce{\epsilon}\rVert^2\), optionally plus \(\lVert\cm{\hat x_0}-\cx{x_0}\rVert^2\). Take a gradient step.

Inference (slides 26–27)

  1. Start from pure noise \(x_T \sim \mathcal N(0, I)\).
  2. For \(t = T-1,\dots,0\): predict \(\cm{\hat\epsilon} = \text{UNet}(x_t, t)\).
  3. Step to a slightly less noisy \(x_{t-1}\) using \(\cm{\hat\epsilon}\).
  4. Either add a little fresh noise (stochastic sampling) or don't (deterministic; the lecture calls this cold diffusion14).
  5. Good results typically take about 1000 steps, which is why diffusion is slow.

Two ways to take the reverse step

Lecture (slide 27). Turn \(\hat\epsilon\) into \(\hat x_0\), then re-noise it to the level of step \(t-1\). This is the deterministic DDIM update13:

$$\cm{\hat x_0} = \frac{x_t - \sqrt{1-\bar\alpha_t}\,\cm{\hat\epsilon}}{\sqrt{\bar\alpha_t}}, \qquad \hat x_{t-1} = \sqrt{\bar\alpha_{t-1}}\;\cm{\hat x_0} + \sqrt{1-\bar\alpha_{t-1}}\;\cm{\hat\epsilon}$$

Worksheet (cell 32, the DDPM update10). Subtract a small amount of the predicted noise, scale up, then add a little noise back:

$$x_{t-1} = \frac{1}{\sqrt{1-\beta_t}}\Big(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\cm{\hat\epsilon}\Big) + \sqrt{\beta_t}\,\ce{z}, \qquad \ce{z}\sim\mathcal N(0,I),\ \ \text{no } z \text{ at } t=0$$

The first form reuses \(\hat\epsilon\) deterministically and lets you skip steps. The second is the classic stochastic sampler. You can run both on the trained network in §5.

A7 · Extensions and applications

Latent, conditional and inpainting diffusion

Diffusion on a latent space (Stable Diffusion15)

image \(x\)
\(512{\times}512{\times}3\approx786\text{k}\)
\(E\)
\(z\)
\(64{\times}64{\times}4\approx16\text{k}\)
\(D\)
diffusion runs here, on \(z\)
\(z_T\to\cdots\to z_0\) with a small U-Net
then decode: \(x = D(z_0)\)
Step 1: train the autoencoder
Step 2: use \(z\) in place of \(x_0\)
Lecture slide 30. The autoencoder \(E, D\) is trained first. Diffusion then runs on \(z\) in place of \(x_0\). The latent is about \(48\times\) smaller here, so the diffusion network can be smaller and faster, and the decoder handles most of the domain's regularities (textures, edges) so the diffusion model doesn't have to.

Conditional diffusion

Train with an extra input alongside the noised image (slides 31–32):

  • Text or label. A one-hot category vector, or a text embedding such as CLIP's16 for “lighthouse on a rocky coast”. In our worksheet U-Net it would be concatenated next to \(t\) before the dense layer.
  • Another image. For segmentation or style transfer, feed the RGB image plus one or more channels of noise, and diffuse only the output channels.
  • Inpainting. Input the image and a mask, and diffuse only in the masked region. One way, RePaint17 (slides 33–34), needs no retraining at all. At every step it replaces the known pixels with the original image noised to the current level. The mask keeps the denoised pixels only where the hole is. The Inpainting lab runs this on the worksheet model.
Slide 33One RePaint iteration
known partsample forward \(q\)
input \(x_0\)→ noise →\(x_{t-1}\sim q\)\(\odot\)\(1-m\)
unknown partsample reverse \(p_\theta\)
\(x_t\)→ denoise →\(x_{t-1}\sim p_\theta\)\(\odot\)\(m\)
combine
\(x_{t-1} = (1-m)\odot x^{\text{known}}_{t-1} + m\odot x^{\text{unknown}}_{t-1}\)→ next iteration

Pros and cons

AdvantagesDisadvantages
High-quality, diverse image outputVery slow at inference: about 1000 network evaluations per sample
Easier and more stable to train than GANs (plain MSE regression, no adversary)Pixel-space models are expensive at high resolution (hence latent diffusion)

Worksheet 10 · notebook walkthrough

0 · Setup

The worksheet runs these demonstrations in order: properties of Gaussian noise, diffusion as a destructive process, training a noise/image estimator, evaluating it on noisy images, and generating images with the reverse process. Code cells below are copied verbatim from worksheet10-diffusion-solution new.ipynb. The cell numbers are the notebook's own cell indices.

In [2]Imports
# Includes importing some of custom packages
# The functions in the custom package are called via myutils.xxx
# e.g. the xx function needs to be called via myutils.xx 

# Imports
import os

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import Model

import numpy as np
import matplotlib.pyplot as plt

print("TensorFlow version: ", tf.__version__)
About the live demos. The browser can't run TensorFlow on the full training set in a workshop slot. So I trained a line-for-line PyTorch port of the cell 21 U-Net (same layers, 2,520,154 parameters, same schedule, loss, Adam optimiser, batch 100, 20 epochs) and the page runs that network's forward pass in plain JavaScript. The dense layer is stored as 8-bit weights to keep the page small. Results will look like the notebook's, not identical.

1 · Properties of Gaussian noise

The forward (destructive) process repeatedly mixes images with Gaussian (normally distributed) noise. First, look at how independent normal variables add up to give normally distributed results.

In [4]A histogram helper
#define a histogram function for showing gaussian variables

def histshow(x):
    # display the values as a histogram
    numbins = 11
    histbins = np.linspace(-5.5,5.5,numbins+1)
    bincentres = (histbins[0:numbins] + histbins[1:numbins+1])/2

    plt.hist(x, bins=histbins, align='left')
    plt.xticks(ticks=bincentres)
    plt.title('Distribution of normally distributed random variables')
    plt.ylabel('Frequency')
    plt.xlabel('Value')
    plt.show()
In [5]10 000 standard normals
# make a tensor filled with normally distributed variables
N = tf.random.normal(shape=[10000], mean=0.0, stddev=1.0)

histshow(N)
1.1

Change the values of a and b so that Nresult has the same (or a very similar) histogram to N above.

  • What should the mean and standard deviation of Nresult be?
  • What rule governs a and b? What should b be if a = 0.5, 0.6 or 0.8?
Show solution

Mean 0 and standard deviation 1, the same as N. For independent \(N_1, N_2\):

$$\operatorname{Var}(aN_1 + bN_2) = a^2\operatorname{Var}(N_1) + b^2\operatorname{Var}(N_2) = a^2 + b^2$$

So we need \(a^2 + b^2 = 1\). That gives \(a=0.5\Rightarrow b\approx0.866\), \(a=0.6\Rightarrow b=0.8\), and \(a=0.8\Rightarrow b=0.6\). This is exactly why the diffusion coefficients are \(\sqrt{1-\beta_t}\) and \(\sqrt{\beta_t}\).

In [7]Mixing two Gaussians
# make wo tensors filled with normally distributed variables
N1 = tf.random.normal(shape=[10000], mean=0.0, stddev=1.0)
N2 = tf.random.normal(shape=[10000], mean=0.0, stddev=1.0)

# YOUR CODE SHOULD MODIFY a AND b
a = 0.5 # or 0.6 or 0.8
b = np.sqrt(1 - a * a)

Nresult = a * N1 + b * N2
histshow(Nresult)

m = np.mean(Nresult)
s = np.std(Nresult)
print(m,s)
Livehistshow(Nresult) with Nresult = a·N1 + b·N2
\(aN_1+bN_2\)expected counts for \(\mathcal N(0,1)\)bell curve \(\mathcal N(0,1)\)bell curve of the blue bars, \(\mathcal N(\mu,\sigma^2)\)
\((a, b)\) with the circle \(a^2+b^2=1\)
Bars show Nresult in the same 11 bins as histshow. The dashed outline is what N (a standard normal) would give, and the plain \(\mathcal N(0,1)\) bell curve is its density. The blue bell curve belongs to the blue bars: \(\mathcal N(\mu,\sigma^2)\) using Nresult's measured mean and std, where theory says \(\sigma=\sqrt{a^2+b^2}\). The two curves coincide only when \(a^2+b^2=1\). It starts at the worksheet's a = 0, b = 0, which puts everything in the 0 bin. The right panel shows \((a, b)\); the histogram matches when the point sits on the unit circle.

2 · Diffusion as a destructive process

First load the MNIST data set. Pixels are scaled from \(\{0,\dots,255\}\) to \([-1, 1)\), which matches the scale of the noise.

In [10]Load and rescale MNIST
(train_images, train_labels), (test_images, test_labels) = keras.datasets.mnist.load_data()
print("train_images shape:", train_images.shape)
print("test_images shape:", test_images.shape)
# Data type adjustment & rescale 
train_images = (train_images.astype(float) / 128) -1.0 #scale the pixel values to be -1 to 1
test_images = (test_images.astype(float) / 128) - 1.0
train_images = np.expand_dims(train_images, axis = -1)
test_images = np.expand_dims(test_images, axis = -1)

Pick a random image and display it. The digit you choose here is used by every live demo further down.

In [12]Show a random digit
sample_id = np.random.choice(train_images.shape[0])
digit = train_images[sample_id]

plt.imshow(digit.squeeze(), cmap="binary")
plt.title("$label = {}$".format(train_labels[sample_id]))
plt.axis('off')
plt.show()
LiveChoose the working digit (40 images from the MNIST test set)
label = ?
or click any thumbnail
Rendered like plt.imshow(..., cmap="binary"): high values are dark. Matplotlib also rescales every image to its own min and max, which is why noisy images look grey rather than saturated.

2.1 Compute the \(\alpha\) and \(\beta\) coefficients

In the original diffusion paper10, images are destroyed over 1000 time steps. The amount of noise added at each step is called \(\beta\). It increases steadily from \(0.0001\) to \(0.02\) over the 1000 steps. The value of \(\bar\alpha\) at each step is the square of the amount of image remaining:

$$\bar\alpha_t = \prod_{t'=0}^{t}(1-\beta_{t'})$$
Typo in the notebook text. The markdown says \(\beta\) rises “to a value of 0.2”, but the code (and the DDPM paper) use \(\mathbf{0.02}\). Also note the code calls this array alpha, without the bar, and indexes from 0, so alpha[t] already includes step \(t\).
In [14]Schedules
#number of time steps
T=1000

# set the beta schedule linearly increasing from 10^-4 to 0.02
beta = np.linspace(0.0001,0.02,T, dtype='float')

# obtain the alpha schedule
alpha = np.cumprod(1.0-beta)

# create arrays with the square roots we will need
sqrtbeta = np.float32(np.sqrt(beta))
sqrt1mbeta = np.float32(np.sqrt(1.0-beta))
sqrtalpha = np.float32(np.sqrt(alpha))
sqrt1malpha = np.float32(np.sqrt(1.0 - alpha))
LiveThe four arrays from cell 14
\(t\)
sqrtalpha \(=\sqrt{\bar\alpha_t}\)sqrt1malpha \(=\sqrt{1-\bar\alpha_t}\)alpha \(=\bar\alpha_t\)beta \(=\beta_t\ (\times 25)\)
Hover the chart to read values.
\(\sqrt{\bar\alpha_t}\) (image kept) and \(\sqrt{1-\bar\alpha_t}\) (noise) cross near \(t\approx255\). By \(t=999\) only \(0.6\%\) of the signal amplitude remains. \(\beta_t\) is drawn \(\times25\) so it is visible.

2.2 Use the \(\alpha\) and \(\beta\) coefficients to create noisy images

We should be able to get noisy images in two ways.

Method 1: repeatedly add noise according to the \(\beta\) schedule, starting each time from the previous noisy image \(x_{t-1}\):

$$x_t = \sqrt{1-\beta_t}\;x_{t-1} + \sqrt{\beta_t}\;\epsilon \qquad \epsilon\sim\mathcal N[0,1]$$

Method 2: add the equivalent amount of noise in a single step, directly from the original image \(x_0\):

$$x_t = \sqrt{\bar\alpha_t}\;x_0 + \sqrt{1-\bar\alpha_t}\;\epsilon \qquad \epsilon\sim\mathcal N[0,1]$$
In [16]Method 1: 1000 small steps
# method 1 (slowly add noise over 1000 timesteps)
# create a list of tensors with all 1000 timesteps
noisy_images1 = [digit * sqrt1mbeta[0] + np.random.normal(size = digit.shape, loc=0.0, scale=sqrtbeta[0])]
for t in range(1,T):
    noisy_images1.append (noisy_images1[t-1] * sqrt1mbeta[t] + np.random.normal(size = digit.shape, loc=0.0, scale=sqrtbeta[t]))


#we're going to sample 10 different timesteps at 0, 100, 200, ..., 800, 900
num_images = 10
timesteps = range(0, T, T//(num_images))

fig, axes = plt.subplots(figsize=(1.5*num_images, 1.5), ncols=num_images)
for i, s in enumerate(timesteps):
    axes[i].imshow(noisy_images1[s], cmap='binary')
    axes[i].set_title("$t = {}$".format(s))
    axes[i].axis('off')
plt.show()
In [17]Method 2: jump straight to t
# method 2 (directly generate noisy images at the timesteps we want)
noisy_images2 = [digit * sqrtalpha[t] + np.random.normal(size = digit.shape, loc=0.0, scale = sqrt1malpha[t]) for t in timesteps]

fig, axes = plt.subplots(figsize=(1.5*num_images, 1.5), ncols=num_images)
for i in range(num_images):
    axes[i].imshow(noisy_images2[i], cmap='binary')
    axes[i].set_title("$t = {}$".format(timesteps[i]))
    axes[i].axis('off')
plt.show()
LiveBoth methods on your digit
Method 1: noisy_images1[t], iterated
Method 2: noisy_images2, closed form
2.3

Compare the two methods. Do they generate images of similar quality? Which method is faster?

Show solution

Both methods look similar. They produce samples from exactly the same distribution, \(\mathcal N(\sqrt{\bar\alpha_t}\,x_0,\,(1-\bar\alpha_t)I)\), just with different random draws. Method 2 is faster: one multiply-add per requested \(t\) instead of \(t\) sequential steps. The timing above shows the gap. This is what makes training practical, because each training example needs a noisy image at a random \(t\).

3 · Training a noise / image estimator

3.1 Define the U-Net

The network takes a noisy image and a timestep, and estimates both the noise and the image. Hover a block in the diagram to highlight its line of code.

Shapes are \((\text{height}\times\text{width}\times\text{channels})\). The time input enters only once, as one extra number appended to the 1568-dim bottleneck before the dense layer. That dense layer holds 2,460,384 of the 2,520,154 parameters (97.6%).
In [21]The U-Net
noisy_input = keras.Input(shape=digit.shape)
time_input = keras.Input(shape=(1,))

# keep separate names for these ones because we will use them on the way up the u-net
x0 = layers.Conv2D(filters=16, kernel_size=[4,4], strides=[2,2], padding='same', activation='relu')(noisy_input)
x1 = layers.Conv2D(filters=32, kernel_size=[4,4], strides=[2,2], padding='same', activation='relu')(x0)
x2 = layers.Flatten()(x1)

t = time_input * 0.002 - 1 # turn the time into a number between -1 and 1

x = layers.Concatenate()([x2,t])

x = layers.Dense(units=x2.shape[-1])(x)

x = layers.Reshape(target_shape = (x1.shape[1],x1.shape[2],x1.shape[3]))(x)
x = layers.Concatenate()([x,x1])
x = layers.Conv2D(filters=64, kernel_size=[3,3], padding='same', activation='relu')(x)
x = layers.Conv2DTranspose(filters=4, kernel_size=[4,4], strides=[2,2], padding='same', activation='relu')(x)
x = layers.Concatenate()([x,x0])
x = layers.Conv2D(filters=32, kernel_size=[3,3], padding='same', activation='relu')(x)
x = layers.Conv2DTranspose(filters=4, kernel_size=[4,4], strides=[2,2], padding='same', activation='relu')(x)
x = layers.Concatenate()([x,noisy_input])
x = layers.Conv2D(filters=16, kernel_size=[3,3], padding='same', activation='relu')(x)
output1 = layers.Conv2D(filters=1, kernel_size=[3,3], strides=[1,1], padding='same')(x) # noise estimate
output2 = layers.Conv2D(filters=1, kernel_size=[3,3], strides=[1,1], padding='same')(x) # image estimate

UNet = keras.Model(inputs=[noisy_input, time_input], outputs=[output1,output2])
LivePeek inside: activations for your digit
noisy_input\(28\times28\times1\)
x0 = Conv2D(16, stride 2) → \(14\times14\times16\)
x1 = Conv2D(32, stride 2) → \(7\times7\times32\)
final Conv2D(16) features → \(28\times28\times16\)
output1noise estimate \(\hat\epsilon\)
output2image estimate \(\hat x_0\)

3.2 Custom layers that make the training data on the fly

These layers generate noise and timesteps, then use them to create noisy images, all inside the Keras graph. Each batch gets fresh random \(t\) and \(\epsilon\) for free.

In [23]SampleNoise, SampleT, MakeNoisyImages
class SampleNoise(layers.Layer):
    def __init__(self):
        super().__init__()

    def compute_output_shape(self, input_shape):
        return input_shape

    # Defines the computation
    def call(self, input1):
        return tf.random.normal(shape = input1.shape, mean=0.0, stddev=1.0)

class SampleT(layers.Layer):
    def __init__(self):
        super().__init__()

    def compute_output_shape(self, input_shape):
        return (input_shape[0],1)

    # Defines the computation
    def call(self, input1):
        batch_size = input1.shape[0]
        return tf.random.uniform(minval=0, maxval=T, shape=(batch_size,1), dtype=tf.int64)

class MakeNoisyImages(layers.Layer):
    def __init__(self):
        super().__init__()

    def compute_output_shape(self, input_shape):
        return input_shape

    # Defines the computation
    def call(self, images, noise, timesteps):
        sqrtalphas = tf.gather(sqrtalpha, timesteps)
        sqrt1malphas = tf.gather(sqrt1malpha, timesteps)
        sqrtalphas = tf.reshape(sqrtalphas,[timesteps.shape[0],1,1,1])
        sqrt1malphas = tf.reshape(sqrt1malphas, [timesteps.shape[0],1,1,1])
        return sqrtalphas * images + sqrt1malphas * noise
LayerDoesMaths
SampleNoisenoise the same shape as the batch\(\ce{\epsilon}\sim\mathcal N(0, I)\)
SampleTone integer timestep per image\(t\sim\mathrm{Uniform}\{0,\dots,999\}\)
MakeNoisyImagesclosed-form noising (method 2), with per-image coefficients via tf.gather\(x_t=\sqrt{\bar\alpha_t}\,\cx{x_0}+\sqrt{1-\bar\alpha_t}\,\ce{\epsilon}\)

3.3 The wrapper that prepares data for the U-Net

In [25]Wrapped model
# wrap the UNet in a model that randomly chooses time, creates noise, adds it and gets the outputs
image_input = keras.Input(shape = digit.shape)
epsilon = SampleNoise()(image_input)
timesteps = SampleT()(image_input)
noised_images = MakeNoisyImages()(image_input, epsilon, timesteps)

timesteps = keras.ops.cast(timesteps, dtype='float32')

noise_estimate, image_estimate = UNet([noised_images, timesteps])

# we're training on both image and noise but only using the noise for inference,
# so we make its loss matter more by multiplying the residual error by a factor of two
noise_estimate = 2 * (noise_estimate - epsilon)

Wrapped = keras.Model(inputs = image_input, outputs = [noise_estimate, image_estimate])

The “subtract inside the model” trick. Keras losses compare an output to a target given to fit. But the true noise \(\epsilon\) is invented inside the model, so fit never sees it. The wrapper therefore outputs the residual \(2(\hat\epsilon-\epsilon)\) and trains it towards a target of zeros (imzeros in cell 27). The total loss is

$$\mathcal L = \underbrace{\operatorname{MSE}\big(2(\cm{\hat\epsilon}-\ce{\epsilon}),\,0\big)}_{=\,4\,\lVert\hat\epsilon-\epsilon\rVert^2/n} + \underbrace{\operatorname{MSE}\big(\cm{\hat x_0},\,\cx{x_0}\big)}_{\text{image head}}$$

The factor of two becomes a weight of four on the noise loss, because that is the head used at inference.

LiveOne training example, as the Wrapped model sees it
Uses your digit as image_input, then SampleT, SampleNoise and MakeNoisyImages, then a forward pass of the trained U-Net. The two losses are the per-example versions of the two Keras outputs.

3.4 Compile and train

About 30 seconds per epoch on an M1 MacBook Air, so roughly 11 minutes for 20 epochs.

In [27]compile + fit
# make a set of zero images as the target for the noise estimate with the true noise subtracted
imzeros = tf.zeros(train_images.shape)

Wrapped.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError(), metrics=['mse','mse'])
history_Wrapped = Wrapped.fit(train_images, [imzeros, train_images], epochs=20, batch_size=100, validation_split=0)
Spot the difference. The worksheet passes loss=tf.keras.losses.MeanSquaredError (the class). The solution passes MeanSquaredError() (an instance). Keras 3 needs a callable loss object, so the worksheet version errors.
MeasuredTraining curves from the port used on this page
noise term \(4\,\mathrm{MSE}(\hat\epsilon,\epsilon)\)image term \(\mathrm{MSE}(\hat x_0,x_0)\)
Epoch-mean losses over 20 epochs (600 steps each). Noise term: Keras' multiply_2_loss, i.e. \(4\cdot\)MSE. Image term: functional_loss. For comparison, the notebook's own log shows 1.297 → 0.392 → 0.315 (noise) and 0.321 → 0.184 → 0.176 (image) for epochs 1–3.

4 · Evaluating on noisy images

Feed the images from section 2.2 into the U-Net and look at the results. Turn the code into a loop that shows timesteps 0, 100, …, 900 as 10 separate columns.

In [29]Noisy input · noise estimate · image estimate
timestep = 100
fig, axes = plt.subplots(figsize=(3*num_images, 10), nrows=3, ncols=10)

for i in range(10):
    noisy_image = noisy_images1[timestep * i]
    noisy_image = np.expand_dims(noisy_image, axis=0)
    noise_estimate, image_estimate = UNet.predict([noisy_image, np.array([[timestep]])])
    axes[0,i].imshow(noisy_image[0], cmap='binary')
    axes[0,i].axis('off')
    axes[1,i].imshow(noise_estimate[0], cmap='binary')
    axes[1,i].axis('off')
    axes[2,i].imshow(image_estimate[0], cmap='binary')
    axes[2,i].axis('off')

plt.show()
A subtle bug worth discussing. The solution loop feeds noisy_images1[timestep * i] but tells the network the time is np.array([[timestep]]), i.e. always 100. The network is told “this image is only lightly noised”, so on heavily noised inputs it under-predicts the noise. The fix is np.array([[timestep * i]]). Toggle both below.
LiveThe \(3\times10\) grid from cell 29
noisy_image
noise_estimate
image_estimate
LiveScrub through t
\(t\)
\(\mathrm{MSE}(\hat\epsilon,\epsilon)\)\(\mathrm{MSE}(\hat x_0,x_0)\), image head\(\mathrm{MSE}(\hat x_0,x_0)\), from \(\hat\epsilon\)
Two ways to get \(\hat x_0\): the image head (output2), or the noise head converted with slide 23's formula \(\hat x_0 = (x_t - \sqrt{1-\bar\alpha_t}\,\hat\epsilon)/\sqrt{\bar\alpha_t}\). The curve sweeps t for this digit. Predicting \(\epsilon\) is hardest at small t, where the noise is a faint layer on top of the digit, and easiest at large t, where \(x_t\) is almost pure noise. Recovering \(x_0\) goes the other way: the image head's error grows with t, and converting \(\hat\epsilon\) to \(\hat x_0\) divides by \(\sqrt{\bar\alpha_t}\), which magnifies any error near \(t=1000\) (dashed curve, clipped).

5 · Generating images with the reverse process

We synthesise an image in 1000 steps, starting from pure noise (first line of code), then iterate timesteps from 999 down to 0. At each step we:

  • subtract a small amount of the noise,
  • scale up the result,
  • add some noise back, because the network's estimate of the noise is inaccurate (but not at the last step).

The code takes a bit more than 30 seconds on an M1 MacBook Air. If it seems slow, set verbose=1 to see per-call progress bars and multiply by 1000.

In [32]The reverse process
noisy_digit = np.random.normal(size = digit.shape, loc=0.0, scale=1.0)
noisy_digit = np.expand_dims(noisy_digit, axis=0)

ndshape = noisy_digit.shape

fig, axes = plt.subplots(figsize=(1.5*num_images, 1.5), ncols=11)
image_no=0
axes[image_no].imshow(noisy_digit[0], cmap='binary')
axes[image_no].set_title("$t = {}$".format(1000))
axes[image_no].axis('off')

for t in range(999,-1,-1):
    tstep = np.array([[t]])
    # noize, im = UNet.predict([noisy_digit, tstep], verbose=0)
    noise_estimate, image_estimate = UNet.predict([noisy_digit, tstep], verbose=0)

    #noisy_digit = (noisy_digit - beta[t] * noise_estimate / np.sqrt(1-alpha[t])) / (np.sqrt(1-beta[t]))
    noisy_digit = (noisy_digit - beta[t] * noise_estimate / sqrt1malpha[t]) / sqrt1mbeta[t]

    if t>0:
        noisy_digit = noisy_digit + np.random.normal(size = ndshape, loc=0.0, scale=sqrtbeta[t])

    if t % 100 == 0:
        image_no = image_no + 1
        axes[image_no].imshow(noisy_digit[0], cmap='binary')
        axes[image_no].set_title("$t = {}$".format(t))
        axes[image_no].axis('off')
plt.show()

The update line is the DDPM step written with the notebook's arrays:

$$\underbrace{x_{t-1}}_{\texttt{noisy\_digit}} = \frac{x_t - \overbrace{\beta_t}^{\texttt{beta[t]}}\,\cm{\hat\epsilon}\,/\,\overbrace{\sqrt{1-\bar\alpha_t}}^{\texttt{sqrt1malpha[t]}}}{\underbrace{\sqrt{1-\beta_t}}_{\texttt{sqrt1mbeta[t]}}} \;+\; \underbrace{\sqrt{\beta_t}}_{\texttt{sqrtbeta[t]}}\,\ce{z}$$
LiveRun cell 32 in the browser
noisy_digit
current guess of \(\hat x_0\)from \(\hat\epsilon\) (slide 23)
snapshots every 100 steps, as plotted by cell 32
5.1

Run this a few times. How often does the result look like a genuine number? (If it's slow, compare your result with others in the workshop.)

Tally it as a class: click a sample once for “digit” (green) and twice for “glyph” (orange). These 16 were generated by the same sampler.

Show solution

The model should make something that looks like a handwritten digit about half the time. The other half it makes something that looks like a written glyph. That's a small, quickly trained model.

5.2 (optional)

Try removing the noise added during generation and run the code at least twice. What happens?

Four workers run in parallel, one per seed.

Shown on a fixed \([-1, 1]\) scale, so a blank image looks blank instead of being stretched to fill the grey range.

Show solution

Without the added noise, the model converges to a single answer whatever the starting noise. With this network the runs end up as the same, almost blank image: the pixel values all sit near \(-1\) and differ between seeds by less than a millionth. Removing the noise term every step makes the update pull towards the average image. The injected noise keeps the samples diverse, so the walk explores the data distribution instead of collapsing onto its average. (The lecture's deterministic slide 27 update avoids this collapse because it re-uses \(\hat\epsilon\) rather than dropping the noise term entirely.)

Bonus · extension from slides 33–34

Inpainting lab

Paint over part of your digit (the painted pixels get regenerated), then run RePaint with the unconditional worksheet U-Net. At every step the unpainted pixels are overwritten with the real digit, noised to the current level. The model only fills the hole, and has to agree with its surroundings.

paint the holedrag to paint
\(x_t\)
resulthole filled by the model
The pixel-space RePaint loop is the slide 33 diagram: known part \(\sim q\) (forward noise), unknown part \(\sim p_\theta\) (one cell 32 step), then combine with the mask. Run it twice to see different plausible completions.
Summary

Cheat sheet

IdeaFormula / factIn the notebook
Forward step\(x_t=\sqrt{1-\beta_t}\,x_{t-1}+\sqrt{\beta_t}\,\epsilon\)cell 16 (method 1)
Variance preserved\(a^2+b^2=1\)Exercise 1.1
Closed form\(x_t=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\epsilon,\ \ \bar\alpha_t=\prod(1-\beta_i)\)cell 14 alpha; cell 17; MakeNoisyImages
Schedule\(\beta\): linear \(10^{-4}\to0.02\), \(T=1000\)cell 14
Network\(\text{U-Net}(x_t, t)\to(\hat\epsilon,\hat x_0)\), with \(t\) scaled to \([-1,1]\)cell 21
Convert\(\hat x_0=(x_t-\sqrt{1-\bar\alpha_t}\,\hat\epsilon)/\sqrt{\bar\alpha_t}\)slide 23; §4 demo
Loss\(4\,\mathrm{MSE}(\hat\epsilon,\epsilon)+\mathrm{MSE}(\hat x_0,x_0)\)cells 25, 27
Sampling step\(x_{t-1}=\big(x_t-\tfrac{\beta_t}{\sqrt{1-\bar\alpha_t}}\hat\epsilon\big)/\sqrt{1-\beta_t}+\sqrt{\beta_t}\,z\)cell 32
Latent diffusiondiffuse \(z=E(x)\), decode with \(D\)Stable Diffusion
Conditioningextra input: one-hot, CLIP text embedding, image, maskslides 31–34

Built for COMP90086 Workshop 10 from the Week 9 diffusion lecture (Ehinger & Drummond) and worksheet 10.

Bibliography

References

  1. J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, L. Fei-Fei. ImageNet: A large-scale hierarchical image database. CVPR 2009. doi.org/10.1109/CVPR.2009.5206848
  2. A. Krizhevsky, I. Sutskever, G. E. Hinton. ImageNet classification with deep convolutional neural networks. NeurIPS 2012. papers.nips.cc/paper/2012/hash/c399862d3b9d6b76c8436e924a68c45b-Abstract.html
  3. G. E. Hinton, R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. Science 313(5786), 2006. doi.org/10.1126/science.1127647
  4. D. P. Kingma, M. Welling. Auto-encoding variational Bayes. ICLR 2014. arXiv:1312.6114. arxiv.org/abs/1312.6114
  5. I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, Y. Bengio. Generative adversarial nets. NeurIPS 2014. arXiv:1406.2661. arxiv.org/abs/1406.2661
  6. T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, X. Chen. Improved techniques for training GANs. NeurIPS 2016. arXiv:1606.03498. arxiv.org/abs/1606.03498
  7. P. Vincent, H. Larochelle, Y. Bengio, P.-A. Manzagol. Extracting and composing robust features with denoising autoencoders. ICML 2008. doi.org/10.1145/1390156.1390294
  8. J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. ICML 2015. arXiv:1503.03585. arxiv.org/abs/1503.03585
  9. Y. Song, S. Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS 2019. arXiv:1907.05600. arxiv.org/abs/1907.05600
  10. J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. NeurIPS 2020. arXiv:2006.11239. arxiv.org/abs/2006.11239
  11. L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR 2017. arXiv:1605.08803. arxiv.org/abs/1605.08803
  12. O. Ronneberger, P. Fischer, T. Brox. U-Net: Convolutional networks for biomedical image segmentation. MICCAI 2015. arXiv:1505.04597. arxiv.org/abs/1505.04597
  13. J. Song, C. Meng, S. Ermon. Denoising diffusion implicit models. ICLR 2021. arXiv:2010.02502. arxiv.org/abs/2010.02502
  14. A. Bansal, E. Borgnia, H.-M. Chu, J. S. Li, H. Kazemi, F. Huang, M. Goldblum, J. Geiping, T. Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise. NeurIPS 2023. arXiv:2208.09392. arxiv.org/abs/2208.09392
  15. R. Rombach, A. Blattmann, D. Lorenz, P. Esser, B. Ommer. High-resolution image synthesis with latent diffusion models. CVPR 2022. arXiv:2112.10752. arxiv.org/abs/2112.10752
  16. A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, I. Sutskever. Learning transferable visual models from natural language supervision. ICML 2021. arXiv:2103.00020. arxiv.org/abs/2103.00020
  17. A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, L. Van Gool. RePaint: Inpainting using denoising diffusion probabilistic models. CVPR 2022. arXiv:2201.09865. arxiv.org/abs/2201.09865
  18. L. Weng. What are diffusion models? Lil'Log, 2021. lilianweng.github.io/posts/2021-07-11-diffusion-models/