8.4 Generative models and diffusion

A generative model learns the distribution of the data well enough to draw new examples from it. A diffusion model learns this backwards: it is trained only to remove a little noise, and by removing noise step after step it turns pure noise into a new image.

import warnings

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from sklearn import datasets
from sklearn.mixture import GaussianMixture
from sklearn.neural_network import MLPRegressor

from statly import style

style.use()
rng = np.random.default_rng(1)
warnings.filterwarnings("ignore")

digits = datasets.load_digits()  # 1,797 handwritten digits, 8 × 8 pixels each
pixels, labels = digits.data / 16, digits.target


def spiral(n: int, local: np.random.Generator) -> np.ndarray:
    """n points along a spiral, with a little noise: a shape no simple distribution has."""
    turn = np.sqrt(local.uniform(0, 1, n)) * 3 * np.pi
    return np.c_[turn * np.cos(turn), turn * np.sin(turn)] / 5 + local.normal(0, 0.05, (n, 2))


def ring(n: int, local: np.random.Generator) -> np.ndarray:
    """n points around a circle, with a little noise."""
    angle = local.uniform(0, 2 * np.pi, n)
    return 0.6 * np.c_[np.cos(angle), np.sin(angle)] + local.normal(0, 0.05, (n, 2))


# The noise schedule: at step t a share beta[t] of fresh noise is mixed in.
STEPS = 50
beta = np.linspace(1e-4, 0.2, STEPS)
alpha = 1 - beta
alpha_bar = np.cumprod(alpha)  # how much of the original is left after t steps


def add_noise(x0: np.ndarray, t: np.ndarray, noise: np.ndarray) -> np.ndarray:
    """Jump straight to step t: sqrt(ᾱ_t) · x0 + sqrt(1 − ᾱ_t) · noise."""
    return np.sqrt(alpha_bar[t])[:, None] * x0 + np.sqrt(1 - alpha_bar[t])[:, None] * noise


def time_features(t: np.ndarray) -> np.ndarray:
    """The step number as a few smooth waves, so the network can tell the steps apart."""
    waves = np.outer(t / STEPS, np.arange(1, 5)) * np.pi
    return np.c_[np.sin(waves), np.cos(waves)]


LABELS = np.eye(3)  # spiral, ring, or no label


def train_denoiser(n: int = 60_000) -> MLPRegressor:
    """A network that sees a noised point, its step and (sometimes) its label, and guesses the noise."""
    local = np.random.default_rng(0)
    shape = local.integers(0, 2, n)  # 0: spiral, 1: ring
    x0 = np.where(shape[:, None] == 0, spiral(n, local), ring(n, local))
    label = np.where(local.uniform(size=n) < 0.2, 2, shape)  # one in five without its label
    t = local.integers(0, STEPS, n)
    noise = local.normal(size=(n, 2))
    inputs = np.c_[add_noise(x0, t, noise), time_features(t), LABELS[label]]
    network = MLPRegressor((128, 128, 128), max_iter=50, batch_size=256, random_state=0)
    return network.fit(inputs, noise)


denoiser = train_denoiser()


def predict_noise(x: np.ndarray, t: int, label: int, guidance: float = 0.0) -> np.ndarray:
    """The network's noise guess; with guidance, pushed further towards the label."""
    steps = time_features(np.full(len(x), t))
    labelled = denoiser.predict(np.c_[x, steps, np.tile(LABELS[label], (len(x), 1))])
    if guidance == 0:
        return labelled
    free = denoiser.predict(np.c_[x, steps, np.tile(LABELS[2], (len(x), 1))])
    return free + (1 + guidance) * (labelled - free)


def generate(
    n: int, label: int = 2, guidance: float = 0.0, keep: tuple[int, ...] = ()
) -> dict[int, np.ndarray]:
    """Start from pure noise and remove the predicted noise step by step (keeping some snapshots)."""
    local = np.random.default_rng(7)
    x = local.normal(size=(n, 2))
    snapshots = {STEPS: x.copy()}
    for t in reversed(range(STEPS)):
        guess = predict_noise(x, t, label, guidance)
        x = (x - beta[t] / np.sqrt(1 - alpha_bar[t]) * guess) / np.sqrt(alpha[t])
        if t > 0:
            x += np.sqrt(beta[t]) * local.normal(size=x.shape)
        if t in keep:
            snapshots[t] = x.copy()
    snapshots[0] = x
    return snapshots


def scatter(ax: Axes, points: np.ndarray, title: str, color: str = style.BLUE) -> None:
    """A square scatter of 2D points on the common scale of this chapter."""
    ax.scatter(points[:, 0], points[:, 1], s=2, color=color, alpha=0.5)
    ax.set(
        xlim=(-2.3, 2.1), ylim=(-1.6, 2.1), aspect="equal", xticks=[], yticks=[], title=title
    )

8.4.1 Generating means sampling

A classifier learns $P(y \mid x)$: which label, given the image. A generative model learns the distribution of the data itself, $P(x)$, well enough to draw new $x$ from it: new images, new sentences, new molecules. The language models of Chapter 8.3 are generative models of text. This chapter turns to images, and to the method behind most of today’s image generators.

The simplest generative model is a named distribution. Fit it, then sample. For data with a complex shape, the choice of distribution decides everything:

data = spiral(2000, rng)
single = rng.multivariate_normal(data.mean(axis=0), np.cov(data.T), 2000)
mixture = GaussianMixture(30, random_state=0).fit(data).sample(2000)[0]

fig, axes = style.subplots(3, height=3.2, width=10.5)
scatter(axes[0], data, "the data", style.GREY)
scatter(axes[1], single, "one normal distribution")
scatter(axes[2], mixture, "a mixture of 30 normals")
fig.tight_layout()
plt.show()

Figure 8.4.1: Left: 2,000 points along a spiral, the data to imitate. Middle: 2,000 draws from a normal distribution fitted to them; it gets the centre and spread right and the shape completely wrong. Right: 2,000 draws from a mixture of 30 normal distributions (Chapter 2.4); the spiral appears, built from overlapping blobs.

A mixture works in two dimensions. An image of 512 × 512 pixels in colour has 786,432 dimensions, and the shapes images take in that space are far beyond any named distribution. The models below let a neural network learn the shape instead.

8.4.2 Autoencoders: squeeze and rebuild

An autoencoder is a network trained to copy its input through a narrow middle layer. The first half, the encoder, squeezes the image into a few numbers, the latent code; the second half, the decoder, rebuilds the image from them. Because the code is small, the network must learn what matters about the images. It is principal components (Chapter 7.1) made nonlinear.

def train_autoencoder(epochs: int = 300) -> list[tuple[np.ndarray, np.ndarray]]:
    """Weights of a 64-128-2-128-64 autoencoder, trained with the Adam variant of gradient descent."""
    local = np.random.default_rng(0)
    sizes = [64, 128, 2, 128, 64]
    layers = [
        (local.normal(0, np.sqrt(2 / a), (a, b)), np.zeros(b)) for a, b in zip(sizes, sizes[1:])
    ]
    params = [p for layer in layers for p in layer]
    mean = [np.zeros_like(p) for p in params]
    square = [np.zeros_like(p) for p in params]
    step = 0
    for _ in range(epochs):
        order = local.permutation(len(pixels))
        for start in range(0, len(order), 64):
            batch = pixels[order[start : start + 64]]
            outputs = run_layers(batch, layers)
            g = (outputs[-1] - batch) / len(batch)  # sigmoid output with cross-entropy loss
            grads = []
            for k in reversed(range(4)):
                grads = [outputs[k].T @ g, g.sum(axis=0)] + grads
                if k:
                    g = (
                        g @ layers[k][0].T * (outputs[k] > 0 if k != 2 else 1)
                    )  # ReLU, except the code
            step += 1
            for p, gr, m, s in zip(params, grads, mean, square):
                m[:] = 0.9 * m + 0.1 * gr
                s[:] = 0.999 * s + 0.001 * gr**2
                p -= 1e-3 * (m / (1 - 0.9**step)) / (np.sqrt(s / (1 - 0.999**step)) + 1e-8)
    return layers


def run_layers(
    x: np.ndarray, layers: list[tuple[np.ndarray, np.ndarray]], first: int = 0
) -> list[np.ndarray]:
    """The output of every layer: ReLU inside, a plain code in the middle, a sigmoid at the end."""
    outputs = [x]
    for k, (W, b) in enumerate(layers[first:], start=first):
        h = outputs[-1] @ W + b
        if k == 3:
            h = 1 / (1 + np.exp(-h))
        elif k != 1:
            h = np.maximum(h, 0)
        outputs.append(h)
    return outputs


autoencoder = train_autoencoder()
codes = run_layers(pixels, autoencoder)[2]
low, high = np.percentile(codes, 3, axis=0), np.percentile(codes, 97, axis=0)
across, up = np.linspace(low[0], high[0], 8), np.linspace(high[1], low[1], 8)

fig = plt.figure(figsize=(10.5, 4.8))
grid = fig.add_gridspec(8, 18)
ax = fig.add_subplot(grid[:, :9])
for d in range(10):
    ax.scatter(*codes[labels == d].T, s=4, alpha=0.6, color=plt.cm.tab10(d), label=str(d))
ax.legend(fontsize=7, markerscale=3, ncol=2, loc="upper left")
ax.set(xlabel="code 1", ylabel="code 2", xticks=[], yticks=[])
for i, y in enumerate(up):
    for j, x in enumerate(across):
        cell = fig.add_subplot(grid[i, 10 + j])
        cell.imshow(
            run_layers(np.array([[x, y]]), autoencoder, first=2)[-1].reshape(8, 8),
            cmap="gray_r",
            vmin=0,
            vmax=1,
        )
        cell.set_axis_off()
plt.show()

Figure 8.4.2: An autoencoder for the digits with a latent code of just two numbers (64 pixels → 128 → 2 → 128 → 64). Left: the code of every digit, coloured by the digit; digits of a kind gather, without the network ever seeing a label. Right: images the decoder makes from an evenly spaced grid of codes across the same area. Moving through the code space morphs one digit into another.

The decoder is a generator: give it a code, get an image. But which codes give good images? The codes of real digits fill an irregular region, and codes outside it decode to smudges. A variational autoencoder fixes this by training the codes to follow a standard normal distribution, so that new images come from decoding random normal draws.

8.4.3 Diffusion: destroying an image, one step at a time

Diffusion models start from the opposite end. Take a training image and add a little noise. Then a little more, for many steps, until nothing but noise is left. In symbols, step t mixes in fresh normal noise ε with weight $\beta_t$, and t steps together have a simple closed form:

$$x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon, \qquad \bar\alpha_t = \prod_{s \le t} (1 - \beta_s), \qquad \varepsilon \sim N(0, I)$$

$\bar\alpha_t$ is how much of the original is left: 1 at the start, close to 0 at the end. This forward process needs no learning; it is just averaging with noise.

SHOWN = (0, 5, 12, 20, 32, 49)
digit = pixels[0] * 2 - 1
points = spiral(2000, rng)

fig, axes = plt.subplots(
    2, len(SHOWN), figsize=(10.5, 3.9), gridspec_kw={"height_ratios": [1, 1.2]}
)
for col, t in enumerate(SHOWN):
    noise = rng.normal(size=64)
    noisy = np.sqrt(alpha_bar[t]) * digit + np.sqrt(1 - alpha_bar[t]) * noise
    axes[0, col].imshow(noisy.reshape(8, 8), cmap="gray_r", vmin=-1.5, vmax=1.5)
    axes[0, col].set_title(f"step {t + 1}: {np.sqrt(alpha_bar[t]):.0%} left", fontsize=9)
    axes[0, col].set_axis_off()
    noisy_points = add_noise(points, np.full(len(points), t), rng.normal(size=points.shape))
    scatter(axes[1, col], noisy_points, "")
    axes[1, col].set(xlim=(-3.2, 3.2), ylim=(-3.2, 3.2))
fig.tight_layout()
plt.show()

Figure 8.4.3: The forward process over 50 steps. Top: a digit, scaled to run from −1 to 1, at six steps. Bottom: the spiral data at the same steps; every point is noised on its own. The titles give the weight of the original, the square root of ᾱ. By step 50 the original has shrunk to 7% under noise of almost full size, and both look like pure normal noise.

8.4.4 Learning to remove noise

The generative part is the way back. A network is shown a noised point $x_t$ and the step t, and asked to guess the noise ε that was added. It is ordinary supervised learning, by least squares (Chapter 4.1), with training pairs made for free by noising data:

$$\text{loss} = \big\| \varepsilon - \hat\varepsilon(x_t, t) \big\|^2$$

Knowing the noise means knowing which way the clean data lies. At every point of the space and at every step, the trained network points towards where the data is likely to be:

mesh = np.stack(
    np.meshgrid(np.linspace(-2, 1.8, 17), np.linspace(-1.4, 1.9, 15)), axis=-1
).reshape(-1, 2)
truth = np.r_[spiral(1000, rng), ring(1000, rng)]

fig, axes = style.subplots(3, height=3.4, width=10.5)
for ax, t in zip(axes, (9, 24, 39)):
    guess = predict_noise(mesh, t, label=2)
    clean = (mesh - np.sqrt(1 - alpha_bar[t]) * guess) / np.sqrt(alpha_bar[t])  # the implied x0
    move = clean * np.sqrt(alpha_bar[t]) - mesh
    scatter(ax, truth, f"step {t + 1}", style.GREY)
    ax.quiver(mesh[:, 0], mesh[:, 1], move[:, 0], move[:, 1], color=style.CORAL, width=0.004)
fig.tight_layout()
plt.show()

Figure 8.4.4: What the trained network has learned, for points noised for 10, 25 and 40 steps (no label given). Each arrow shows the direction and size of the network’s guess at the clean point from that position, scaled for display; grey dots are the data. Early in the reverse process (right, much noise) the arrows point broadly towards the middle; late (left, little noise) they pull sharply onto the spiral and the ring.

8.4.5 Generating: from noise to data

To generate, start from pure normal noise and run the steps backwards. At each step, subtract the predicted noise (scaled to that step), and add back a little fresh noise so that the samples keep their variety:

$$x_{t-1} = \frac{1}{\sqrt{1 - \beta_t}} \Big( x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\, \hat\varepsilon(x_t, t) \Big) + \sqrt{\beta_t}\, z, \qquad z \sim N(0, I)$$

snapshots = generate(2000, keep=(20, 10, 4))

fig, axes = style.subplots(5, height=2.6, width=11)
for ax, t in zip(axes, (STEPS, 20, 10, 4, 0)):
    scatter(
        ax,
        snapshots[t],
        "pure noise" if t == STEPS else f"{STEPS - t} steps back" if t else "generated",
    )
    if t:
        ax.set(xlim=(-3.2, 3.2), ylim=(-3.2, 3.2))
fig.tight_layout()
plt.show()

Figure 8.4.5: Generating 2,000 new points with the trained network, without a label. From pure noise (left), each step removes a little predicted noise; the cloud contracts, and the spiral and the ring emerge. The final points (right) are new: none is a copy of a training point.

An image generator does exactly this in hundreds of thousands of dimensions instead of two, with a convolutional or transformer network (Chapters 8.2 and 8.3) as the noise predictor, and 20 to 1,000 steps.

8.4.6 Conditioning: asking for something

A useful generator makes what you ask for. The noise predictor is given the request as an extra input: a class label here, the embedding of a text prompt in a text-to-image model. During training the label is sometimes left out, so the same network also learns the unlabelled version. Guidance then pushes each step further away from the unlabelled guess and towards the labelled one:

$$\hat\varepsilon_{\text{guided}} = \hat\varepsilon(x_t) + (1 + w)\,\big(\hat\varepsilon(x_t \mid \text{label}) - \hat\varepsilon(x_t)\big)$$

REQUESTS = [
    ("no label", 2, 0.0),
    ("label: spiral", 0, 0.0),
    ("label: ring", 1, 0.0),
    ("ring, guidance 3", 1, 3.0),
]

fig, axes = style.subplots(4, height=2.9, width=10.5)
for ax, (title, label, guidance) in zip(axes, REQUESTS):
    scatter(ax, generate(1500, label, guidance)[0], title, style.PURPLE)
fig.tight_layout()
plt.show()

Figure 8.4.6: One network, four requests. Without a label it makes both shapes; with the label ‘spiral’ or ‘ring’ it makes only that one. Guidance with w = 3 sharpens the ring further: more typical, less varied. Text-to-image models use the same trick, with a text embedding as the label.

8.4.7 The perfect denoiser memorises

There is a statistical catch. For a finite training set, the best possible noise predictor is known exactly: the clean image is a weighted average of the training images, weighted by how likely each one is to have produced the noisy one. A denoiser that achieved this perfect fit would only ever generate copies of its training images. Generating digits with that perfect denoiser shows it:

train_images = pixels[:1000] * 2 - 1


def perfect_noise_guess(x: np.ndarray, t: int) -> np.ndarray:
    """The noise implied by the exact posterior mean of the clean image, given the training set."""
    scale, spread = np.sqrt(alpha_bar[t]), 1 - alpha_bar[t]
    distance = ((x[:, None, :] - scale * train_images[None]) ** 2).sum(axis=2)
    log_weight = -distance / (2 * spread)
    weight = np.exp(log_weight - log_weight.max(axis=1, keepdims=True))
    clean = (weight / weight.sum(axis=1, keepdims=True)) @ train_images
    return (x - scale * clean) / np.sqrt(spread)


x = np.random.default_rng(3).normal(size=(8, 64))
for t in reversed(range(STEPS)):
    x = (x - beta[t] / np.sqrt(1 - alpha_bar[t]) * perfect_noise_guess(x, t)) / np.sqrt(
        alpha[t]
    )
    if t > 0:
        x += np.sqrt(beta[t]) * rng.normal(size=x.shape)
nearest = train_images[((x[:, None, :] - train_images[None]) ** 2).sum(axis=2).argmin(axis=1)]

fig, axes = plt.subplots(2, 8, figsize=(10, 2.8))
for col in range(8):
    for row, (image, title) in enumerate(
        ((x[col], "generated"), (nearest[col], "closest training"))
    ):
        axes[row, col].imshow(image.reshape(8, 8), cmap="gray_r", vmin=-1, vmax=1)
        axes[row, col].set_axis_off()
        if col == 0:
            axes[row, col].set_title(title, fontsize=9, loc="left")
fig.tight_layout()
plt.show()

Figure 8.4.7: Top: eight digits generated by the exact optimal denoiser for 1,000 training digits. Bottom: the closest training digit to each. They are copies. Real diffusion models are trained networks that cannot fit their training set perfectly, and that imperfection is what lets them make new images; but with small or repetitive training sets, they too can reproduce training images.

This is overfitting in a new guise, with real consequences: generated images that reproduce training images can leak private photos or copyrighted work. Checking generated samples against their nearest training examples, as in the figure, is a basic test.

8.4.8 Other generative models

Model	How it generates	Strengths	Weaknesses
Autoregressive (Chapter 8.3)	one token after another, each from $P(\text{next} \mid \text{past})$	exact likelihood; text, code	slow for long outputs
Variational autoencoder	decode a random normal code	fast; smooth latent space	blurry images
Generative adversarial network	a generator against a discriminator that tells real from fake	sharp, fast images	unstable training; misses variety
Normalising flow	transform normal noise by invertible steps	exact likelihood	restricted networks
Diffusion	remove predicted noise step by step	best image quality; stable training	many steps per sample

Judging a generative model is itself a statistical problem: compare the distribution of generated samples with that of real held-out data, not single examples. Common scores compare the mean and covariance of network features of real and generated images, a two-sample comparison in the spirit of Part 3.

8.4.9 In one picture

• A generative model learns $P(x)$ and samples from it; simple distributions fail on complex shapes.
• An autoencoder squeezes data into a small latent code and rebuilds it; a variational autoencoder makes the codes normal so that random codes decode to new data.
• Diffusion: the forward process $x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t}\,\varepsilon$ destroys data; a network learns to predict ε by least squares; running the steps backwards from noise generates new data.
• Conditioning and guidance steer generation with a label or a text prompt.
• A perfect fit to the training set would only memorise; check generated samples against their nearest training examples.

Next: reinforcement learning, learning from rewards instead of answers.