TinyDiffusion: four 2D checkpoints

README.md

@@ -0,0 +1,148 @@
+---
+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.

config.json

@@ -0,0 +1,56 @@
+{
+  "model_type": "tinydiffusion",
+  "description": "DDPM on 2D points. Not images. The denoiser is an MLP, never a UNet.",
+  "architecture": {
+    "timestep_embedding": {
+      "type": "nn.Embedding",
+      "num_embeddings": 100,
+      "dim": 32
+    },
+    "layers": [
+      "Linear(2+32, 128)",
+      "SiLU",
+      "Linear(128, 128)",
+      "SiLU",
+      "Linear(128, 2)"
+    ],
+    "note": "Output is the predicted 2D noise, raw. No activation on the final layer.",
+    "params": 24450
+  },
+  "diffusion": {
+    "T": 100,
+    "objective": "epsilon (predict the noise that was added)",
+    "sampler": "ancestral DDPM, no noise added on the final step",
+    "schedules": {
+      "linear": {
+        "betas": "torch.linspace(1e-4, 0.10, 100)",
+        "alpha_bar_T": 0.0056,
+        "warning": "beta_T is 0.10, NOT the usual 0.02. The 0.02 default is tuned for T=1000; at T=100 it stalls at alpha_bar_T=0.36 and never reaches pure noise."
+      },
+      "cosine": {
+        "betas": "Nichol & Dhariwal: alpha_bar from a cosine curve, betas derived",
+        "alpha_bar_T": 0.0
+      }
+    }
+  },
+  "checkpoints": {
+    "dot": {
+      "schedule": "linear"
+    },
+    "line": {
+      "schedule": "linear"
+    },
+    "moons-linear": {
+      "schedule": "linear"
+    },
+    "moons-cosine": {
+      "schedule": "cosine"
+    }
+  },
+  "data_space": {
+    "dot": "Gaussian, centre (2.0, 2.0), std 0.1. NOT normalised.",
+    "line": "y = x, u ~ Uniform[-2, 2], perpendicular Gaussian fuzz std 0.1. NOT normalised.",
+    "moons": "sklearn make_moons(noise=0.05), normalised to zero mean / unit std per axis. Generated points come out in THAT space, not raw make_moons space."
+  },
+  "license": "mit"
+}

dot.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88ea4e7688b2068ae686b4124acd28ab7b9a0a54e5959ca9eab5c1674dd0deff
+size 98336

line.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc49fb34be6d674bc71f5a1f250fe48c18ca21e03756757616d190e7f55f42e8
+size 98336

moons-cosine.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:407c996a9de30284fa8f531ce05b3a4546bfcf15347d1966abd844205bf8e6d6
+size 98336

moons-linear.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:866b9e18703f2691546d5a2b9565ee6cef46aaf6591b794bcb75e06f98bc13bf
+size 98336