🎓 A1Learner
🔍

PyTorch The Diffusion Idea

A diffusion model generates data by learning to remove noise step by step. Training data takes the forward road first, where small noise steps turn a data point into almost pure noise.

One noising step

One step shrinks the point by the square root of one minus beta. It also adds beta's share of one fixed noise draw. With beta at 0.10, the point moves only a little.

torch.manual_seed(0)
point = torch.tensor([1.5, -0.5])
noise = torch.randn(2)
beta = 0.10
noisy = (1 - beta) ** 0.5 * point + beta ** 0.5 * noise
print([round(v, 2) for v in noisy.tolist()])

The signal share shrinks

alpha_bars is the running product of one minus beta over all steps so far. The point at any step is scaled by the square root of that share.

betas = torch.linspace(0.10, 0.50, 10)
alpha_bars = torch.cumprod(1.0 - betas, dim=0)
print(round(alpha_bars[0].item(), 2))
print(round(alpha_bars[-1].item(), 4))

Every step from the same start

Each row of the closing table starts again from the original point. The formula jumps straight to any step with that step's alpha_bars value.

for t in range(steps):
    ab = alpha_bars[t]
    noisy = ab.sqrt() * point + (1.0 - ab).sqrt() * noise

Example

import torch

torch.manual_seed(0)
point = torch.tensor([1.5, -0.5])
noise = torch.randn(2)

steps = 10
betas = torch.linspace(0.10, 0.50, steps)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)

print("Start point:", [round(v, 2) for v in point.tolist()])
print("Step  Beta  Point after this step")
for t in range(steps):
    ab = alpha_bars[t]
    noisy = ab.sqrt() * point + (1.0 - ab).sqrt() * noise
    print(f"{t + 1:>4}  {betas[t].item():.2f}  "
          f"({noisy[0].item():.2f}, {noisy[1].item():.2f})")
print("Signal left at the end:", round(alpha_bars[-1].sqrt().item(), 2))

Key points

  • Beta: the noise share added at one step, rising from 0.10 to 0.50 here
  • Use torch.cumprod to build alpha_bars from the beta list
  • At the last step the point keeps a small share of its size and noise takes the rest
  • A model that removes one step of noise can walk the road backwards

In this example, the start point is (1.5, -0.5), and one seeded noise draw gives (1.54, -0.29). The schedule holds ten betas rising from 0.10 to 0.50, and each row of the table prints its beta. The first row matches the snippet, landing at (1.91, -0.57). The point then drifts as the noise share grows row by row. By the last row it sits at (1.75, -0.37). That last point is mostly the noise draw, since the noise share there is near one. The signal left prints as 0.15, the square root of the last alpha_bars value. The next lesson trains a model to reverse this road.

Note

The schedule here uses ten large steps, while real diffusion models often use many small ones

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does alpha_bars hold at the final step of a noising schedule?