AI for Everyone, Part 6 — Drawing Digits From Words With a 1.11M Diffusion Model

• The homework Part 5 left behind
• The core idea — destroying is easy, restoring is hard
  • Forward — adding the noise
  • The objective — predict the noise, not the image
• Conditioning — instructing the picture with a word
• Reverse — walking back 400 times
• Training log
• Results
• Summary
• 🧠 Comprehension quiz
• References

The homework Part 5 left behind

Part 5's colourisation preserved shape well but the colours came out washed out. The cause was the loss, not the model. A greyscale car could be red or blue, and since L1 asks to minimise absolute error, predicting grey — the middle of all valid answers — becomes the safest strategy.

This part's problem is also many-to-one. There are thousands of handwritten 3s that correspond to the word "three". Ask a regression loss to "predict the image of a 3 directly", as in Part 5, and you get a blurry smear: the average of every 3.

Diffusion sidesteps this head-on. Instead of producing the answer in one shot, it turns generation into gradually clearing away noise.

The core idea — destroying is easy, restoring is hard

Adding noise to a clean image until it becomes pure static is easy: one formula does it. The hard direction is the other way.

Diffusion's insight is this: break one hard problem into 400 easy ones. Making an image from static in a single step is hard, but "make this slightly cleaner" is learnable. Repeat it 400 times.

Forward — adding the noise

T = 400
beta = torch.linspace(1e-4, 0.02, T)   # how much noise to mix at each step
alpha = 1 - beta
abar = torch.cumprod(alpha, 0)         # cumulative product

abar[t] expresses "how much of the original survives up to time t". Near 1 at t=0 (the original), near 0 at t=399 (pure static).

Here comes diffusion's first trick. You do not walk the 400 steps in order — you can jump straight to any t.

t = torch.randint(0, T, (256,), device=dev)
eps = torch.randn_like(x0)                          # standard normal noise
a = abar[t][:, None, None, None]
xt = a.sqrt() * x0 + (1 - a).sqrt() * eps           # t steps of noise in one line

That single line is what makes training efficient. Each batch samples different values of t, so every timestep is learned in parallel.

The objective — predict the noise, not the image

loss = F.mse_loss(model(xt, t.float(), w), eps)

What the model predicts is not the original image x0 but the noise eps that was mixed in.

Mathematically the two are equivalent: knowing xt and t lets you compute x0 from eps and vice versa. In practice, though, predicting noise works far better.

The reason lies in the nature of the target. Ask for x0 and, when t is large and almost nothing but static remains, the model must invent an image from no information — which is exactly where Part 5's regression-to-the-mean returns. The noise eps, by contrast, is always standard normal at every t. A target with constant scale and distribution makes training stable.

Conditioning — instructing the picture with a word

The timestep and the word are added in the same vector space and injected into the middle of the U-Net.

class CondUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.temb = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 128))
        self.wemb = nn.Embedding(10, 128)          # zero~nine
        self.e1 = block(1, 48); self.e2 = block(48, 96); self.e3 = block(96, 192)
        self.cond3 = nn.Linear(128, 192)
        self.d2 = block(192 + 96, 96); self.d1 = block(96 + 48, 48)
        self.out = nn.Conv2d(48, 1, 1)

    def forward(self, x, t, w):
        half = torch.exp(-np.log(10000) * torch.arange(32, device=x.device) / 32)
        te = torch.cat([torch.sin(t[:, None] * half), torch.cos(t[:, None] * half)], 1)
        c = self.temb(te) + self.wemb(w)           # timestep + word
        s1 = self.e1(x)
        s2 = self.e2(F.max_pool2d(s1, 2))
        h = self.e3(F.max_pool2d(s2, 2)) + self.cond3(c)[:, :, None, None]
        ...

Using sines and cosines for the timestep embedding follows the same reasoning as positional encoding in transformers. Feeding the raw integer t gives the network an awkward scale, while spreading it across sine waves of different frequencies makes neighbouring timesteps share similar vectors.

The injection point — the deepest layer, the output of e3 — is deliberate too. At 8×8 the resolution is smallest and the information density per channel highest, so conditioning here propagates through the entire decoder.

The two Nones in self.cond3(c)[:, :, None, None] reshape a (batch, channel) vector into (batch, channel, 1, 1) so it broadcasts across the spatial dimensions. In other words, the condition is added identically at every position.

Reverse — walking back 400 times

@torch.no_grad()
def sample(words):
    x = torch.randn(len(words), 1, 28, 28, device=dev)      # start from pure static
    w = torch.tensor([WORDS.index(s) for s in words], device=dev)
    for ti in reversed(range(T)):                            # 399 -> 0
        t = torch.full((len(words),), ti, device=dev, dtype=torch.float32)
        eps = model(x, t, w)                                 # predict the noise
        a, ab = alpha[ti], abar[ti]
        x = (x - (1 - a) / (1 - ab).sqrt() * eps) / a.sqrt() # remove one step
        if ti > 0:
            x = x + beta[ti].sqrt() * torch.randn_like(x)    # inject fresh noise
    return ((x.clamp(-1, 1) + 1) / 2).squeeze(1).cpu().numpy()

That last injection may look perverse. Why put noise back after working to remove it?

Because it is the source of diversity. Fresh randomness at every step means the same word yields a different 3 each time — the exact opposite of Part 5, where regression got stuck on one average. Diffusion is a structure for probabilistically picking one answer among many. Only the final step (ti == 0) skips the injection, so the result comes out clean.

Training log

[    9.2s] params {"params": 1111777, "T": 400, "n": 60000}
[   10.1s] step      0 mse=1.1502
[   12.2s] step    100 mse=0.0807
[  959.7s] step  44200 mse=0.0324
[  960.6s] SUMMARY {"steps": 44215, "final_mse": 0.0320}

Within 100 steps — two seconds — MSE fell from 1.15 to 0.081. Starting near 1 is natural given the target is standard normal. It then descended gently to 0.032.

Results

Each row runs "zero" through "nine"; the two rows are different samples for the same words.

Row one is correct for all ten digits. Stroke thickness and slant look genuinely handwritten.

Row two has two errors. The "five" slot produced something closer to a 9, and the "nine" slot something closer to a 7. Eight out of ten — 80%.

Comparing the rows reveals diffusion's character. The same word produced different handwriting. The 2s differ in how sharply the curve bends; the 7s differ in whether they carry a crossbar. Nothing was averaged away — each sample is one concrete piece of handwriting. That is precisely the contrast with Part 5's drained colours.

The confused 5 and 9 are digits that resemble each other in form. With 1.11 million parameters and 16 minutes, this much confusion remains.

Summary

Three lines capture it.

Diffusion splits one hard generation problem into 400 easy denoising problems. The forward process is one formula, and training parallelises by jumping to arbitrary t.

The model predicts noise, not images. Because the target is standard normal at every t, training stays stable and the regression-to-the-mean of predicting x0 directly is avoided.

Noise re-injection during reverse creates diversity. Producing a different result each time under the same condition is the decisive difference from a regression model.

This implementation is the basic form, though. Walking all 400 steps is slow, and there is no dial for how strongly to enforce the condition. Part 11 addresses both with classifier-free guidance and DDIM — and shows, with measurements, that the two are not independent.

🧠 Comprehension quiz

References

• MNIST dataset
• Denoising Diffusion Probabilistic Models (DDPM)
• Diffusion Models Beat GANs on Image Synthesis
• Attention Is All You Need — sinusoidal positional encoding