Diffusion models from scratch

Open in Colab View on GitHub DownloadGPU

In this notebook we build a diffusion model by hand: the noise schedule, the loss, the network and three samplers, each in a few lines of JAX and Flax. Only the training loop comes from Dew, as a custom Objective handed to its Trainer. Notebook 02 trains the same kind of model with Dew’s built-in pieces, so you can compare the two.

I have tried to explain the concepts as simply as I can, but some of them need a little mathematics. Where the notebook simplifies, it links the paper that does it properly. The model trains on Oxford Flowers at 64x64 and expects one NVIDIA GPU.

Diffusion and score-based models

Several families of generative models came before diffusion: variational autoencoders, generative adversarial networks, autoregressive models and normalizing flows. GANs came closest in quality, but they were painful and unstable to train. Autoregressive image models were slow, flows were expensive, and VAEs produced blurry samples.

The idea behind diffusion is that generating an image in one step is too hard, so we generate it over many small steps. We treat the data as samples from a complicated distribution, start from a simple one such as a standard normal, and learn to move a sample from the simple distribution towards the data a little at a time. Starting from Gaussian noise, each step removes some noise, until what is left looks like it came from the data.

Two lines of work arrived at this independently. Score-based models learn the gradient of the log-density of the data, $\nabla_x \log p(x)$ , which points from a sample towards more likely samples, and take small steps along it. Diffusion models (Ho et al., 2020) learn to denoise a sample a little at each discrete time step. The two are equivalent: they differ in how they write the loss, and Song et al. (2021) put both under one framework of stochastic differential equations.

The diffusion process

The forward process gradually adds noise to a data sample $x_0$ . As a stochastic differential equation (SDE):

$dx_t = f(t, x_t)\,dt + g(t)\,dW_t$

$f$ is the drift, the deterministic part; $g$ is the diffusion coefficient, which scales the increments $dW_t$ of a Wiener process (Brownian motion). Every forward SDE has a reverse-time SDE (Anderson, 1982) that runs from noise back to data:

$dx_t = \left[f(t, x_t) - g(t)^2\,\nabla_{x_t} \log p_t(x_t)\right]dt + g(t)\,d\bar{W}_t$

The only unknown in it is the score, $\nabla_{x_t} \log p_t(x_t)$ , and that is what the network learns. The reverse SDE also has a probability flow ODE with the same marginal densities at every time:

$\frac{dx_t}{dt} = f(t, x_t) - \frac{1}{2} g(t)^2\,\nabla_{x_t} \log p_t(x_t)$

Solving the SDE adds fresh noise at every step; ancestral samplers such as DDPM and Euler ancestral do that. Solving the ODE is deterministic; DDIM, Euler and Heun do that. We write samplers of both kinds at the end of this notebook.

In practice we never simulate the forward SDE step by step. For the processes used here, a noisy sample at any time has a closed form:

$x_t = \alpha_t x_0 + \sigma_t \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I)$

$\alpha_t$ is the signal rate and $\sigma_t$ the noise rate. A noise schedule is the choice of these two functions of $t$ .

%pip install -q "dew-ml[cuda13,streaming] @ git+https://github.com/AshishKumar4/dew"

Settings

Every number the notebook uses is here. The images are 64x64, a batch holds 64 of them, and we train for 6,000 steps, as notebook 02 does.

import os

IMAGE_SIZE = 64
BATCH_SIZE = 64
STEPS = 6000
LEARNING_RATE = 2e-4
TIMESTEPS = 1000
FEATURES = (64, 128, 256)
SAMPLES = 32
RUN_DIR = "runs/01-from-scratch"
SEED = 0

# The tutorial test sets DEW_TUTORIAL_SMOKE=1 to run every cell in minutes on a CPU.
if os.environ.get("DEW_TUTORIAL_SMOKE") == "1":
    BATCH_SIZE, STEPS, FEATURES, SAMPLES = 8, 4, (16, 32, 64), 4

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

print(jax.devices())

from dew.artifacts import uint8_pixels


def show_images(images, columns=8, title=None):
    images = np.asarray(images)
    if images.dtype != np.uint8:
        images = uint8_pixels(images)
    rows = (len(images) + columns - 1) // columns
    figure, axes = plt.subplots(rows, columns, figsize=(columns * 1.2, rows * 1.2), squeeze=False)
    for axis in np.ravel(axes):
        axis.axis("off")
    for axis, image in zip(np.ravel(axes), images):
        axis.imshow(image)
    if title:
        figure.suptitle(title)
    plt.show()

The data

Oxford Flowers has about 8,000 photos of 102 kinds of flower. HFImages reads a copy from the Hugging Face Hub, resizes each photo to 64x64 and flips half of them at random. The first run downloads about 270 MB. The data arrives as uint8 pixels, and the model works in [-1, 1].

from dew.data import DataPartition, HFImages, Loading

data = HFImages(
    name="pranked03/flowers-blip-captions",
    image_size=IMAGE_SIZE,
    val_batches=0,
    loading=Loading(workers=0, threads=16, read_buffer=64),
).load(batch=BATCH_SIZE)

batch = next(iter(data.train(DataPartition())))
images = batch["image"].astype(np.float32) / 127.5 - 1
print(images.shape, images.min(), images.max())
show_images(images[:16])

A noise schedule

A schedule gives the signal rate $\alpha_t$ and the noise rate $\sigma_t$ for each step $t$ . We use variance-preserving schedules, where $\alpha_t^2 + \sigma_t^2 = 1$ at every step, so the variance of $x_t$ stays at that of the data as noise replaces signal. As $t$ grows, $\alpha_t$ falls towards 0 and $\sigma_t$ rises towards 1.

The original DDPM paper wrote the forward process one step at a time, as a Markov chain:

$q(x_t \mid x_{t-1}) = \mathcal{N}\left(x_t;\ \sqrt{1-\beta_t}\,x_{t-1},\ \beta_t I\right)$

that is, $x_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\epsilon_t$ . Chaining these steps gives the closed form above, with

$\alpha_t = \prod_{s \le t} \sqrt{1-\beta_s}, \qquad \sigma_t = \sqrt{1-\alpha_t^2}$

(The DDPM paper uses $\alpha$ for $1-\beta_t$ and $\bar\alpha$ for its product, so its $\bar\alpha_t$ is our $\alpha_t^2$ . Be careful with the notation when you read it.)

DDPM used a linear schedule for $\beta_t$ . The cosine schedule of Nichol and Dhariwal (2021) trains better; it chooses the betas so that $\alpha_t^2$ follows a squared cosine.

def cosine_beta_schedule(timesteps, s=0.008, max_beta=0.999):
    t = np.linspace(0, 1, timesteps + 1, dtype=np.float64)
    alpha_bar = np.cos((t + s) / (1 + s) * np.pi / 2) ** 2
    alpha_bar = alpha_bar / alpha_bar[0]
    return np.clip(1 - alpha_bar[1:] / alpha_bar[:-1], 0, max_beta)


betas = cosine_beta_schedule(TIMESTEPS)
signal = np.sqrt(np.cumprod(1 - betas))  # alpha_t
noise = np.sqrt(1 - signal**2)           # sigma_t

steps = np.arange(TIMESTEPS)
plt.figure(figsize=(6, 3))
plt.plot(steps, signal, label="signal rate")
plt.plot(steps, noise, label="noise rate")
plt.plot(steps, signal**2 + noise**2, "--", label="variance")
plt.xlabel("step t")
plt.legend()
plt.show()

The same schedule can be written without betas, directly as $\alpha(t) = \cos\left(\frac{\pi t}{2T}\right)$ and $\sigma(t) = \sin\left(\frac{\pi t}{2T}\right)$ . It is variance preserving because $\cos^2 + \sin^2 = 1$ , and it is defined for any real $t$ , which is how continuous-time formulations use it. The two curves differ only by the small offset $s$ near $t = 0$ .

Dew has the tabulated version as dew.diffusion.CosineNoiseScheduler. Its rates should equal ours to float32 precision:

from dew.diffusion import CosineNoiseScheduler

continuous_signal = np.cos(np.pi * steps / (2 * TIMESTEPS))
print("largest gap to the continuous cosine:", float(np.abs(continuous_signal - signal).max()))

dew_signal, dew_noise = CosineNoiseScheduler(TIMESTEPS).rates(steps)
print("largest gap to Dew's table:", float(np.abs(np.ravel(dew_signal) - signal).max()),
      float(np.abs(np.ravel(dew_noise) - noise).max()))

Adding and removing noise

Here are eight flowers at increasing noise levels. Because $x_t = \alpha_t x_0 + \sigma_t \epsilon$ , anyone who knows $\epsilon$ can take it back out exactly: $x_0 = (x_t - \sigma_t \epsilon) / \alpha_t$ . That is the whole trick. If a network can guess $\epsilon$ from $x_t$ alone, it can denoise.

eps = np.asarray(jax.random.normal(jax.random.key(SEED), images[:8].shape))
levels = [0, 100, 250, 500, 750, 999]
noisy = {t: signal[t] * images[:8] + noise[t] * eps for t in levels}
show_images(np.concatenate([np.clip(noisy[t], -1, 1) for t in levels]),
            title="One row per step: t = 0, 100, 250, 500, 750 and 999")

recovered = (noisy[500] - noise[500] * eps) / signal[500]
print("largest error after removing the known noise at t = 500:", float(np.abs(recovered - images[:8]).max()))

The network

The network $\epsilon_\theta(x_t, t)$ predicts the noise in a noisy image. Predicting $\epsilon$ is one parameterization; predicting the clean image or a velocity also works, and the early diffusion papers found the noise the easiest target. The score follows from it, $\nabla_{x_t} \log p_t(x_t) \approx -\epsilon_\theta(x_t, t) / \sigma_t$ .

We use a small U-Net. An encoder halves the resolution twice while it widens the channels, and a decoder mirrors it, with skip connections that carry each encoder level’s features to the decoder level of the same size. Each level has residual blocks with group normalization, and the two lowest levels add self-attention over all positions. The time step enters every residual block through a sinusoidal embedding, the same kind of encoding transformers use for positions.

import flax.linen as nn


def time_embedding(t, dim):
    half = dim // 2
    freqs = jnp.exp(-jnp.log(10_000.0) * jnp.arange(half) / (half - 1))
    angles = t.astype(jnp.float32)[:, None] * freqs[None, :]
    return jnp.concatenate([jnp.sin(angles), jnp.cos(angles)], axis=-1)


class ResBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x, temb):
        h = nn.Conv(self.features, (3, 3))(nn.swish(nn.GroupNorm(8)(x)))
        h = h + nn.Dense(self.features)(temb)[:, None, None, :]
        h = nn.Conv(self.features, (3, 3))(nn.swish(nn.GroupNorm(8)(h)))
        if x.shape[-1] != self.features:
            x = nn.Conv(self.features, (1, 1))(x)
        return x + h


class Attention(nn.Module):
    heads: int = 4

    @nn.compact
    def __call__(self, x):
        b, h, w, c = x.shape
        y = nn.GroupNorm(8)(x).reshape(b, h * w, c)
        q, k, v = jnp.split(nn.DenseGeneral((3 * self.heads, c // self.heads))(y), 3, axis=2)
        y = jax.nn.dot_product_attention(q, k, v)
        return x + nn.DenseGeneral(c, axis=(-2, -1))(y).reshape(b, h, w, c)


class UNet(nn.Module):
    features: tuple[int, ...]

    @nn.compact
    def __call__(self, x, t):
        width = 4 * self.features[0]
        temb = nn.Dense(width)(nn.swish(nn.Dense(width)(time_embedding(t, width))))
        h = nn.Conv(self.features[0], (3, 3))(x)
        skips = [h]
        last = len(self.features) - 1
        for level, features in enumerate(self.features):
            h = ResBlock(features)(ResBlock(features)(h, temb), temb)
            if level == last:
                h = Attention()(h)
            skips.append(h)
            if level < last:
                h = nn.Conv(features, (3, 3), strides=(2, 2))(h)
        h = ResBlock(self.features[-1])(Attention()(ResBlock(self.features[-1])(h, temb)), temb)
        for level in reversed(range(len(self.features))):
            features = self.features[level]
            h = ResBlock(features)(jnp.concatenate([h, skips.pop()], axis=-1), temb)
            h = ResBlock(features)(h, temb)
            if level == last:
                h = Attention()(h)
            if level > 0:
                b, height, width_, c = h.shape
                h = jax.image.resize(h, (b, 2 * height, 2 * width_, c), "nearest")
                h = nn.Conv(self.features[level - 1], (3, 3))(h)
        h = ResBlock(self.features[0])(jnp.concatenate([h, skips.pop()], axis=-1), temb)
        return nn.Conv(3, (3, 3), kernel_init=nn.initializers.zeros)(nn.swish(nn.GroupNorm(8)(h)))


model = UNet(FEATURES)
shapes = jax.eval_shape(model.init, jax.random.key(0), jnp.zeros((1, IMAGE_SIZE, IMAGE_SIZE, 3)),
                        jnp.zeros((1,), jnp.int32))
print(f"{sum(x.size for x in jax.tree_util.tree_leaves(shapes)) / 1e6:.1f}M parameters")

The loss, as a Dew objective

Training is short to write down. Take a clean batch, pick a random step $t$ for each image, draw noise $\epsilon$ , make $x_t$ , and ask the network for $\epsilon$ :

$L = \mathbb{E}_{x_0, t, \epsilon}\left\| \epsilon_\theta(\alpha_t x_0 + \sigma_t \epsilon,\ t) - \epsilon \right\|^2$

This is the simplified DDPM loss. An Objective says how to initialize the variables and computes this loss; step.key is a fresh random key for every step. ema asks the trainer to keep an exponential moving average of the weights, which usually samples better than the live weights.

import optax
from dew import Aux, EMASpec, Field, InputSpec, Objective


class NoisePrediction(Objective):
    inputs = InputSpec(Field("image", (IMAGE_SIZE, IMAGE_SIZE, 3)))
    ema = EMASpec(decay=optax.constant_schedule(0.999))

    def __init__(self, model):
        self.model = model

    def init(self, key, variables=None):
        return self.model.init(key, jnp.zeros((1, IMAGE_SIZE, IMAGE_SIZE, 3)), jnp.zeros((1,), jnp.int32))

    def loss(self, variables, batch, step):
        x0 = batch["image"].astype(jnp.float32) / 127.5 - 1
        t_key, noise_key = jax.random.split(step.key)
        t = jax.random.randint(t_key, (x0.shape[0],), 0, TIMESTEPS)
        eps = jax.random.normal(noise_key, x0.shape)
        a = jnp.asarray(signal, jnp.float32)[t][:, None, None, None]
        s = jnp.asarray(noise, jnp.float32)[t][:, None, None, None]
        loss = jnp.mean((self.model.apply(variables, a * x0 + s * eps, t) - eps) ** 2)
        return loss, Aux(metrics={"mse": loss})

Training

Trainer computes the gradients, applies AdamW, updates the moving average and writes a checkpoint at the end. LocalTracker records the logged numbers so we can plot them.

from dew import Checkpoints, LocalTracker, Trainer

objective = NoisePrediction(model)
trainer = Trainer(objective, optax.adamw(LEARNING_RATE), key=jax.random.key(SEED),
                  checkpoints=Checkpoints(RUN_DIR), tracker=LocalTracker(f"{RUN_DIR}/tracking"))
state = trainer.fit(data, steps=STEPS, log_every=max(1, STEPS // 12), checkpoint_every=STEPS)

import json

rows = [json.loads(line) for line in open(f"{RUN_DIR}/tracking/scalars.jsonl")]
rows = [row for row in rows if "train/loss" in row["scalars"]]
plt.figure(figsize=(6, 3))
plt.plot([row["step"] for row in rows], [row["scalars"]["train/loss"] for row in rows], marker=".")
plt.xlabel("step")
plt.ylabel("loss")
plt.show()

Sampling

To generate, we start from pure noise at the last step and walk the steps back to zero. Every sampler below takes the same ingredients: the network’s noise estimate $\hat\epsilon$ at the current step, the clean image it implies, $\hat x_0 = (x_t - \sigma_t \hat\epsilon) / \alpha_t$ , and a rule for the next point. Each walk runs in one compiled jax.lax.scan, with the moving-average weights.

DDIM

The simplest rule goes to the predicted clean image, then adds back the predicted noise at the next, lower level, $x_{t'} = \alpha_{t'} \hat x_0 + \sigma_{t'} \hat\epsilon$ (Song, Meng and Ermon, 2021). It adds no new randomness, so it solves the probability flow ODE, and it is exactly Euler’s method on that ODE written in $\sigma / \alpha$ . Each step jumps straight to the next level, so the walk can skip steps; we use 100 of the 1,000.

params = state.averaged
alphas = jnp.asarray(signal, jnp.float32)
sigmas = jnp.asarray(noise, jnp.float32)


def predict_noise(params, x, t):
    return model.apply(params, x, jnp.full((x.shape[0],), t, jnp.int32))


def to_image(params, x, t):
    return jnp.clip((x - sigmas[t] * predict_noise(params, x, t)) / alphas[t], -1, 1)


@jax.jit
def ddim(params, key, grid):
    def step(x, pair):
        t, t_next = pair
        eps = predict_noise(params, x, t)
        x0 = jnp.clip((x - sigmas[t] * eps) / alphas[t], -1, 1)
        return alphas[t_next] * x0 + sigmas[t_next] * eps, None

    x = jax.random.normal(key, (SAMPLES, IMAGE_SIZE, IMAGE_SIZE, 3))
    x, _ = jax.lax.scan(step, x, (grid[:-1], grid[1:]))
    return to_image(params, x, grid[-1])


grid = jnp.linspace(TIMESTEPS - 1, 0, 100).round().astype(jnp.int32)
show_images(ddim(params, jax.random.key(1), grid), title="DDIM, 100 steps")

Euler’s method

DDIM is a numerical ODE solver in disguise, so it helps to see what such a solver does. Take something whose rate of change we know: an object under constant acceleration $a$ , starting at rest. Its speed is $v(t) = a t$ , and its position solves $\frac{dx}{dt} = a t$ , which integrates to $x(t) = \frac{1}{2} a t^2$ .

Euler’s method integrates by adding up small changes: from $x$ at time $t$ , step to $x + \frac{dx}{dt}\,\Delta t$ . With large steps the sum trails the true parabola, because each step uses the speed at the start of its interval. Higher-order methods such as Heun’s correct for that.

a, dt, horizon = 1.0, 1.0, 10.0
times = np.arange(0, horizon + dt, dt)
positions = [0.0]
for t in times[:-1]:
    positions.append(positions[-1] + a * t * dt)

fine = np.linspace(0, horizon, 200)
plt.figure(figsize=(6, 3))
plt.plot(fine, 0.5 * a * fine**2, label="exact, x = a t² / 2")
plt.plot(times, positions, "o--", label=f"Euler, step {dt}")
plt.xlabel("t")
plt.ylabel("x")
plt.legend()
plt.show()

DDPM

DDPM (Ho et al., 2020) solves the reverse SDE instead: every step removes the predicted noise and then adds fresh noise of a smaller size. Walking from step $t$ to $t-1$ , with $\beta_t$ the schedule’s per-step noise,

$x_{t-1} = \frac{1}{\sqrt{1-\beta_t}}\left(x_t - \frac{\beta_t}{\sigma_t}\,\hat\epsilon\right) + \sqrt{\tilde\beta_t}\,z, \qquad \tilde\beta_t = \frac{\sigma_{t-1}^2}{\sigma_t^2}\,\beta_t, \quad z \sim \mathcal{N}(0, I)$

$\tilde\beta_t$ is the variance of the true posterior $q(x_{t-1} \mid x_t, x_0)$ . DDPM takes all 1,000 steps, ten times the network evaluations of the DDIM walk above.

betas_j = jnp.asarray(betas, jnp.float32)


@jax.jit
def ddpm(params, key):
    def step(x, inputs):
        t, step_key = inputs
        eps = predict_noise(params, x, t)
        mean = (x - betas_j[t] / sigmas[t] * eps) / jnp.sqrt(1 - betas_j[t])
        variance = jnp.where(t > 0, sigmas[jnp.maximum(t - 1, 0)] ** 2 / sigmas[t] ** 2 * betas_j[t], 0.0)
        return mean + jnp.sqrt(variance) * jax.random.normal(step_key, x.shape), None

    start, walk = jax.random.split(key)
    x = jax.random.normal(start, (SAMPLES, IMAGE_SIZE, IMAGE_SIZE, 3))
    ts = jnp.arange(TIMESTEPS - 1, -1, -1)
    x, _ = jax.lax.scan(step, x, (ts, jax.random.split(walk, TIMESTEPS)))
    return jnp.clip(x, -1, 1)


show_images(ddpm(params, jax.random.key(2)), title="DDPM, 1,000 steps")

Heun

Heun’s method is a second-order ODE solver: take an Euler step, evaluate the slope again where it lands, and step with the average of the two slopes (Karras et al., 2022, Algorithm 1). Writing the ODE in $\tilde x = x / \alpha$ and $\tilde\sigma = \sigma / \alpha$ makes it simple, $d\tilde x / d\tilde\sigma = \hat\epsilon$ . Each step costs two network evaluations, but it is accurate enough that 20 steps, 39 evaluations in all, give images as clean as DDIM’s 100.

@jax.jit
def heun(params, key, grid):
    def step(x, pair):
        t, t_next = pair
        scaled, s, s_next = x / alphas[t], sigmas[t] / alphas[t], sigmas[t_next] / alphas[t_next]
        eps = predict_noise(params, x, t)
        euler = scaled + (s_next - s) * eps
        eps_next = predict_noise(params, euler * alphas[t_next], t_next)
        return (scaled + (s_next - s) * 0.5 * (eps + eps_next)) * alphas[t_next], None

    x = jax.random.normal(key, (SAMPLES, IMAGE_SIZE, IMAGE_SIZE, 3))
    x, _ = jax.lax.scan(step, x, (grid[:-1], grid[1:]))
    return to_image(params, x, grid[-1])


grid = jnp.linspace(TIMESTEPS - 1, 0, 20).round().astype(jnp.int32)
show_images(heun(params, jax.random.key(3), grid), title="Heun, 20 steps")

Where to go next

Everything in this notebook exists in Dew as a configured component. dew.diffusion has the schedules and the prediction transforms, including $\epsilon$ , $x_0$ , velocity and flow targets. dew.sampling has these three samplers and fourteen more, all driven by one sample call. DiffusionObjective holds the loss with its weighting and conditions. Notebook 02 trains a diffusion transformer with them, and Diffusion processes and solvers explains how they fit together.