# 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__)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.
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.
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:
LimitThe output is a label \(y\). The model never learns what images look like, so it cannot produce a new \(x\).
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:
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.
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):
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.
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:
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.
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:
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.
Where diffusion sits among generative models
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)\).
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.
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.
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.
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.
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.
The two loops
Training (slides 19–25)
- Take a real image \(\cx{x_0}\) from the dataset.
- Draw \(t \sim \mathrm{Uniform}\{0,\dots,T-1\}\) and \(\ce{\epsilon}\sim\mathcal N(0,I)\).
- Make \(x_t = \sqrt{\bar\alpha_t}\,\cx{x_0} + \sqrt{1-\bar\alpha_t}\,\ce{\epsilon}\) in one shot.
- Predict \(\cm{\hat\epsilon}, \cm{\hat x_0} = \text{UNet}(x_t, t)\).
- 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)
- Start from pure noise \(x_T \sim \mathcal N(0, I)\).
- For \(t = T-1,\dots,0\): predict \(\cm{\hat\epsilon} = \text{UNet}(x_t, t)\).
- Step to a slightly less noisy \(x_{t-1}\) using \(\cm{\hat\epsilon}\).
- Either add a little fresh noise (stochastic sampling) or don't (deterministic; the lecture calls this cold diffusion14).
- 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.
Latent, conditional and inpainting diffusion
Diffusion on a latent space (Stable Diffusion15)
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.
Pros and cons
| Advantages | Disadvantages |
|---|---|
| High-quality, diverse image output | Very 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) |
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.
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.
#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()# make a tensor filled with normally distributed variables
N = tf.random.normal(shape=[10000], mean=0.0, stddev=1.0)
histshow(N)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
Nresultbe? - What rule governs
aandb? What shouldbbe ifa= 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\):
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}\).
# 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)
# 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
b = 0
Nresult = a * N1 + b * N2
histshow(Nresult)
m = np.mean(Nresult)
s = np.std(Nresult)
print(m,s)
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.
(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.
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()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'})$$alpha, without the bar, and indexes from 0, so alpha[t] already includes step \(t\).#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))sqrtalpha \(=\sqrt{\bar\alpha_t}\)sqrt1malpha \(=\sqrt{1-\bar\alpha_t}\)alpha \(=\bar\alpha_t\)beta \(=\beta_t\ (\times 25)\)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]$$# 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()# 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()noisy_images1[t], iteratednoisy_images2, closed formCompare 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.
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])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.
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| Layer | Does | Maths |
|---|---|---|
SampleNoise | noise the same shape as the batch | \(\ce{\epsilon}\sim\mathcal N(0, I)\) |
SampleT | one integer timestep per image | \(t\sim\mathrm{Uniform}\{0,\dots,999\}\) |
MakeNoisyImages | closed-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
# 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
The factor of two becomes a weight of four on the noise loss, because that is the head used at inference.
3.4 Compile and train
About 30 seconds per epoch on an M1 MacBook Air, so roughly 11 minutes for 20 epochs.
# 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)
# 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)
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.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.
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()
timestep = 100
noisy_image = noisy_images1[timestep]
noisy_image = np.expand_dims(noisy_image, axis=0)
noise_estimate, image_estimate = UNet.predict([noisy_image, np.array([[timestep]])])
fig, axes = plt.subplots(figsize=(3*num_images, 3), nrows=3, ncols=1)
image_no=0
axes[0].imshow(noisy_image[0], cmap='binary')
axes[0].axis('off')
axes[1].imshow(noise_estimate[0], cmap='binary')
axes[1].axis('off')
axes[2].imshow(image_estimate[0], cmap='binary')
axes[2].axis('off')
plt.show()
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.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.
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}$$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.
Try removing the noise added during generation and run the code at least twice. What happens?
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.)
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.
Cheat sheet
| Idea | Formula / fact | In 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 diffusion | diffuse \(z=E(x)\), decode with \(D\) | Stable Diffusion |
| Conditioning | extra input: one-hot, CLIP text embedding, image, mask | slides 31–34 |
Built for COMP90086 Workshop 10 from the Week 9 diffusion lecture (Ehinger & Drummond) and worksheet 10.
References
- 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
- 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
- G. E. Hinton, R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. Science 313(5786), 2006. doi.org/10.1126/science.1127647
- D. P. Kingma, M. Welling. Auto-encoding variational Bayes. ICLR 2014. arXiv:1312.6114. arxiv.org/abs/1312.6114
- 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
- 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
- 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
- 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
- Y. Song, S. Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS 2019. arXiv:1907.05600. arxiv.org/abs/1907.05600
- J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. NeurIPS 2020. arXiv:2006.11239. arxiv.org/abs/2006.11239
- L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR 2017. arXiv:1605.08803. arxiv.org/abs/1605.08803
- O. Ronneberger, P. Fischer, T. Brox. U-Net: Convolutional networks for biomedical image segmentation. MICCAI 2015. arXiv:1505.04597. arxiv.org/abs/1505.04597
- J. Song, C. Meng, S. Ermon. Denoising diffusion implicit models. ICLR 2021. arXiv:2010.02502. arxiv.org/abs/2010.02502
- 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
- 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
- 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
- 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
- L. Weng. What are diffusion models? Lil'Log, 2021. lilianweng.github.io/posts/2021-07-11-diffusion-models/