🎓 A1Learner
🔍

PyTorch A Tiny Diffusion Model

A tiny diffusion model is a small network that predicts the noise inside a noised point at a given step. Removing the predicted noise step by step turns pure noise back into data-like points.

A random step per row

Each training row draws its own step number from 0 to 19. The model must handle every noise level, not one fixed level.

torch.manual_seed(0)
print(torch.randint(0, 20, (4,)).tolist())

The model predicts the noise

The target for the model is the exact noise that was added. Mean squared error scores the prediction against that noise.

ab = alpha_bars[t]
noisy = ab.sqrt().unsqueeze(1) * x0 + (1.0 - ab).sqrt().unsqueeze(1) * noise
loss = nn.functional.mse_loss(model(noisy, t), noise)

Walking backwards

Sampling starts from pure noise and removes one predicted step at a time. Fresh noise is added back at every step except the last one.

x = (x - beta / (1.0 - ab).sqrt() * noise_hat) / alpha.sqrt()
if ti > 0:
    x = x + beta.sqrt() * torch.randn(5, 2)

Example

import torch
from torch import nn

steps_total = 20
betas = torch.linspace(0.05, 0.40, steps_total)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)

torch.manual_seed(0)
per_blob = 256
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, 0.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 0.0])
data = torch.cat([blob_a, blob_b])


class Denoiser(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(3, 64), nn.ReLU(),
            nn.Linear(64, 64), nn.ReLU(),
            nn.Linear(64, 2))

    def forward(self, x, t):
        level = (t.float() / steps_total).unsqueeze(1)
        return self.net(torch.cat([x, level], dim=1))


torch.manual_seed(3)
model = Denoiser()
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
batch = 256
for step in range(400):
    rows = torch.randint(0, data.shape[0], (batch,))
    x0 = data[rows]
    t = torch.randint(0, steps_total, (batch,))
    noise = torch.randn(batch, 2)
    ab = alpha_bars[t]
    noisy = ab.sqrt().unsqueeze(1) * x0 + (1.0 - ab).sqrt().unsqueeze(1) * noise
    loss = nn.functional.mse_loss(model(noisy, t), noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print("Final noise-prediction loss:", round(loss.item(), 4))

torch.manual_seed(42)
x = torch.randn(5, 2)
with torch.no_grad():
    for ti in reversed(range(steps_total)):
        t = torch.full((5,), ti, dtype=torch.long)
        noise_hat = model(x, t)
        beta = betas[ti]
        alpha = alphas[ti]
        ab = alpha_bars[ti]
        x = (x - beta / (1.0 - ab).sqrt() * noise_hat) / alpha.sqrt()
        if ti > 0:
            x = x + beta.sqrt() * torch.randn(5, 2)

print("Generated points:")
for row in x.tolist():
    print(f"({row[0]:.2f}, {row[1]:.2f})")
print("Real points:")
for i in [0, 100, 256, 356, 511]:
    row = data[i].tolist()
    print(f"({row[0]:.2f}, {row[1]:.2f})")

Key points

  • Denoiser: the network that predicts the noise in a noised point
  • Training covers all 20 steps because each row draws a random step
  • The generated points land near the two blob centres at (-2, 0) and (2, 0)
  • Five generated points print next to five real points in the closing program

In this example, the schedule holds 20 steps with betas rising from 0.05 to 0.40. The data is 512 points in two blobs around (-2, 0) and (2, 0). The denoiser is a small MLP that takes the two point coordinates and the step fraction as its three inputs. Each of the 400 rounds trains on a batch of 256 rows. Every row draws a random step and fresh noise. The final noise-prediction loss prints as 0.3036. Sampling then starts five points as pure noise under a seed and walks all 20 steps backwards. The five generated points print next to five real points. Most generated points land near a blob centre. Two land in the inner tails between the centre and a blob.

Note

Twenty steps and 400 training rounds keep this model tiny, so some generated points land between the blobs

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a diffusion denoiser learn to output for a noised point?