Diffusion Models

Creating images from noise.

⏱ 3 min read▶ Video at 1:21:54🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

Here's a thought experiment. Imagine dropping ink in water. It slowly spreads until the water's evenly grey. That's easy. But could you run it backward, and watch the grey water gather itself back into a drop? If you could learn to do that, you could create things. That's diffusion.

Ink spreads through water
Could you run it backward?
That's diffusion

Noise to picture

Diffusion has two parts. Forward is easy. Take an image, and add a little noise, over and over, until only noise remains. Backward is the clever part. We train a network to look at a noisy image, and to predict the noise in it. That's a denoiser, just as you built earlier. Then we start with pure noise, and remove a bit at a time. Step by step, an image appears. It doesn't copy anything. It's growing a new image, using what it knows about real ones.

Destroying an image

The forward process has a neat formula. Alpha bar tells us how much of the original signal is left at any step. So we can jump straight to any noise level, without looping. Here we track a signal through a thousand steps. Run it. The signal starts near one, and fades to almost nothing, leaving pure noise with a spread of one. That's what makes this so useful. Random noise is the destination of the forward process, so it can be the start of the backward process.

forward.pypython

import torch

torch.manual_seed(0)
x = torch.ones(10000)

T = 1000
betas = torch.linspace(1e-4, 0.02, T)
alpha_bar = torch.cumprod(1 - betas, dim=0)

for t in [0, 200, 500, 800, 999]:
    noise = torch.randn_like(x)
    noisy = alpha_bar[t].sqrt() * x + (1 - alpha_bar[t]).sqrt() * noise
    print(f"step {t:>3}  signal left {alpha_bar[t].sqrt():.3f}"
          f"  mean {noisy.mean():.2f}  spread {noisy.std():.2f}")

Output

step   0  signal left 1.000  mean 1.00  spread 0.01
step 200  signal left 0.810  mean 0.80  spread 0.59
step 500  signal left 0.279  mean 0.26  spread 0.95
step 800  signal left 0.039  mean 0.05  spread 1.00
step 999  signal left 0.006  mean 0.01  spread 0.99

A real diffusion model

Here's a small but real diffusion model. The network takes a noisy image, and the time step, and it predicts the noise. Training is just five lines. Pick a random time, add noise, and ask the network to predict that noise. To generate, we start with pure noise, and walk backward through all two hundred steps. In every step, the network removes a little noise. Run it. Digits, grown out of pure randomness. Add text guidance and a much larger network, and you have the tools that make images from a sentence.

diffusion.pypython

import torch
from torch import nn
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

tf = transforms.Compose([transforms.Resize(16), transforms.ToTensor()])
train = datasets.MNIST("data", train=True, transform=tf)
train = torch.utils.data.Subset(train, range(20000))
dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
T = 200
betas = torch.linspace(1e-4, 0.05, T)
abar = torch.cumprod(1 - betas, 0)

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.t = nn.Embedding(T, 64)
        self.a = nn.Sequential(nn.Conv2d(1, 64, 3, padding=1), nn.GELU())
        self.b = nn.Sequential(nn.Conv2d(64, 128, 3, 2, 1), nn.GELU())
        self.c = nn.Sequential(nn.Conv2d(128, 128, 3, padding=1), nn.GELU())
        self.d = nn.Sequential(nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.GELU())
        self.e = nn.Conv2d(128, 1, 3, padding=1)
    def forward(self, x, t):
        a = self.a(x) + self.t(t)[:, :, None, None]
        y = self.d(self.c(self.b(a)))
        return self.e(torch.cat([y, a], 1))

torch.manual_seed(0)
net = Net()
opt = torch.optim.Adam(net.parameters(), lr=2e-3)
for epoch in range(5):
    for x, _ in dl:
        x = x * 2 - 1
        t = torch.randint(0, T, (len(x),))
        noise = torch.randn_like(x)
        ab = abar[t][:, None, None, None]
        noisy = ab.sqrt() * x + (1 - ab).sqrt() * noise
        loss = ((net(noisy, t) - noise) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()

x = torch.randn(16, 1, 16, 16)
with torch.no_grad():
    for i in reversed(range(T)):
        t = torch.full((16,), i)
        eps = net(x, t)
        x = (x - betas[i] / (1 - abar[i]).sqrt() * eps) / (1 - betas[i]).sqrt()
        if i > 0:
            x = x + betas[i].sqrt() * torch.randn_like(x)
fig, ax = plt.subplots(2, 8, figsize=(9, 2.6))
for a, img in zip(ax.ravel(), x[:, 0]):
    a.imshow(img.clamp(-1, 1), cmap="gray"); a.axis("off")
plt.tight_layout()
plt.savefig("diffusion.png", dpi=120)
print("Grown from pure noise, one small step at a time")

Output

Grown from pure noise, one small step at a time

Summary

Destroy images with noise, then teach a network to reverse it. That simple idea powers the image generators that amazed the world. Aria can now see, read, write, and imagine. Next, the most useful practical skill of them all. Standing on the shoulders of giants.

Key takeaways

Forward: add noise
Train: predict the noise
Generate: remove noise step by step

Next up: Transfer Learning — standing on the shoulders of giants

❮ GANs Transfer Learning ❯