---
license: mit
library_name: pytorch
language:
  - en
tags:
  - pytorch
  - safetensors
  - diffusion
  - ddpm
  - tiny-model
  - educational
  - toy-dataset
---

# TinyDiffusion

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

Sibling of [TinyNet](https://huggingface.co/shahfazal/tinynet-relu-v1), 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.

```python
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.