Stable Diffusion Models

diffusion models generate by learning to undo noise: destroy an image with a fixed gaussian corruption process, train a network to reverse one small step of the destruction, then chain the reversals from pure noise back to data (ho et al. 2020, denoising diffusion probabilistic models). 𐃏 stable diffusion (rombach et al. 2022, high-resolution image synthesis with latent diffusion models) runs this machinery not on pixels but in the latent space of an autoencoder, with a text-conditioned u-net doing the denoising. first the maths, then the architecture.

the forward (noising) process

fix a variance schedule \(\beta_1, \dots, \beta_T \in (0,1)\) and define a markov chain that gradually gaussianises the data:

\begin{equation} q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t;\; \sqrt{1 - \beta_t}\, x_{t-1},\; \beta_t I\right). \end{equation}

the \(\sqrt{1-\beta_t}\) shrinkage is chosen so variance is preserved: if \(\operatorname{Var}(x_{t-1}) = I\) then \(\operatorname{Var}(x_t) = (1-\beta_t) I + \beta_t I = I\). because compositions of these gaussians are gaussian, the chain telescopes into a closed form — with \(\alpha_t = 1 - \beta_t\) and \(\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s\),

\begin{equation} q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\; \sqrt{\bar{\alpha}_t}\, x_0,\; (1 - \bar{\alpha}_t) I\right) \quad\Longleftrightarrow\quad x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1 - \bar{\alpha}_t}\, \varepsilon,\;\; \varepsilon \sim \mathcal{N}(0, I). \end{equation}

this is the engine of cheap training: any noise level \(t\) is reachable from \(x_0\) in one step, no simulation of the chain required. as \(t \to T\) (with \(\bar{\alpha}_T \approx 0\)), \(x_T\) is indistinguishable from \(\mathcal{N}(0, I)\) — the data has been fully forgotten.

the diffusion markov chain: the fixed forward kernel q destroys structure left to right; the learned reverse kernel p rebuilds it right to left.

from variational bound to epsilon-prediction

the model is a latent-variable model \(p_\theta(x_0) = \int p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1} \mid x_t)\, dx_{1:T}\) with gaussian reverse kernels, trained by the usual variational bound on negative log-likelihood (Goodfellow, Ian, 2016):

\begin{equation} \mathbb{E}\big[-\log p_\theta(x_0)\big] \le \mathbb{E}_q\!\left[ \underbrace{\mathrm{KL}\big(q(x_T \mid x_0) \,\|\, p(x_T)\big)}_{L_T \approx 0} + \sum_{t>1} \underbrace{\mathrm{KL}\big(q(x_{t-1} \mid x_t, x_0) \,\|\, p_\theta(x_{t-1} \mid x_t)\big)}_{L_{t-1}} - \underbrace{\log p_\theta(x_0 \mid x_1)}_{L_0} \right]. \end{equation}

the pivotal fact: the forward posterior \(q(x_{t-1} \mid x_t, x_0)\) is itself gaussian with known mean

\begin{equation} \tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\, \beta_t}{1 - \bar{\alpha}_t}\, x_0 + \frac{\sqrt{\alpha_t}\, (1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}\, x_t, \end{equation}

so each \(L_{t-1}\) is a kl between gaussians — a weighted squared distance between means. now reparameterise: substitute \(x_0 = \big(x_t - \sqrt{1-\bar\alpha_t}\,\varepsilon\big)/\sqrt{\bar\alpha_t}\) and let the network predict the noise \(\varepsilon_\theta(x_t, t)\) rather than the mean. the bound becomes a weighted sum of terms \(\lVert \varepsilon - \varepsilon_\theta(x_t, t) \rVert^2\). the ddpm result, stated honestly: ho et al. drop the weights and train on the uniform-weight surrogate

\begin{equation} L_{\text{simple}} = \mathbb{E}_{t \sim \mathcal{U}\{1..T\},\, x_0,\, \varepsilon} \left[ \big\lVert \varepsilon - \varepsilon_\theta\big(\sqrt{\bar\alpha_t}\, x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon,\; t\big) \big\rVert^2 \right], \end{equation}

which is not the elbo but a reweighting of it that downweights the easy low-\(t\) terms — and produces better samples than the exact bound. the entire training loop is: pick an image, pick a timestep, add closed-form noise, regress the noise. no adversary, no equilibrium, compare gans.

the reverse process and sampling

the learned reverse kernel is \(p_\theta(x_{t-1} \mid x_t) = \mathcal{N}\big(x_{t-1};\, \mu_\theta(x_t, t),\, \sigma_t^2 I\big)\) with the mean expressed through the predicted noise:

\begin{equation} \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\, \varepsilon_\theta(x_t, t) \right), \qquad \sigma_t^2 = \beta_t \;\text{(or the posterior variance; both work)}. \end{equation}

ancestral sampling is then a loop:

  • init \(x_T \sim \mathcal{N}(0, I)\).
  • loop for \(t = T, \dots, 1\): predict \(\varepsilon_\theta(x_t, t)\), form \(\mu_\theta\), set \(x_{t-1} = \mu_\theta + \sigma_t z\) with \(z \sim \mathcal{N}(0,I)\) (no noise at \(t=1\)).
  • return \(x_0\).

that is \(T\) network evaluations per sample — ddpm used \(T = 1000\), which is the method’s honest weakness. ddim (song et al. 2020, denoising diffusion implicit models) reinterprets the same trained \(\varepsilon_\theta\) inside a non-markovian family whose deterministic member permits taking large strides over the schedule — 20–50 steps for near-identical quality, and, being deterministic, an invertible map between noise and image. 𐃏

latent diffusion: the stable part

pixel-space diffusion at 512px spends most of its capacity on imperceptible high-frequency detail. latent diffusion moves the chain into a compressed space:

  • vae latent space. a perceptually-trained autoencoder (kl-regularised, with a patch-adversarial loss on the decoder) maps a \(512 \times 512 \times 3\) image to a \(64 \times 64 \times 4\) latent — spatial factor 8, ~48x fewer numbers. diffusion happens entirely on these latents; see autoencoders.
  • u-net denoiser. \(\varepsilon_\theta\) is a u-net — downsampling and upsampling convolutional stages with skip connections — with the timestep injected as an embedding into every residual block.
  • cross-attention conditioning. the prompt goes through a frozen clip text encoder into a sequence of embeddings; at several u-net resolutions, attention layers use image features as queries and text embeddings as keys and values, so every spatial location can consult every word.
  • decode. after the reverse loop finishes, the vae decoder maps the final latent back to pixels — one pass, at the very end.
stable diffusion at inference: the reverse chain runs entirely in the vae’s latent space, steered by cross-attention to the text embedding; the decoder renders pixels once at the end.

classifier-free guidance

conditioning alone produces washed-out adherence to the prompt. cfg (ho and salimans 2022, classifier-free diffusion guidance) trains one network for both conditional and unconditional prediction by randomly dropping the conditioning during training, then at sampling extrapolates past the conditional prediction:

\begin{equation} \tilde{\varepsilon}_\theta(z_t, c) = \varepsilon_\theta(z_t, \varnothing) + w \big( \varepsilon_\theta(z_t, c) - \varepsilon_\theta(z_t, \varnothing) \big), \end{equation}

with guidance scale \(w > 1\) (stable diffusion defaults to \(w \approx 7.5\)). the difference vector points from “any image” toward “images matching \(c\)”; scaling it trades diversity for prompt fidelity, and overdoing it gives the over-saturated, over-contrasty look every heavy user recognises. cost: two u-net evaluations per step.

toy diffusion in 1d

the full pipeline — closed-form forward noising, an mlp \(\varepsilon\)-predictor with hand-written backprop, ancestral sampling — on the most brutal test of mode coverage: a dataset of exactly two points, \(x_0 \in \{-1, +1\}\). a mode-collapsing model would pick one.

import numpy as np

rng = np.random.default_rng(0)

# --- schedule: T steps, linear betas, closed-form alpha-bars
T = 50
betas = np.linspace(1e-4, 0.25, T)
alphas = 1 - betas
abar = np.cumprod(alphas)                  # \bar\alpha_t

# --- data: two points, x0 = -1 or +1 with equal probability
sample_data = lambda n: rng.choice([-1.0, 1.0], size=(n, 1))

# --- epsilon-predictor: tiny mlp, input (x_t, t/T), output eps_hat
p = {"W1": rng.normal(0, 0.5, (2, 32)), "b1": np.zeros(32),
     "W2": rng.normal(0, 0.5, (32, 1)), "b2": np.zeros(1)}
m = {k: 0.0 for k in p}; v = {k: 0.0 for k in p}

def forward(inp):
    h = np.tanh(inp @ p["W1"] + p["b1"])
    return h, h @ p["W2"] + p["b2"]

def adam_step(g, t, lr=2e-3):
    for k in p:
        m[k] = 0.9*m[k] + 0.1*g[k]
        v[k] = 0.999*v[k] + 0.001*g[k]**2
        p[k] -= lr * (m[k]/(1-0.9**t)) / (np.sqrt(v[k]/(1-0.999**t)) + 1e-8)

# --- training: sample t, noise with the closed form, regress the noise
B = 256
for step in range(1, 8001):
    x0 = sample_data(B)
    t = rng.integers(0, T, B)
    eps = rng.normal(size=(B, 1))
    xt = np.sqrt(abar[t])[:, None]*x0 + np.sqrt(1 - abar[t])[:, None]*eps
    inp = np.hstack([xt, (t[:, None] + 1)/T])
    h, eps_hat = forward(inp)
    d = 2*(eps_hat - eps)/B                       # d(mse)/d(eps_hat)
    g = {"W2": h.T @ d, "b2": d.sum(0)}
    dh = (d @ p["W2"].T) * (1 - h**2)
    g["W1"], g["b1"] = inp.T @ dh, dh.sum(0)
    adam_step(g, step)
    if step in (1, 1000, 4000, 8000):
        print(f"step {step:>5}  mse {((eps_hat - eps)**2).mean():.4f}")

# --- ancestral sampling: run the learned reverse chain from pure noise
n = 4000
x = rng.normal(size=(n, 1))                       # x_T ~ N(0, 1)
for t in reversed(range(T)):
    _, eps_hat = forward(np.hstack([x, np.full((n, 1), (t + 1)/T)]))
    x = (x - betas[t]/np.sqrt(1 - abar[t]) * eps_hat) / np.sqrt(alphas[t])
    if t > 0:
        x += np.sqrt(betas[t]) * rng.normal(size=(n, 1))

near_pos = (np.abs(x - 1) < 0.3).mean()
near_neg = (np.abs(x + 1) < 0.3).mean()
print(f"\nsamples: mean {x.mean():+.3f}, std {x.std():.3f}")
print(f"fraction within 0.3 of +1: {near_pos:.3f}, of -1: {near_neg:.3f}")
print("mean of |x| (target 1.0):", f"{np.abs(x).mean():.3f}")

step     1 mse 2.8974
step  1000  mse 0.3682
step  4000  mse 0.2710
step  8000  mse 0.1701

samples: mean +0.051, std 1.005
fraction within 0.3 of +1: 0.473, of -1: 0.435
mean of |x| (target 1.0):  0.989

starting from a single gaussian blob, the reverse chain splits the mass onto both modes — 47% lands on \(+1\), 44% on \(-1\), mean of \(|x|\) is \(0.989\) against a target of \(1.0\). the residual mse never reaches zero and should not: at high \(t\) the noise is genuinely unrecoverable, so a fraction of \(\lVert\varepsilon\rVert^2\) is irreducible.

see also

  • gans — the adversarial predecessor, and what mode collapse looks like
  • autoencoders — the compression stage that makes latent diffusion cheap
  • transformers — the cross-attention mechanism doing the text conditioning
  • cnn from scratch — the convolutional anatomy of the u-net

References

Goodfellow, Ian (2016). Deep Learning, MIT Press.