Generative Adversarial Networks and Diffusion

Week 13 of 18 · Generative and self-supervised · 7 days

Full curriculum
Week 13 · Generative and self-supervised

Generative Adversarial Networks and Diffusion

Week 13 · Day 1 of 7

Two Networks, One Game

Replacing a pixel loss with an opponent, and the first samples

By 1156 words

Week 12 ended with a complaint. The variational autoencoder could generate digits from noise, but they came out soft, and the reason was the loss: averaging over every plausible digit is what minimises a pixel-wise error, and the average of several plausible digits is a blur. No amount of extra capacity fixes a loss that rewards hedging.

So replace the loss. Instead of scoring an output against a target pixel by pixel, train a second network whose only job is to tell real images from generated ones, and score the generator by whether it fools that network. An average of two digits does not fool anybody, so hedging stops paying.

Generative adversarial network: Two networks trained against each other. The generator turns a random vector into an image. The discriminator looks at an image and outputs a score for whether it is real. The generator's loss is the discriminator's mistake, so improving one makes the other's job harder.

The two networks

Neither is unusual on its own. The generator is a stack that widens from a small random vector up to 784 numbers, ending in Tanh because the images have been scaled to the range minus one to one. The discriminator is an ordinary binary classifier that happens to be looking at images.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
torch.manual_seed(0)
G, D = Gen(), Disc()
G.eval()
with torch.no_grad():
fake = G(torch.randn(2, 32))
print('the discriminator scores this noise at %.3f'
% D(fake).mean())
show(fake[0])
the discriminator scores this noise at -0.083
=++==++==+=++===+++==++++==+
+==+++=+++=+=++++++===++=+++
++++++===*+===+=++==++=+++++
+=+====+++++=+=====++++++===
++=+=+=+=++=+====+=++++===++
+=++===+==++=+==+=+==+=+==+=
+++==+++==++==++==+++-=++==+
======++=++=+=+==+++=+==+=++
===++=++++=+=++===++=-=====+
=====+==++++++=+=====+++==++
==++==+-+++=+-==++=++==+=+=+
+=++=-++==+++*++=+==+==-+++=
==++=+==++==++==+==++=++=+==
===+=+===++*==+=+===+====++=

Static, and a discriminator that has not been trained either, sitting near zero because it has no opinion yet. A score of zero is the boundary: the network outputs a logit, so positive means it thinks the image is real and negative means it does not.

The game

Each batch runs two updates, in this order:

  1. Show the discriminator a batch of real images labelled one and a batch of generated images labelled zero, and take one step on its parameters. The generated images are detach()ed, because this step must not change the generator.
  2. Generate a fresh batch, score it with the discriminator, and take one step on the generator's parameters in the direction that makes the discriminator call those images real. The discriminator is not updated here even though the gradient passes straight through it.

The detach is not an optimisation

Leaving it out means the discriminator step also updates the generator, in the direction that makes its own output easier to detect. The run does not crash and the losses look plausible. It simply never produces anything, which is the worst kind of bug to have in a model that is expected to look bad for the first few epochs anyway.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
bce = nn.BCEWithLogitsLoss()

def train_gan(epochs=30, zdim=32, d_steps=1, saturating=False,
lr_g=2e-4, lr_d=2e-4, seed=0, trace_every=0):
torch.manual_seed(seed)
G, D = Gen(zdim), Disc()
og = torch.optim.Adam(G.parameters(), lr=lr_g, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=lr_d, betas=(0.5, 0.999))
for ep in range(1, epochs + 1):
dsum = gsum = acc = n = 0.0
for xb, _ in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
for _ in range(d_steps):
fake = G(torch.randn(len(real), zdim)).detach()
loss_d = bce(D(real), ones) + bce(D(fake), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim))
score = D(fake)
loss_g = -bce(score, zeros) if saturating else bce(score, ones)
og.zero_grad()
loss_g.backward()
og.step()
dsum += loss_d.item()
gsum += loss_g.item()
acc += ((D(real) > 0).float().mean().item()
+ (score < 0).float().mean().item()) / 2
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %12.3f %12.3f %14.2f'
% (ep, dsum / n, gsum / n, acc / n))
G.eval()
return G, D
print('%6s %12s %12s %14s' % ('epoch', 'D loss', 'G loss', 'D accuracy'))
G, D = train_gan(epochs=30, trace_every=5)
torch.manual_seed(7)
with torch.no_grad():
for img in G(torch.randn(3, 32)):
show(img)
print()
epoch D loss G loss D accuracy
5 1.189 0.833 0.87
10 1.153 0.984 0.79
15 1.084 1.157 0.81
20 1.085 1.123 0.78
25 1.006 1.209 0.80
30 1.007 1.230 0.80
. .- .
. + . : .
-. . %%.: .
. +@@@@@@@@@@%.
..#@@. +%.
- :*# .**@++.
: :+. .-%%@%+: .
.-: =@@%*+##--
+=-=*@%...:@%# :
= + . *=-.:
.. . : :#@@% +
- *:%@%@@@@@@@% *
.::@=:=
: .

.
-
. . :: .
. . . - . . +
#+@@==.:.:
--+-%- . .
.: -.. .
+ =#:=
+ .*%@@=
. .=#%@@@= +
*%@@@: . .
: .:*: .
. ..
= -

.

. -. .
.##@=--=:
---%=#=.: --
..: :#. :-.
--@*+-: +
- %%-#..
:-==
. . ..=#++.-%##
:.%%@#:.*@#:
-=%@+*-
:

Thirty epochs on eight thousand images, from a generator that is four dense layers. What comes out has the thickness and curvature of handwriting, closed loops in roughly the right places, and a dark background it was never told about. What it does not have is a particular digit you would confidently name. That is the honest description of this budget, and it is still a long way from the noise the same network produced a few blocks ago, with no target image anywhere in the loss.

The idea in one line

A pixel loss rewards hedging, so replace it with a network that is trying to catch you hedging. Everything difficult about GANs follows from the fact that this replacement loss is not fixed: it is another model, and it is moving.
Week 13 · Day 2 of 7

Reading a Run That Has No Validation Loss

Why the losses say so little, and the gradient that vanishes

By 1137 words

Look again at the trace from yesterday. Over thirty epochs the discriminator loss drifted down a little, the generator loss drifted up by about half, and discriminator accuracy slid from 0.87 to 0.80. Meanwhile the samples went from static to digit-shaped.

Sit with that for a moment. The generator's loss got worse while the generator got better. In any of the first twelve weeks that sentence would be a bug report.

Why the losses tell you so little

Both numbers are measured against an opponent that is changing. If the generator loss falls, that can mean the generator improved, or it can mean the discriminator got worse and is now being fooled by the same images it caught last epoch. The two are indistinguishable from the loss alone, and they call for opposite responses.

There is no validation loss for a GAN

You cannot hold out data and watch a number to decide when to stop, because there is no target to compare against. This is a real and permanent difference from everything in the first twelve weeks, and it is why the rest of this week is spent building measurements that do not come from the loss.

The one number worth reading

Discriminator accuracy, which the loop above tracks as the average of how often it calls real images real and generated images fake. At the start it should be high, because early fakes are obvious. If it stays pinned at one, the generator is getting no useful signal. If it falls to a half the discriminator has stopped discriminating and the generator is optimising against nothing. Somewhere in between, and moving slowly, is what a working run looks like.

The loss the original paper wrote down, and the one everybody uses

Stated as a game, the generator should minimise the probability the discriminator assigns to fake. Written that way it has a flat spot exactly where you need it most: when the discriminator is confident, the gradient reaching the generator goes to almost nothing, and the generator is stuck precisely because it is bad. The fix is to flip the problem round and have the generator maximise the probability of real instead. The two have the same optimum and completely different gradients.

That is the usual argument. It is also measurable, so here it is measured, with the discriminator given a single epoch of head start:

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
bce = nn.BCEWithLogitsLoss()

def train_gan(epochs=30, zdim=32, d_steps=1, saturating=False,
lr_g=2e-4, lr_d=2e-4, seed=0, trace_every=0):
torch.manual_seed(seed)
G, D = Gen(zdim), Disc()
og = torch.optim.Adam(G.parameters(), lr=lr_g, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=lr_d, betas=(0.5, 0.999))
for ep in range(1, epochs + 1):
dsum = gsum = acc = n = 0.0
for xb, _ in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
for _ in range(d_steps):
fake = G(torch.randn(len(real), zdim)).detach()
loss_d = bce(D(real), ones) + bce(D(fake), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim))
score = D(fake)
loss_g = -bce(score, zeros) if saturating else bce(score, ones)
og.zero_grad()
loss_g.backward()
og.step()
dsum += loss_d.item()
gsum += loss_g.item()
acc += ((D(real) > 0).float().mean().item()
+ (score < 0).float().mean().item()) / 2
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %12.3f %12.3f %14.2f'
% (ep, dsum / n, gsum / n, acc / n))
G.eval()
return G, D
torch.manual_seed(0)
G, D = Gen(), Disc()
# Give the discriminator one epoch of head start. This is not a contrived
# situation: at the beginning of any run the fakes are obvious noise and
# the discriminator gets very good very quickly.
od = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
for xb, _ in loader:
real = xb.flatten(1)
fake = G(torch.randn(len(real), 32)).detach()
loss_d = (bce(D(real), torch.ones(len(real), 1))
+ bce(D(fake), torch.zeros(len(real), 1)))
od.zero_grad()
loss_d.backward()
od.step()

def g_grad_norm(saturating):
G.zero_grad()
score = D(G(torch.randn(256, 32)))
zeros, ones = torch.zeros(256, 1), torch.ones(256, 1)
loss = -bce(score, zeros) if saturating else bce(score, ones)
loss.backward()
return torch.cat([p.grad.flatten()
for p in G.parameters()]).norm().item()

with torch.no_grad():
caught = (D(G(torch.randn(512, 32))) < 0).float().mean()
print('the discriminator already spots %.1f%% of the fakes' % (100 * caught))
print()
print('%-22s %16s' % ('generator loss', 'gradient norm'))
print('%-22s %16.6f' % ('saturating', g_grad_norm(True)))
print('%-22s %16.6f' % ('non-saturating', g_grad_norm(False)))
the discriminator already spots 100.0% of the fakes

generator loss gradient norm
saturating 0.068186
non-saturating 4.905675

The discriminator is catching every single fake after one epoch, which is exactly the regime the argument is about. Same generator, same discriminator, same batch of noise, and the gradient reaching the generator differs by a factor of roughly seventy. The saturating version is not slightly worse, it is delivering almost nothing to learn from at the moment the generator most needs a direction.

In code the whole change is bce(score, ones) in place of -bce(score, zeros). Two short expressions, both of which read as reasonable, which is why it is easy to write the wrong one and very hard to notice afterwards.

What to do when a run goes wrong

  • Discriminator accuracy pinned at 1.0. It is winning too easily. Slow it down: a lower learning rate for it than for the generator, or fewer of its steps per generator step.
  • Discriminator accuracy at 0.5 and generator loss falling steadily. The generator has found a hole rather than learned the data. Look at the samples before believing the number.
  • Both losses flat and samples static. Check the detach, check that both optimisers are stepping, and check that the generator's final activation matches the range the data was scaled to.
  • Everything looks fine and the samples are all the same digit. That is tomorrow.
Week 13 · Day 3 of 7

Mode Collapse

The failure the loss cannot show you, and how to count it

By 1146 words

A generator has an easy way to satisfy its loss that has no analogue in anything so far. If one particular output reliably fools the discriminator, producing that output every time scores well, and the loss has no term that asks for variety. The generator is being paid per image, not per distribution.

Mode collapse: When a generator produces only a narrow slice of the data, ignoring the rest. The loss does not show it, the samples look fine one at a time, and it is only visible when you count what comes out of many draws.

You have to count

There is no loss to read, so the measurement has to be built. Train an ordinary digit classifier on the real data, ask the generator for two thousand images, and count what the classifier says they are. A generator covering the data should produce something close to an even spread across ten digits, which has an entropy of ln(10), about 2.303. The further below that, the narrower the generator.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
bce = nn.BCEWithLogitsLoss()

def train_gan(epochs=30, zdim=32, d_steps=1, saturating=False,
lr_g=2e-4, lr_d=2e-4, seed=0, trace_every=0):
torch.manual_seed(seed)
G, D = Gen(zdim), Disc()
og = torch.optim.Adam(G.parameters(), lr=lr_g, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=lr_d, betas=(0.5, 0.999))
for ep in range(1, epochs + 1):
dsum = gsum = acc = n = 0.0
for xb, _ in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
for _ in range(d_steps):
fake = G(torch.randn(len(real), zdim)).detach()
loss_d = bce(D(real), ones) + bce(D(fake), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim))
score = D(fake)
loss_g = -bce(score, zeros) if saturating else bce(score, ones)
og.zero_grad()
loss_g.backward()
og.step()
dsum += loss_d.item()
gsum += loss_g.item()
acc += ((D(real) > 0).float().mean().item()
+ (score < 0).float().mean().item()) / 2
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %12.3f %12.3f %14.2f'
% (ep, dsum / n, gsum / n, acc / n))
G.eval()
return G, D
def train_judge():
"""A plain digit classifier, used only to score what the GAN makes."""
torch.manual_seed(1)
clf = nn.Sequential(nn.Flatten(),
nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 10))
opt = torch.optim.AdamW(clf.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
for _ in range(12):
for xb, yb in loader:
opt.zero_grad()
lf(clf(xb), yb).backward()
opt.step()
clf.eval()
test = Subset(datasets.MNIST('data', train=False, download=True,
transform=tf), range(2000))
tl = DataLoader(test, batch_size=512)
right = seen = 0
with torch.no_grad():
for xb, yb in tl:
right += (clf(xb).argmax(1) == yb).sum().item()
seen += len(yb)
return clf, right / seen

def coverage(G, clf, n=2000):
with torch.no_grad():
imgs = G(torch.randn(n, G.zdim)).reshape(-1, 1, 28, 28)
pred = clf(imgs).argmax(1)
counts = torch.bincount(pred, minlength=10).float()
frac = counts / counts.sum()
ent = -(frac[frac > 0] * frac[frac > 0].log()).sum().item()
return counts.long().tolist(), ent
clf, acc = train_judge()
print('judge accuracy on real digits %.4f' % acc)
print()
print('%-26s %-34s %8s' % ('recipe', 'digits produced out of 2000', 'entropy'))
for name, kw in [('standard, z=32', {}),
('saturating G loss', {'saturating': True}),
('z=2', {'zdim': 2}),
('5 D steps per G step', {'d_steps': 5}),
('D learns 10x faster', {'lr_d': 2e-3})]:
G, _ = train_gan(epochs=30, **kw)
counts, ent = coverage(G, clf)
print('%-26s %-34s %8.3f'
% (name, ' '.join('%3d' % c for c in counts), ent))
print()
print('for reference, a perfectly even spread scores %.3f'
% float(torch.tensor(10.0).log()))
judge accuracy on real digits 0.9075

recipe digits produced out of 2000 entropy
standard, z=32 167 178 172 189 273 184 218 222 279 118 2.275
saturating G loss 62 82 191 170 239 336 265 243 275 137 2.206
z=2 1 467 30 208 646 47 22 276 274 29 1.752
5 D steps per G step 3 485 126 90 335 154 40 304 257 206 2.026
D learns 10x faster 9 179 105 124 416 162 102 445 217 241 2.080

for reference, a perfectly even spread scores 2.303

The standard recipe comes out close to even, which is a better result than a four-layer generator deserves. Everything below it is narrower, and the clearest case is the one with a two-dimensional latent vector: whole digits nearly vanish from its output while others appear several hundred times. A latent space that small cannot be stretched over ten separated regions of digit space, so it covers a few and abandons the rest.

The judge is not perfect either

It is a two-layer classifier trained on the same eight thousand images, and it gets about nine real digits in ten right. So a count of zero in some column is partly the generator and partly the judge being wrong about odd-looking inputs. Read the table as a ranking, which it is reliable for, rather than as an exact census.

What actually helped, and what did not

The two adjustments most often recommended for stabilising a GAN, giving the discriminator extra steps and letting it learn faster, both came out narrower than doing neither. That is worth being direct about: they are remedies for a discriminator that is losing, and this discriminator was not losing. Applied to a run that does not have that problem they make the generator's job harder for no gain.

The saturating loss did less damage to coverage here than the reputation suggests, which is a reminder that day 2's measurement was about the gradient at the start of a run, not a guarantee about where the run ends up. A small model on an easy dataset recovers from a bad gradient that would sink a larger one.

The habit to keep

Any generative model needs a measurement that looks at many samples at once. Counting what a classifier calls them is crude and it costs one small model, and it catches the failure that eyeballing three pictures never will.
Week 13 · Day 4 of 7

Asking For a Particular Digit

Conditioning both networks, and how much control that really buys

By 1768 words

So far the generator takes noise and returns whatever it likes. Usually you want to ask for something specific, and the change needed is smaller than it sounds: give both networks the label.

The generator receives the label alongside its noise vector, so it can learn a different mapping for each one. The discriminator receives the label alongside the image, which is the part that matters. Without it the discriminator would only ask whether the image looks like a digit, and the generator could ignore the label entirely. With it, the question becomes whether this image looks like a real example of that label, and producing a convincing seven when asked for a three now loses.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
bce = nn.BCEWithLogitsLoss()

def train_gan(epochs=30, zdim=32, d_steps=1, saturating=False,
lr_g=2e-4, lr_d=2e-4, seed=0, trace_every=0):
torch.manual_seed(seed)
G, D = Gen(zdim), Disc()
og = torch.optim.Adam(G.parameters(), lr=lr_g, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=lr_d, betas=(0.5, 0.999))
for ep in range(1, epochs + 1):
dsum = gsum = acc = n = 0.0
for xb, _ in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
for _ in range(d_steps):
fake = G(torch.randn(len(real), zdim)).detach()
loss_d = bce(D(real), ones) + bce(D(fake), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim))
score = D(fake)
loss_g = -bce(score, zeros) if saturating else bce(score, ones)
og.zero_grad()
loss_g.backward()
og.step()
dsum += loss_d.item()
gsum += loss_g.item()
acc += ((D(real) > 0).float().mean().item()
+ (score < 0).float().mean().item()) / 2
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %12.3f %12.3f %14.2f'
% (ep, dsum / n, gsum / n, acc / n))
G.eval()
return G, D
class CondGen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.emb = nn.Embedding(10, 10)
self.net = nn.Sequential(
nn.Linear(zdim + 10, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z, y):
return self.net(torch.cat([z, self.emb(y)], 1))

class CondDisc(nn.Module):
def __init__(self):
super().__init__()
self.emb = nn.Embedding(10, 10)
self.net = nn.Sequential(
nn.Linear(784 + 10, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x, y):
return self.net(torch.cat([x.flatten(1), self.emb(y)], 1))

def train_cgan(epochs=30, zdim=32, seed=0):
torch.manual_seed(seed)
G, D = CondGen(zdim), CondDisc()
og = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
for _ in range(epochs):
for xb, yb in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
fake = G(torch.randn(len(real), zdim), yb).detach()
loss_d = bce(D(real, yb), ones) + bce(D(fake, yb), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim), yb)
loss_g = bce(D(fake, yb), ones)
og.zero_grad()
loss_g.backward()
og.step()
G.eval()
return G
print(CondGen())
CondGen(
(emb): Embedding(10, 10)
(net): Sequential(
(0): Linear(in_features=42, out_features=128, bias=True)
(1): LeakyReLU(negative_slope=0.2)
(2): Linear(in_features=128, out_features=256, bias=True)
(3): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(4): LeakyReLU(negative_slope=0.2)
(5): Linear(in_features=256, out_features=784, bias=True)
(6): Tanh()
)
)

Does the label actually control the output

The same judge from yesterday answers this directly. Ask for two hundred of each digit and count how often the classifier agrees with what was requested.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
class Gen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.net = nn.Sequential(
nn.Linear(zdim, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z):
return self.net(z)

class Disc(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x):
return self.net(x.flatten(1))
bce = nn.BCEWithLogitsLoss()

def train_gan(epochs=30, zdim=32, d_steps=1, saturating=False,
lr_g=2e-4, lr_d=2e-4, seed=0, trace_every=0):
torch.manual_seed(seed)
G, D = Gen(zdim), Disc()
og = torch.optim.Adam(G.parameters(), lr=lr_g, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=lr_d, betas=(0.5, 0.999))
for ep in range(1, epochs + 1):
dsum = gsum = acc = n = 0.0
for xb, _ in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
for _ in range(d_steps):
fake = G(torch.randn(len(real), zdim)).detach()
loss_d = bce(D(real), ones) + bce(D(fake), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim))
score = D(fake)
loss_g = -bce(score, zeros) if saturating else bce(score, ones)
og.zero_grad()
loss_g.backward()
og.step()
dsum += loss_d.item()
gsum += loss_g.item()
acc += ((D(real) > 0).float().mean().item()
+ (score < 0).float().mean().item()) / 2
n += 1
if trace_every and ep % trace_every == 0:
print('%6d %12.3f %12.3f %14.2f'
% (ep, dsum / n, gsum / n, acc / n))
G.eval()
return G, D
class CondGen(nn.Module):
def __init__(self, zdim=32):
super().__init__()
self.zdim = zdim
self.emb = nn.Embedding(10, 10)
self.net = nn.Sequential(
nn.Linear(zdim + 10, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2),
nn.Linear(256, 784), nn.Tanh())

def forward(self, z, y):
return self.net(torch.cat([z, self.emb(y)], 1))

class CondDisc(nn.Module):
def __init__(self):
super().__init__()
self.emb = nn.Embedding(10, 10)
self.net = nn.Sequential(
nn.Linear(784 + 10, 256), nn.LeakyReLU(0.2),
nn.Linear(256, 128), nn.LeakyReLU(0.2),
nn.Linear(128, 1))

def forward(self, x, y):
return self.net(torch.cat([x.flatten(1), self.emb(y)], 1))

def train_cgan(epochs=30, zdim=32, seed=0):
torch.manual_seed(seed)
G, D = CondGen(zdim), CondDisc()
og = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
od = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
for _ in range(epochs):
for xb, yb in loader:
real = xb.flatten(1)
ones = torch.ones(len(real), 1)
zeros = torch.zeros(len(real), 1)
fake = G(torch.randn(len(real), zdim), yb).detach()
loss_d = bce(D(real, yb), ones) + bce(D(fake, yb), zeros)
od.zero_grad()
loss_d.backward()
od.step()
fake = G(torch.randn(len(real), zdim), yb)
loss_g = bce(D(fake, yb), ones)
og.zero_grad()
loss_g.backward()
og.step()
G.eval()
return G
def train_judge():
"""A plain digit classifier, used only to score what the GAN makes."""
torch.manual_seed(1)
clf = nn.Sequential(nn.Flatten(),
nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 10))
opt = torch.optim.AdamW(clf.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
for _ in range(12):
for xb, yb in loader:
opt.zero_grad()
lf(clf(xb), yb).backward()
opt.step()
clf.eval()
test = Subset(datasets.MNIST('data', train=False, download=True,
transform=tf), range(2000))
tl = DataLoader(test, batch_size=512)
right = seen = 0
with torch.no_grad():
for xb, yb in tl:
right += (clf(xb).argmax(1) == yb).sum().item()
seen += len(yb)
return clf, right / seen

def coverage(G, clf, n=2000):
with torch.no_grad():
imgs = G(torch.randn(n, G.zdim)).reshape(-1, 1, 28, 28)
pred = clf(imgs).argmax(1)
counts = torch.bincount(pred, minlength=10).float()
frac = counts / counts.sum()
ent = -(frac[frac > 0] * frac[frac > 0].log()).sum().item()
return counts.long().tolist(), ent
clf, _ = train_judge()
G = train_cgan(epochs=30)
torch.manual_seed(5)
print('%8s %10s' % ('asked for', 'agreement'))
total = 0.0
with torch.no_grad():
for digit in range(10):
y = torch.full((200,), digit, dtype=torch.long)
imgs = G(torch.randn(200, 32), y).reshape(-1, 1, 28, 28)
agree = (clf(imgs).argmax(1) == digit).float().mean().item()
total += agree / 10
print('%8d %10.3f' % (digit, agree))
print('%8s %10.3f' % ('overall', total))
torch.manual_seed(9)
with torch.no_grad():
for digit in [3, 7]:
print()
print('asked for a %d' % digit)
show(G(torch.randn(1, 32), torch.tensor([digit]))[0])
asked for agreement
0 0.915
1 0.825
2 0.755
3 0.590
4 0.575
5 0.740
6 0.215
7 0.255
8 0.605
9 0.370
overall 0.584

asked for a 3
. - .
@
: . .:= . .
* . #-.=%@%-# .
* =%@%%-.@*.
=@#. =. . #
+:@@@. # *
+@#@@@@.- ..=
: : : ..-. ++ :*#% .
=.. @#+. .@@@@@@+ .
.. :@@@% .#@@@@:: -
.- .@@@+% =
. . #: % :


asked for a 7
.
- - . : .
.
+.#%*.: . .
. +=+@=%%+@@-. .#%:
=-*#-: .#%*-:. %. :.
. .=+.=:*-=%-:.%::--
. .-:.+@*@%@%@-. .
.. : =**%*:*
. :- -@@%- . .
.-%%-.. : .
. .%@%==: -
.-. %%%@@:. . =
. . -. :: .

Against a floor of 0.1 for a generator that ignored the label completely, and a ceiling of about 0.91 set by the judge's own accuracy on real digits, the overall figure is a little under 0.6. So the conditioning works, and calling it working needs qualifying: some digits come back reliably and others hardly ever do, with the worst column barely above a fifth.

Why the columns differ so much

Nothing in the training gave equal weight to being right about each digit. The generator is rewarded for fooling the discriminator, and for a label whose real examples vary a lot in shape, a vague blob that leans towards the class average is often enough to pass. The judge, which was trained on clean data, then declines to call it that digit. The discriminator and the judge are asking different questions, and the gap between them is where the missing forty percent sits.

This is the pattern, not the exception

Every conditional generator has this gap: the training signal is plausibility, and what you actually want is plausibility and obedience. Larger models close it rather than remove it. When a text-to-image model gives you four of the five things you asked for, this is the same failure at a different scale.

Week 13 · Day 5 of 7

Destroying an Image on Purpose

The forward process, the schedule, and the jump that makes it train

By 878 words

Everything in the first four days rests on one uncomfortable arrangement: the loss is another network, and it is moving. Diffusion gets to the same place without that. It defines a process that destroys an image, one small step at a time, until nothing is left but noise, and then trains a network to undo a single step of it. There is no opponent, and the loss is an ordinary mean squared error that means what it says.

The forward process: A fixed recipe, with no learned parameters, that takes an image and repeatedly shrinks it slightly towards zero while adding a little Gaussian noise. After enough steps the result is indistinguishable from pure noise, and every step along the way is a training example.

The schedule

One number per step, conventionally called beta, says how much noise that step adds. Start it small so the early steps barely disturb the image, and grow it so the later ones finish the job. What the network actually needs is not beta but its running product, written abar: the fraction of the original image still standing at step t.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
x0 = train_set[1][0].flatten()
torch.manual_seed(0)
for t in [0, 40, 100, 199]:
eps = torch.randn_like(x0)
xt = abar[t].sqrt() * x0 + (1 - abar[t]).sqrt() * eps
print('t = %-4d abar = %.4f mean %.3f std %.3f'
% (t, abar[t], xt.mean(), xt.std()))
show(xt, step=3)
print()
t = 0 abar = 0.9999 mean -0.688 std 0.658


:%@@@@@:
.@@@#=@%.-@*
#@@: . @@#
-@@= @@+
-@+ .*@*
-@@@%%@@@#+



t = 40 abar = 0.8102 mean -0.614 std 0.731
: .: : ::.: - +.
:-. - . + : . =
.. : : -@%%*=@: = :
.-: .: .:*+@=%***-=%# : -
.. : =@@ .. - @@#:
:.*.+%@%+ . ..:-@+* -
.: -.=@# .: = =@#. - :
: --.#@@*@%%@@@ . -= :
: : .. : %= .. .
. . :* - - . : ..

t = 100 abar = 0.2760 mean -0.424 std 0.925
.+.. =. %: = = = % : %-::
=- +.+ -= - @ %.*: @.#=%
= % :-+- : +@=@#@==* :-
# %: @-@:%#@.@*+--@= : *-
*: *. -:@--: . #@:#% .:-+
:= =@**-: :#+.# + @*@ %
* .*.+ @*.-@** *%%@ +-%+#
.:@% =*#@=.#@:@@= + :+ .
- =+# :=:: -:-+ + :@. -
:* : *--: @ .*- : = *.

t = 199 abar = 0.0061 mean -0.080 std 0.976
-@%@ : @=.#@:@- @*+@-: *-
@@. %@@=#**.@ %=**- :* %
@ +% = *= *@@@=@ *%##@#*:@@
@: @ **@= =-%+-.* -:@- *:@ -
@*.+= @.@ = %*:@@+@-++ *=
%@. = * # @@=@*+@:#+%:# @
# = @.@. + .-#*-@+=*.*%@ +
@ -:+#@@.-.# #: =@.@#:%=@ =
+:#-+%.* #@ @# @=# : +@ %%=
#-@-:@*@-#@.%: # %=-*- = -

The digit is intact at the start, still legible under the noise at step 40, arguably visible at step 100 if you already know it is there, and gone by step 199. Watch the two statistics as well: the last row has a mean near zero and a standard deviation near one, which is to say it has become a sample from a standard normal distribution. That is the whole point of the schedule, because it is the distribution sampling will start from later.

Check that your schedule actually finishes

The first version of this week used a gentler schedule whose abar at the final step was 0.13, not 0.006. That leaves an eighth of the original image still present, so the model was never trained on anything resembling pure noise, and sampling began from a state it had never seen. It produced static. Print abar at your last step before you train anything, and if it is not close to zero, your schedule is too short or too gentle.

The jump that makes any of this trainable

Taken literally, producing a training example at step 500 means running 500 steps. Do that for every image in every batch and there is no training run. It is unnecessary: a chain of Gaussian steps is itself Gaussian, and abar is exactly the term that lets you skip straight there in one line.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
x0 = train_set[1][0].flatten()
torch.manual_seed(0)

# the slow way: 80 actual steps of the process, 400 times over
runs = []
for _ in range(400):
x = x0.clone()
for t in range(80):
x = alpha[t].sqrt() * x + beta[t].sqrt() * torch.randn_like(x)
runs.append(x)
slow = torch.stack(runs)

# the fast way: straight to t=79 in one line
fast = torch.stack([abar[79].sqrt() * x0
+ (1 - abar[79]).sqrt() * torch.randn_like(x0)
for _ in range(400)])

print('%-24s %12s %12s' % ('', 'mean pixel', 'std pixel'))
print('%-24s %12.4f %12.4f' % ('80 small steps', slow.mean(), slow.std()))
print('%-24s %12.4f %12.4f' % ('one closed-form jump', fast.mean(),
fast.std()))
mean pixel std pixel
80 small steps -0.4620 0.8624
one closed-form jump -0.4603 0.8642

Same distribution, one line instead of eighty. This is why training picks a random t for every image in the batch and jumps there directly, and it is the single fact that turns diffusion from a thought experiment into something you can run.

Week 13 · Day 6 of 7

Learning to Undo One Step

Predicting the noise, the loss floor, and watching a digit appear

By 1585 words

The forward process has no parameters. Everything learned sits in one network with one job: given a noisy image and a number saying how noisy it is, say which noise was added.

Predicting the noise rather than the clean image looks like an odd choice, and the two are algebraically interchangeable, since knowing either one and the noisy image gives you the other. Predicting the noise turns out to be better behaved across the whole range of t, and it is what the standard formulation does.

Telling the network what time it is

The same noisy image means different things at different t, so the step number has to reach the network. Here it is an embedding looked up per step and added inside every block, rather than only at the input. That detail matters more than it looks: a first attempt that conditioned only the input layer underfit badly.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
class Block(nn.Module):
"""One residual block, told at every level how noisy the input is."""
def __init__(self, h):
super().__init__()
self.norm = nn.LayerNorm(h)
self.tproj = nn.Linear(h, h)
self.net = nn.Sequential(nn.SiLU(), nn.Linear(h, h),
nn.SiLU(), nn.Linear(h, h))

def forward(self, x, temb):
return x + self.net(self.norm(x) + self.tproj(temb))

class Denoiser(nn.Module):
"""Given a noisy image and how noisy it is, predict the noise."""
def __init__(self, hidden=768, blocks=3):
super().__init__()
self.inp = nn.Linear(784, hidden)
self.temb = nn.Embedding(T, hidden)
self.blocks = nn.ModuleList([Block(hidden) for _ in range(blocks)])
self.out = nn.Sequential(nn.LayerNorm(hidden), nn.SiLU(),
nn.Linear(hidden, 784))

def forward(self, x, t):
temb = self.temb(t)
h = self.inp(x)
for block in self.blocks:
h = block(h, temb)
return self.out(h)

def add_noise(x0, t, eps):
a = abar[t].reshape(-1, 1)
return a.sqrt() * x0 + (1 - a).sqrt() * eps

def train_diffusion(epochs=200, seed=0, trace_every=0):
torch.manual_seed(seed)
net = Denoiser()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
for ep in range(1, epochs + 1):
total = n = 0.0
for xb, _ in loader:
x0 = xb.flatten(1)
t = torch.randint(0, T, (len(x0),))
eps = torch.randn_like(x0)
loss = ((net(add_noise(x0, t, eps), t) - eps) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
sched.step()
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
net.eval()
return net

@torch.no_grad()
def sample(net, n=8, seed=3, keep=()):
torch.manual_seed(seed)
x = torch.randn(n, 784)
frames = {}
for step in range(T - 1, -1, -1):
t = torch.full((n,), step, dtype=torch.long)
# what the network thinks the clean image is, from here
x0 = ((x - (1 - abar[step]).sqrt() * net(x, t))
/ abar[step].sqrt()).clamp(-1, 1)
# the posterior mean: a small step from x towards that guess
mean = (abar_prev[step].sqrt() * beta[step] * x0
+ alpha[step].sqrt() * (1 - abar_prev[step]) * x
) / (1 - abar[step])
x = mean if step == 0 else mean + beta[step].sqrt() * torch.randn_like(x)
if step in keep:
frames[step] = x.clone()
return (x, frames) if keep else x
print('a network that always guesses zero scores %.4f'
% (torch.randn(4000, 784) ** 2).mean())
print()
print('%6s %14s' % ('epoch', 'noise MSE'))
net = train_diffusion(epochs=200, trace_every=25)
torch.save(net.state_dict(), 'denoiser.pt')
print()
for img in sample(net, n=3):
show(img)
print()
a network that always guesses zero scores 0.9993

epoch noise MSE
25 0.2013
50 0.1796
75 0.1707
100 0.1539
125 0.1544
150 0.1438
175 0.1431
200 0.1380

--=-:----::----=:-:--:=--.--
----:-:::-=-+-+-: ::::-.--:
--==------:=:-- :=%***+==:-
-:-=---.==%@%***@@@@@#%#=--:
:--:. . -*%*+-=#+=:=+: . =:
::-. .+#*#= **= .==*+-:==
-::-::=@@@@=.*##+: .:@#*-==-
-:--:=*@@%=*+#=:*%= :*#*+-:-
:-:--**@*: **++@@%:..+*==:.:
::-=+@@%: :-#%--:=+***+==--
:--==@@*#. .#%=#%@#-:-=-=:-:
::--:==#@**=*=:::*. :------:
-.=:--:+==*- -*+--:::::.--
---:---==--..:.:-==:::-:-=:-

*++++++*+=++**=+++=+*****+=+
****=*=*=++**+*+*##**+*+*#=+
++++**+*=+=: -+%@%==#**+*+*
*+*#***+:. :::+.%@@@#***+
*++#%@@@@@#+#:-:+.+-::#*****
#**%#@@@#*-+=##%%#@: -====+
=+++%@@@#=:-*@@%**:=::++*+++
*#++#@@@@%--: :.. .###****=*
++=#*#%%@#*:: . .#@@@@+%+++
*+*=++-=##*:--=...=*@#+*+*++
*++##+. . . .=-=*%#*####*+++
#++=*#%: .=%@%@*#*#+*++++*
#+++=*++.=-#@@#*#*=+%+++++**
*+*++*=#+==+%%+++*##*++***=*

****+*=+==*=**=*+*=*++**+=*=
=++#+#+#+*++=+=**++++*++*++=
*+****=+++*###%@*=-*@@##*=+=
++*+++=**%++=#@@@%#%%%@#-+*+
*++#*##%*- .##*#.::++=*++*+*
*=*+==#*. .@%%=. --+:=*+*+++
=+**=#%+=--@@@%#=@+ :%#*+++
+*+*+@%%+ -+*+%%#@: +@%*++*
=#+#+@@#- .:-=+=:---+%%#+++
**+++@@@- .-#%*. =@@@%#*+##
*+*#@%@+-=@@@%= :+@@@%*++*+
*+**#%%@@@@@@@*-*%**##*+=+++
*+***+=====#: :*+=#+++=+=++
++*=++#+=+-++**#++#++#+*++-+

Why the loss does not go anywhere near zero

Noise drawn from a standard normal has a mean squared value of one, so guessing zero every time scores about 1.0, and that is the number to beat. Reaching 0.14 is a large win, and no amount of training would reach zero, because a good part of this loss is irreducible.

Consider a very small t. The noisy image is almost the clean image, with a whisper of noise on top. Recovering which noise that was would mean knowing the clean image to a precision far beyond what is there to see, so the best possible prediction is close to zero and the loss at those steps stays close to one. The average over all t therefore has a floor well above zero, and comparing runs is the only way to read it.

The samples are digits. They also carry a haze across the background that a real MNIST image does not have, which is what an under-trained denoiser looks like: each reverse step leaves a small error, two hundred steps accumulate it, and nothing in the process cleans it up at the end. The loss was still falling slowly at epoch 200, so this is a compute budget, not a broken method.

Running the process backwards

Sampling starts from pure noise and walks back down. At each step the network predicts the noise, which is rearranged into a guess at the clean image, clamped to the valid pixel range, and then used to take one small step back towards it, with fresh noise added. Reading it as guess where this is going, move a little way there, jitter is a fair description of the arithmetic.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
class Block(nn.Module):
"""One residual block, told at every level how noisy the input is."""
def __init__(self, h):
super().__init__()
self.norm = nn.LayerNorm(h)
self.tproj = nn.Linear(h, h)
self.net = nn.Sequential(nn.SiLU(), nn.Linear(h, h),
nn.SiLU(), nn.Linear(h, h))

def forward(self, x, temb):
return x + self.net(self.norm(x) + self.tproj(temb))

class Denoiser(nn.Module):
"""Given a noisy image and how noisy it is, predict the noise."""
def __init__(self, hidden=768, blocks=3):
super().__init__()
self.inp = nn.Linear(784, hidden)
self.temb = nn.Embedding(T, hidden)
self.blocks = nn.ModuleList([Block(hidden) for _ in range(blocks)])
self.out = nn.Sequential(nn.LayerNorm(hidden), nn.SiLU(),
nn.Linear(hidden, 784))

def forward(self, x, t):
temb = self.temb(t)
h = self.inp(x)
for block in self.blocks:
h = block(h, temb)
return self.out(h)

def add_noise(x0, t, eps):
a = abar[t].reshape(-1, 1)
return a.sqrt() * x0 + (1 - a).sqrt() * eps

def train_diffusion(epochs=200, seed=0, trace_every=0):
torch.manual_seed(seed)
net = Denoiser()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
for ep in range(1, epochs + 1):
total = n = 0.0
for xb, _ in loader:
x0 = xb.flatten(1)
t = torch.randint(0, T, (len(x0),))
eps = torch.randn_like(x0)
loss = ((net(add_noise(x0, t, eps), t) - eps) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
sched.step()
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
net.eval()
return net

@torch.no_grad()
def sample(net, n=8, seed=3, keep=()):
torch.manual_seed(seed)
x = torch.randn(n, 784)
frames = {}
for step in range(T - 1, -1, -1):
t = torch.full((n,), step, dtype=torch.long)
# what the network thinks the clean image is, from here
x0 = ((x - (1 - abar[step]).sqrt() * net(x, t))
/ abar[step].sqrt()).clamp(-1, 1)
# the posterior mean: a small step from x towards that guess
mean = (abar_prev[step].sqrt() * beta[step] * x0
+ alpha[step].sqrt() * (1 - abar_prev[step]) * x
) / (1 - abar[step])
x = mean if step == 0 else mean + beta[step].sqrt() * torch.randn_like(x)
if step in keep:
frames[step] = x.clone()
return (x, frames) if keep else x
net = Denoiser()
net.load_state_dict(torch.load('denoiser.pt'))
net.eval()
final, frames = sample(net, n=1, seed=11, keep=(150, 100, 50, 20))
for step in [150, 100, 50, 20]:
print('%d steps still to go' % step)
show(frames[step][0], step=3)
print()
print('finished')
show(final[0], step=3)
150 steps still to go
@-: **%.%= @= @%@:+ @ -@-@
@ @ :+===+ = : -@=@@@= @=-
. @=-=@* +.@ #-@*%-@=- # -@+
@#=@#+.#.::@*+@@= +@@@@ @ %@
*@ -%+ @.:#@@ @ @@ .=. *
.@#+ # . @@ @@=- %. -=#:*%
%@+: +-:@@ @+@=@ @@ +%%@+=**
=+@ ++ =%=* @@@:=:@:- ::=.
@+: *#:#@:@@#+*@- -.*%-+=+
.@# +:@%@ .% .@@@:*%:=@*#@:

100 steps still to go
@-*:.**%-%=:=#@+**: @ @-
:#@ :* *. :=- :@@##+@.%*
.=@. --+ @%@ #@@#%# . # #+:
##-*+.%*+ .@+@.@ -@+@+ # .
*=-:@* : +=:# = # - . *+.-
*@. +% = ##@@-@-* # =%@@ #
* -: @@+-%*#+# *+%:%%## .
%:@ #@ @@=+@@*+. +@:#.+.%**
@ @%@+.:%++ @ -%:-=-* # =
*:-:- #@#% @. -=+*=* %**

50 steps still to go
+=- --- :==-=#%***-. +.==#
:#+.:. =.*: :- +#@ -#- . =
.:.:+:=- ++- :@#@#++*: :-:-
*-+==#=*+:@@@@@@ =*@+**:. *-
:#**-%+ . -+ . @.-.-.-=+.
+=-:#-=-+ +*@@%@% .. #=++=
:= . +%%* @*= .-=.#- .:
# *.:@ @@@%@@## .:@@@ *-=-*
=:-*##.=: --:#@+.@:==- -
*: . :#@-*--*@ + +#.*+-+ @%

20 steps still to go
.. :::.:---.=:*-=-=-.::+=:+:
.-=.==-=+ .. -=+-=*=:=+* :::
-::-----::=.:-+*#+#+=-.:-:*=
+=:-*==***@#@@@@:*@@@@++=..:
=---=+*. - +-..:+-:::
---:=..:-+#+@@@@@ ..= =*==-.
.-.-.-.*@@=+@%= . -*#:=-:::=
+:+-:+:%%@%@%-- =+@@++=**-=:
.::-=+:-*-:-.:=*+###-*.:+:.=
--=-.--=.-:=++:=:-=.=-=== :-

finished
-:--:-:-:--:--------:---==-:
--=:::==+--.:--=+*+=----::--
-:----==--- .=++%###=:.-:--=
-=:==:=#%@%@%@@%-+%@@@+*+:--
-=--=+*: . ..=-..-*=::-
:--::...-=#+%@@@@ ... =+==:-
::-- --*@%:+@@+ . -+=:==--:=
-:==:--%@@@@@=. =*@@#=-=---:
-::-=-==#+:- .-=+*+*+=::=-:-
:-=:::--:=:-===--==:-=---:-=

With 150 and 100 steps left there is nothing to see. By 50 a rough layout has appeared, by 20 it is clearly a digit, and the remaining steps only tidy it. Almost all of the visible decision happens in a narrow band in the middle of the chain, which is the observation behind every fast sampler: most of those 200 network calls are spent on very little.

Week 13 · Day 7 of 7

Which One, and What It Costs

Coverage and compute side by side, and the checklist

By 1894 words

Two methods, the same dataset, the same judge from day 3. Here is what each one bought.

Coverage

The same count as day 3, over the same number of samples, so the entropies are directly comparable.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
class Block(nn.Module):
"""One residual block, told at every level how noisy the input is."""
def __init__(self, h):
super().__init__()
self.norm = nn.LayerNorm(h)
self.tproj = nn.Linear(h, h)
self.net = nn.Sequential(nn.SiLU(), nn.Linear(h, h),
nn.SiLU(), nn.Linear(h, h))

def forward(self, x, temb):
return x + self.net(self.norm(x) + self.tproj(temb))

class Denoiser(nn.Module):
"""Given a noisy image and how noisy it is, predict the noise."""
def __init__(self, hidden=768, blocks=3):
super().__init__()
self.inp = nn.Linear(784, hidden)
self.temb = nn.Embedding(T, hidden)
self.blocks = nn.ModuleList([Block(hidden) for _ in range(blocks)])
self.out = nn.Sequential(nn.LayerNorm(hidden), nn.SiLU(),
nn.Linear(hidden, 784))

def forward(self, x, t):
temb = self.temb(t)
h = self.inp(x)
for block in self.blocks:
h = block(h, temb)
return self.out(h)

def add_noise(x0, t, eps):
a = abar[t].reshape(-1, 1)
return a.sqrt() * x0 + (1 - a).sqrt() * eps

def train_diffusion(epochs=200, seed=0, trace_every=0):
torch.manual_seed(seed)
net = Denoiser()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
for ep in range(1, epochs + 1):
total = n = 0.0
for xb, _ in loader:
x0 = xb.flatten(1)
t = torch.randint(0, T, (len(x0),))
eps = torch.randn_like(x0)
loss = ((net(add_noise(x0, t, eps), t) - eps) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
sched.step()
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
net.eval()
return net

@torch.no_grad()
def sample(net, n=8, seed=3, keep=()):
torch.manual_seed(seed)
x = torch.randn(n, 784)
frames = {}
for step in range(T - 1, -1, -1):
t = torch.full((n,), step, dtype=torch.long)
# what the network thinks the clean image is, from here
x0 = ((x - (1 - abar[step]).sqrt() * net(x, t))
/ abar[step].sqrt()).clamp(-1, 1)
# the posterior mean: a small step from x towards that guess
mean = (abar_prev[step].sqrt() * beta[step] * x0
+ alpha[step].sqrt() * (1 - abar_prev[step]) * x
) / (1 - abar[step])
x = mean if step == 0 else mean + beta[step].sqrt() * torch.randn_like(x)
if step in keep:
frames[step] = x.clone()
return (x, frames) if keep else x
def train_judge():
"""A plain digit classifier, used only to score what the GAN makes."""
torch.manual_seed(1)
clf = nn.Sequential(nn.Flatten(),
nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 10))
opt = torch.optim.AdamW(clf.parameters(), lr=1e-3)
lf = nn.CrossEntropyLoss()
for _ in range(12):
for xb, yb in loader:
opt.zero_grad()
lf(clf(xb), yb).backward()
opt.step()
clf.eval()
test = Subset(datasets.MNIST('data', train=False, download=True,
transform=tf), range(2000))
tl = DataLoader(test, batch_size=512)
right = seen = 0
with torch.no_grad():
for xb, yb in tl:
right += (clf(xb).argmax(1) == yb).sum().item()
seen += len(yb)
return clf, right / seen

def coverage(G, clf, n=2000):
with torch.no_grad():
imgs = G(torch.randn(n, G.zdim)).reshape(-1, 1, 28, 28)
pred = clf(imgs).argmax(1)
counts = torch.bincount(pred, minlength=10).float()
frac = counts / counts.sum()
ent = -(frac[frac > 0] * frac[frac > 0].log()).sum().item()
return counts.long().tolist(), ent
clf, acc = train_judge()
net = Denoiser()
net.load_state_dict(torch.load('denoiser.pt'))
net.eval()
imgs = torch.cat([sample(net, n=500, seed=s) for s in [21, 22, 23, 24]])
with torch.no_grad():
pred = clf(imgs.reshape(-1, 1, 28, 28)).argmax(1)
counts = torch.bincount(pred, minlength=10).float()
frac = counts / counts.sum()
ent = -(frac[frac > 0] * frac[frac > 0].log()).sum().item()
print('judge accuracy on real digits %.4f' % acc)
print()
print('%-26s %-34s %8s' % ('model', 'digits produced out of 2000', 'entropy'))
print('%-26s %-34s %8.3f'
% ('diffusion', ' '.join('%3d' % c for c in counts.long().tolist()),
ent))
judge accuracy on real digits 0.9075

model digits produced out of 2000 entropy
diffusion 245 42 259 224 129 282 190 359 221 49 2.167

Cost

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)

RAMP = ' .:-=+*#%@'
def show(flat, step=2):
"""Print one flattened image, scaled from [-1, 1] back to characters."""
t = (flat.detach().reshape(28, 28) + 1) / 2
for row in t[::step]:
print(''.join(RAMP[min(9, max(0, int(v * 9.999)))]
for v in row.tolist()))
T = 200
beta = torch.linspace(1e-4, 0.05, T)
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)
abar_prev = torch.cat([torch.ones(1), abar[:-1]])
class Block(nn.Module):
"""One residual block, told at every level how noisy the input is."""
def __init__(self, h):
super().__init__()
self.norm = nn.LayerNorm(h)
self.tproj = nn.Linear(h, h)
self.net = nn.Sequential(nn.SiLU(), nn.Linear(h, h),
nn.SiLU(), nn.Linear(h, h))

def forward(self, x, temb):
return x + self.net(self.norm(x) + self.tproj(temb))

class Denoiser(nn.Module):
"""Given a noisy image and how noisy it is, predict the noise."""
def __init__(self, hidden=768, blocks=3):
super().__init__()
self.inp = nn.Linear(784, hidden)
self.temb = nn.Embedding(T, hidden)
self.blocks = nn.ModuleList([Block(hidden) for _ in range(blocks)])
self.out = nn.Sequential(nn.LayerNorm(hidden), nn.SiLU(),
nn.Linear(hidden, 784))

def forward(self, x, t):
temb = self.temb(t)
h = self.inp(x)
for block in self.blocks:
h = block(h, temb)
return self.out(h)

def add_noise(x0, t, eps):
a = abar[t].reshape(-1, 1)
return a.sqrt() * x0 + (1 - a).sqrt() * eps

def train_diffusion(epochs=200, seed=0, trace_every=0):
torch.manual_seed(seed)
net = Denoiser()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
for ep in range(1, epochs + 1):
total = n = 0.0
for xb, _ in loader:
x0 = xb.flatten(1)
t = torch.randint(0, T, (len(x0),))
eps = torch.randn_like(x0)
loss = ((net(add_noise(x0, t, eps), t) - eps) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
sched.step()
if trace_every and ep % trace_every == 0:
print('%6d %14.4f' % (ep, total / n))
net.eval()
return net

@torch.no_grad()
def sample(net, n=8, seed=3, keep=()):
torch.manual_seed(seed)
x = torch.randn(n, 784)
frames = {}
for step in range(T - 1, -1, -1):
t = torch.full((n,), step, dtype=torch.long)
# what the network thinks the clean image is, from here
x0 = ((x - (1 - abar[step]).sqrt() * net(x, t))
/ abar[step].sqrt()).clamp(-1, 1)
# the posterior mean: a small step from x towards that guess
mean = (abar_prev[step].sqrt() * beta[step] * x0
+ alpha[step].sqrt() * (1 - abar_prev[step]) * x
) / (1 - abar[step])
x = mean if step == 0 else mean + beta[step].sqrt() * torch.randn_like(x)
if step in keep:
frames[step] = x.clone()
return (x, frames) if keep else x
import time
net = Denoiser()
net.load_state_dict(torch.load('denoiser.pt'))
net.eval()
start = time.perf_counter()
_ = sample(net, n=64, seed=1)
diff_t = time.perf_counter() - start
params = sum(p.numel() for p in net.parameters())
print('the denoiser holds %d parameters' % params)
print('one reverse pass runs the network %d times' % T)
print('64 images took %.1f seconds on this machine' % diff_t)
print('a generator produces 64 images in a single forward pass')
the denoiser holds 6680848 parameters
one reverse pass runs the network 200 times
64 images took 2.2 seconds on this machine
a generator produces 64 images in a single forward pass

Timings are the one number here that is not portable

Everything else on this page is arithmetic that will reproduce. A wall-clock reading depends on the machine, the thread count and what else is running. Take the ratio as the finding, not the seconds.

How to choose

GANDiffusion
Training stabilityAn equilibrium between two networks, with failure modes that do not appear in the lossOne ordinary regression, which either converges or does not
Loss you can readNoYes, against a known baseline of 1.0
Sampling costOne forward passOne forward pass per step, 200 of them here
Compute to reach these samples30 epochs200 epochs, and still improving
Mode coverage here, out of 2.3032.2752.167
What goes wrongSilent collapse onto part of the dataBlur and residual noise, visible immediately

Two things in that count are worth naming. Diffusion came out slightly lower than the GAN, which is the opposite of the usual claim, and the honest reading is that the difference is small, the judge is imperfect, and the diffusion model was still improving when the budget ran out. More interesting is where it lost: ones and nines arrive about a fifth as often as they should. A one is mostly background, so a model that leaves a haze everywhere blurs the thing that makes a one a one. The failure has a cause, and it is the same residual noise day 6 pointed at.

What the table does support is the shape of the trade. The GAN reached comparable coverage in a seventh of the epochs, and the diffusion model reached it without any of the instability, without a second network, and with a loss that could be read against a known baseline the whole way.

That is roughly why the field moved. At small scale a GAN is quick and often good enough. As models grow, an adversarial equilibrium becomes steadily harder to hold and a regression stays a regression, and sampling cost turns out to be the easier problem to attack, which is what the fast samplers did.

The checklist

  • Before anything else, get a baseline number for your loss. For a denoiser it is 1.0. For a GAN there is not one, which is itself the thing to plan around.
  • Never judge a generative model by three pictures. Count what a classifier calls a few thousand samples and read the spread.
  • For a GAN, watch discriminator accuracy rather than either loss, and treat 1.0 or 0.5 as the two ways it is going wrong.
  • Use the non-saturating generator loss. The alternative delivers around a seventieth of the gradient when it matters most.
  • For diffusion, print abar at your final step and confirm it is near zero before you spend anything on training.
  • Condition both networks on the label, not just the generator, or the label will be ignored.
  • Expect conditioning to be partial. Measure agreement against a judge and know the number before you promise anybody control.

Where this leaves the course

Weeks 12 and 13 have covered generating data without labels. The next step is to stop treating the generated thing as the product at all, and use the same label-free training to learn representations, which is where the largest models actually get their general ability from.