shahfazal/tiny-diffusion

TinyDiffusion

DDPM on 2D points. Not images. The denoiser is a 24,450-parameter MLP, trained on CPU in seconds.

Sibling of TinyNet, and built for the same reason: strip everything that isn't the mechanism. The forward noising, the epsilon objective, the ancestral sampler and the time conditioning are unchanged from the ones in Stable Diffusion. Only the denoiser changes with the data type, and a UNet would be the largest thing on screen while having nothing to do with diffusion. So: an MLP, and 2D points, and the entire data distribution fits on one scatter plot.

• Code: https://github.com/shahfazal/tiny-diffusion
• Write-up: https://shahfazal.com/posts/tiny-diffusion/

Checkpoints

file	distribution	schedule	steps	train loss
dot.safetensors	tight Gaussian at (2, 2), σ=0.1	linear	3,000	0.066
line.safetensors	y = x over [-2, 2], σ=0.1	linear	3,000	0.258
moons-linear.safetensors	two crescents (v0.1)	linear	15,000	0.306
moons-cosine.safetensors	two crescents (v0.1.1)	cosine	15,000	0.413

Unconditional, one distribution per checkpoint. moons-cosine is the one to use: better samples than moons-linear despite the worse loss (see below).

Three things that will bite you

Everything else is standard DDPM. These are not.

1. beta_T = 0.10, not 0.02. The canonical 1e-4 → 0.02 range is tuned for T=1000. At T=100 it injects a tenth of the total noise and stalls at alpha_bar_T = 0.36, so the forward process never reaches N(0, I) while sampling still starts there. With beta_T = 0.10, alpha_bar_T = 0.0056.

2. Match the schedule to the checkpoint. Loading moons-cosine and sampling it on the linear schedule fails silently, producing plausible-looking garbage. See config.json.

3. moons lives in normalised space. It was trained on make_moons(noise=0.05) standardised to zero mean / unit std per axis, so generated points come out in that space, not raw make_moons coordinates. dot and line are unnormalised.

Usage

No diffusers, no transformers, no from_pretrained — there is no library that knows this architecture, so the module definition below is the API.

import math, torch, torch.nn as nn
from safetensors.torch import load_file

T = 100

class TinyDiffusion(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed  = nn.Embedding(T, 32)
        self.layer1 = nn.Linear(2 + 32, 128)
        self.layer2 = nn.Linear(128, 128)
        self.layer3 = nn.Linear(128, 2)
        self.act    = nn.SiLU()

    def forward(self, x, t):
        h = torch.cat([x, self.embed(t)], dim=1)
        h = self.act(self.layer1(h))
        h = self.act(self.layer2(h))
        return self.layer3(h)

def linear_betas():
    return torch.linspace(1e-4, 0.10, T)

def cosine_betas(s=0.008):                       # Nichol & Dhariwal
    t  = torch.linspace(0, T, T + 1) / T
    ab = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    ab = ab / ab[0]
    return (1 - ab[1:] / ab[:-1]).clamp(max=0.999)

@torch.no_grad()
def sample(model, betas, n=512):
    alphas, ab = 1 - betas, torch.cumprod(1 - betas, dim=0)
    x = torch.randn(n, 2)
    for s in reversed(range(T)):
        t = torch.full((n,), s, dtype=torch.long)
        eps = model(x, t)
        mean = (1 / alphas[s].sqrt()) * (x - (betas[s] / (1 - ab[s]).sqrt()) * eps)
        x = mean + betas[s].sqrt() * torch.randn_like(x) if s > 0 else mean
    return x

model = TinyDiffusion().eval()
model.load_state_dict(load_file("moons-cosine.safetensors"))
pts = sample(model, cosine_betas())              # (512, 2)

Architecture

Embedding(100, 32) → concat with the 2D point → Linear(34, 128) → SiLU → Linear(128, 128) → SiLU → Linear(128, 2). Output is epsilon, raw. 24,450 params, a third of which were the timestep embedding before the width went to 128.

Trained with Adam (lr 1e-3, batch 256) on freshly sampled data every step — the dataset is a generator, not an array, so memorisation is off the table and extra capacity is free.

make_moons returns which crescent each point came from. It's discarded: the model has no idea there are two modes, which is what makes mode-dropping possible at all. Using it would be conditioning.

The loss ranks these backwards

Why moons-cosine ships despite the worse loss. Three seeds each:

schedule	train loss	off-manifold	minority moon
real	–	0.014	50%
linear	0.330	0.046	47.8%
cosine	0.404	0.038	48.8%

(off-manifold = mean distance from each generated point to the nearest real one.)

Cosine: 18% better samples, 22% worse loss. Every seed, no overlap. Select on the loss curve and you select linear, which is worse.

The floor is set by how much of x0 is still recoverable from x_t. Cosine keeps signal alive longer (SNR crosses 1 at t=49 vs t=37 for linear), so mid-trajectory the model is asked a harder question and scores worse on it. The problem moved, not the quality. The same effect makes the loss incomparable across the ladder: dot 0.06, line 0.27, moons 0.35 — all three succeeded.

Limitations

• 2D points. This will not generate images, by design.
• Unconditional. You can sample the moons distribution; you cannot ask for the left crescent.
• Samples are mildly over-dispersed (dot: σ 0.12 generated vs 0.10 real) — reverse-step error accumulating over 100 steps.

License

MIT.