Latent diffusion and diffusion transformers: autoencoders, the scale factor, cross-attention and DiT

Before this lesson: lesson 02 (VAEs and VQ-VAE: the first stage), lesson 08 (SNR, log-SNR $\lambda$, zero terminal SNR) and lesson 11 (classifier-free guidance and the null condition). From deep learning: attention (queries, keys, values), the transformer and ViT (patch tokens) and normalization and residual blocks (layer norm). Time: about 70 minutes, plus the notebook.

---

Why this matters #

Latent diffusion is the design of Stable Diffusion 1–3, of the diffusion transformer (DiT) and its descendants, and of most current image and video generators. It is also used in science: LDCast runs latent diffusion for precipitation nowcasting, and latent models make large weather and climate ensembles affordable.

Two choices decide quality and cost: the autoencoder (factor $f$, channels, regularization, scale) and the backbone with its conditioning (a U-Net with cross-attention, or a transformer with adaLN-Zero or joint attention). In interviews you are expected to estimate token counts and attention costs quickly, by hand. This lesson gives you the formulas and the habits.

1. Predict first #

2. The picture #

Two kinds of compression #

Rombach et al. (2022) split the bits of an image into two kinds. Perceptual detail is high-frequency texture: a likelihood model must pay for it, but it matters little to the eye. Semantic content is the layout and the objects. An autoencoder can remove most perceptual bits cheaply. A diffusion model is good at semantic generation. Latent diffusion gives each part to the model that is good at it.

Two stages #

1. Train an autoencoder once. The encoder $\mathcal E$ maps $x \in \mathbb R^{H \times W \times C}$ to $z = \mathcal E(x) \in \mathbb R^{h \times w \times c}$ with $h = H/f$, $w = W/f$. The integer $f$ is the downsampling factor. The decoder $\mathcal D$ maps back: $\tilde x = \mathcal D(z) \approx x$.
2. Freeze the autoencoder. Train a diffusion model on the latents $z$. To sample, run the whole reverse process on $z$ and call the decoder once at the end.

The latent is still an image. The encoder is convolutional, so $z$ keeps a spatial layout: position $(i, j)$ of $z$ describes an $f \times f$ patch of $x$. A U-Net or a patch transformer can use the same 2-D structure as in pixel space.

Tokens. A transformer cuts the latent into $p \times p$ patches, one token per patch (as in ViT). The number of tokens is

and self-attention compares every token with every other token: $T^2$ pairs per head and layer.

Interactive figure: loading…

Three entry points for a condition #

A condition can enter the denoiser in three ways. The shape of the condition chooses the way:

• a map on the same grid (a low-resolution image, a mask, a topography map): concatenate it to the input channels;
• a single vector (the timestep, a class, a pooled text vector, scenario parameters): modulate the normalization layers (adaLN);
• a sequence of tokens (text, a variable set of observations): attention, either cross-attention (the image reads the text) or joint attention (both read each other).

3. Definitions and main results #

Stage 1: the autoencoder #

• Losses (Rombach et al., section 3.1 and appendix G): reconstruction ($\ell_1$) + perceptual (LPIPS) + a patch-based adversarial loss (lesson 05)
  • a small regularizer. There are two regularizer types: a KL term toward $\mathcal N(0, I)$ with a weight of about $10^{-6}$ (KL-reg.), or a vector-quantization layer absorbed into the decoder (VQ-reg., lesson 02).
• Why the KL weight is tiny. The goal is high-fidelity reconstruction, so the code should be almost deterministic. The KL term only stops the latent from drifting to an arbitrary scale. Modeling the distribution of $z$ is the job of the diffusion model, not of the prior.
• Compression factor in number of values:

Stable Diffusion's KL-f8 autoencoder has $f = 8$, $C = 3$, $c = 4$, so the factor is $48$. Stable Diffusion 3 uses $c = 16$, so $12$.

• The $f$ trade-off (Rombach et al., section 4.1, with $f \in \{1, 2, 4, 8, 16, 32\}$ and the same compute for all). Small $f$ (1–2): training is slow, because the diffusion model must still do the perceptual compression. Large $f$ (32): fidelity stops improving, because the autoencoder has thrown away information. $f = 4$ to $16$ balances the two. After 2M steps on ImageNet, the FID of LDM-8 was 38 points better than that of pixel-space LDM-1.
• The decoder is a ceiling. No diffusion model can produce detail that $\mathcal D$ cannot reconstruct. More latent channels improve reconstruction, but they make the latent harder to model. Esser et al. (2024) found that $c = 16$ pays off with large diffusion models, and chose it for Stable Diffusion 3.

The latent scale factor #

The noise schedule was designed for unit-variance data (lesson 08). Encoder outputs need not have unit variance.

Proposition. Let the latents have standard deviation $\sigma_z$ per coordinate, and use a VP process $z_t = \sqrt{\bar\alpha_t}\,z + \sqrt{1 - \bar\alpha_t}\,\varepsilon$. Then the effective signal-to-noise ratio is

Proof.

1. The signal part $\sqrt{\bar\alpha_t}\,z$ has variance $\bar\alpha_t\sigma_z^2$ per coordinate.
2. The noise part $\sqrt{1 - \bar\alpha_t}\,\varepsilon$ has variance $1 - \bar\alpha_t$.
3. Their ratio is $\sigma_z^2\,\mathrm{SNR}(t)$. Take the log: $\log(\sigma_z^2\,\mathrm{SNR}) = \log\mathrm{SNR} + 2\ln\sigma_z$. $\blacksquare$

So Rombach et al. (appendix G) rescale the KL-reg. latents: estimate one standard deviation $\hat\sigma$ over the first batch, and train on $z/\hat\sigma$. Stable Diffusion v1 stores the scale factor $0.18215$ ($\hat\sigma \approx 5.49$); SDXL stores $0.13025$ ($\hat\sigma \approx 7.68$). Newer autoencoders also store a shift (a mean), which is subtracted before scaling. VQ-reg. latents already have a variance near 1.

What goes wrong without it. With $\sigma_z = 5.49$, every noise level has a $5.49^2 = 30.1$ times higher SNR than designed.

• At $t = T$ the Stable Diffusion schedule then leaves a signal with standard deviation $0.375$ in $z_T$ (instead of $0.068$). But sampling starts from pure noise. This is a large train–test mismatch at the first sampling step: the zero-terminal-SNR problem of lesson 08, made much worse.
• SNR = 1 moves from $t \approx 355$ to $t \approx 816$. Few training steps fall at the low SNR where the global layout is decided, so the model learns the layout poorly.

Interactive figure: loading…

Conditioning mechanisms #

Cross-attention (Rombach et al., section 3.3). Flatten an intermediate U-Net feature map to $\varphi(z_t) \in \mathbb R^{N \times d}$ ($N = h_i w_i$ latent tokens at that level). A text encoder gives $\tau(y) \in \mathbb R^{M \times d_\tau}$ (Stable Diffusion v1: a frozen CLIP ViT-L/14 text encoder, $M = 77$, $d_\tau = 768$). Then

and the output is added to the residual stream. Queries come from the latent tokens; keys and values come from the text tokens. Each latent position takes a weighted mix of text tokens.

• Cost per layer: $NMd$ multiply-adds for $QK^\top$ and $NMd$ for the product with $V$, plus the projections. Self-attention over the same latent costs $N^2 d$ for each of the two products. The attention matrix has $N \times M$ entries per head, against $N \times N$.
• Information flows one way: the text tokens are not updated.
• For classifier-free guidance, the null condition is the embedding of the empty prompt (lesson 11).
• In Stable Diffusion v1, each attention block is self-attention, then cross-attention, then an MLP, at the latent resolutions $64^2$, $32^2$ and $16^2$, and in the middle block at $8^2$ (for $512^2$ images).

Interactive figure: loading…

Concatenation. For a condition $c$ that lives on the same grid as $z$ (a low-resolution image for super-resolution, a mask and a masked image for inpainting, a semantic map), feed $[z_t, c]$ along the channel axis. It costs only extra input channels in the first layer, and it keeps exact spatial alignment. Rombach et al. use it for super-resolution and semantic synthesis.

adaLN and adaLN-Zero (Peebles and Xie 2023, section 3.2). A global vector $\mathbf c$ (the sum of the timestep and class embeddings) goes through a small MLP. Per block, it outputs a shift $\beta$, a scale $\gamma$ and a gate $\alpha$ for the attention branch and for the MLP branch:

Here LN is layer norm without learned parameters, and $\odot$ multiplies element by element.

• "Zero": the last linear layer of that MLP starts at zero. Then $\alpha_1 = \alpha_2 = 0$, so $h = x$ and $\mathrm{out} = h = x$: every block is the identity at initialization. (The final output layer is also zero-initialized.) This follows the ResNet practice of starting residual blocks as the identity.
• The gates learn first. At initialization the output does not depend on the attention and MLP weights (they are multiplied by $\alpha = 0$), so their gradients are zero. The gradient with respect to $\alpha$ is $\mathrm{Attn}(\cdot)$ or $\mathrm{MLP}(\cdot)$, which is not zero. So the first updates open the gates.
• The same modulation acts on every token. It suits a global vector; a token sequence needs attention.
• DiT's ablation (DiT-XL/2, 400k steps, FID without guidance): adaLN-Zero 19.5, adaLN 25.2, cross-attention 26.1 (with about 15% more compute), in-context tokens 35.2.

The diffusion transformer (DiT) #

1. Input. The noised latent, $I \times I \times c$. For $256 \times 256$ images this is $32 \times 32 \times 4$ (the frozen Stable Diffusion autoencoder).
2. Patchify. Cut it into $p \times p$ patches and embed each one linearly (a convolution with kernel and stride $p$): $T = (I/p)^2$ tokens of width $d$, plus fixed 2-D sine–cosine position embeddings. DiT uses $p \in \{2, 4, 8\}$. Halving $p$ multiplies $T$ by 4 and the compute by at least 4, with almost no change in the number of parameters.
3. Blocks. $L$ adaLN-Zero blocks. Configurations S, B, L, XL (XL: $L = 28$, $d = 1152$, 16 heads).
4. Output. A final adaLN and a linear layer map each token to $p \times p \times 2c$ values (the noise prediction and a diagonal covariance), rearranged back to $I \times I \times 2c$.
5. Scaling. FID falls steadily as the forward-pass compute grows, through model size or smaller patches. DiT-XL/2 (118.6 Gflops) reached FID 2.27 on ImageNet $256 \times 256$ with classifier-free guidance (scale 1.5), the best result at the time.

Counting a block. Attention has $4d^2$ weights ($Q$, $K$, $V$, output), the MLP $8d^2$ (width $4d$), and the adaLN-Zero modulation $6d^2$: about $18d^2$ per block. For DiT-XL: $28 \cdot 18 \cdot 1152^2 \approx 669$M, plus embeddings: the 675M of the paper. The modulation is a third of the parameters but almost no compute, because it runs once per sample, not once per token. The compute per block and token is $12d^2$ multiply-adds for the linear layers and $2Td$ for attention (scores and values).

MM-DiT: joint attention (Esser et al. 2024, Stable Diffusion 3) #

• Two token streams. Text tokens (from two CLIP encoders and T5-XXL) and image tokens ($2 \times 2$ patches of the 16-channel latent).
• Separate weights, joint attention. Each stream has its own projections, MLPs, norms and modulation. For attention, the two sequences are concatenated: $Q = [Q_{\text{txt}}; Q_{\text{img}}]$, and the same for $K$ and $V$. This equals two transformers that share one attention operation. Text tokens now read the image too (two-way), unlike in cross-attention.
• Global conditioning. The timestep and pooled CLIP text vectors enter through adaLN-style modulation. The pooled vector holds only coarse information, so the token sequence is still needed.
• Cost. Attention over $N + M$ tokens: $(N + M)^2 d$ multiply-adds per product.
• Training. Rectified flow (lesson 10). In their comparison on CC12M, MM-DiT beat CrossDiT (cross-attention to text), UViT and a DiT with concatenated text tokens and shared weights.

Training details of large latent transformers

• Timestep sampling. SD3 samples $t$ from a logit-normal distribution, which puts more training steps at intermediate noise levels (lesson 08: the sampling density of $t$ is a loss weight).
• Resolution-dependent shift. More tokens average the noise, so the same $t$ destroys less structure at high resolution. SD3 shifts the timesteps toward more noise for larger images (shift 3.0 at $1024^2$), the latent analog of the shifted schedules of lesson 08.
• QK-normalization. RMSNorm on $Q$ and $K$ before the dot product stops the attention logits from growing during mixed-precision training at high resolution.
• Store the latents. The encoder is frozen, so encode the data set once and store the latents. Apply augmentations before encoding, or store several augmented latents per example.

Answers to section 1. (1) $196{,}608 \to 32 \cdot 32 \cdot 4 = 4{,}096$ values: 48 times fewer ($f^2C/c = 64 \cdot 3/4$). Positions: $65{,}536 \to 1{,}024$ ($f^2 = 64$ times fewer), so the $N^2$ term falls $64^2 = 4{,}096$ times. (2) The global layout. The effective SNR is $5.5^2 \approx 30$ times higher at every step (the Proposition), so the model rarely trains at the noise levels where the layout is decided, and its first sampling step sees an input with much less signal than in training. (3) The image (latent) tokens give the queries; the text tokens give keys and values.

4. Worked example #

(a) DiT on ImageNet 256.

1. Image: $256 \cdot 256 \cdot 3 = 196{,}608$ values. Latent ($f = 8$, $c = 4$): $32 \cdot 32 \cdot 4 = 4{,}096$ values, 48 times fewer.
2. Tokens: $p = 2$ gives $(32/2)^2 = 16^2 = 256$; $p = 4$ gives $8^2 = 64$; $p = 8$ gives $4^2 = 16$.
3. Attention-matrix entries per head and layer: $256^2 = 65{,}536$, $64^2 = 4{,}096$ and $16^2 = 256$.

(b) The compute of DiT-XL/2 ($d = 1152$, $T = 256$, 28 blocks), in multiply-adds.

1. Linear layers: $12d^2 = 12 \cdot 1{,}327{,}104 = 15{,}925{,}248$ per token and block.
2. Attention: $2Td = 2 \cdot 256 \cdot 1152 = 589{,}824$ per token and block (only 4% of the linear part at 256 tokens).
3. All blocks: $28 \cdot 256 \cdot (15{,}925{,}248 + 589{,}824) = 118.38$G.
4. adaLN modulation, once per sample: $28 \cdot 6d^2 = 0.22$G. Patch embedding and output layer: $0.02$G.
5. Total: $118.62$G. DiT reports $118.64$ "Gflops", so the paper counts one multiply-add as one operation.

(c) Stable Diffusion v1 at $512 \times 512$. Latent $64 \times 64 \times 4$. At the highest-resolution attention level, $N = 4096$ tokens of width $d = 320$; text $M = 77$.

1. Self-attention: $N^2 d = 4096^2 \cdot 320 = 5.37$G multiply-adds for $QK^\top$ (10.7G with the value product).
2. Cross-attention: $NMd = 4096 \cdot 77 \cdot 320 = 0.10$G for $QK^\top$ (0.20G with values): $N/M = 53$ times cheaper.
3. Memory: one $4096 \times 4096$ attention matrix in fp16 is $4096^2 \cdot 2$ bytes $= 33.5$ MB; 8 heads give 268 MB per layer and image, unless a memory-efficient (flash) kernel avoids storing it.
4. Pixel tokens instead: $512^2 = 262{,}144$ tokens, $6.9 \times 10^{10}$ attention entries per head.

(d) The scale factor.

1. $\hat\sigma = 1/0.18215 = 5.49$.
2. Log-SNR shift: $2\ln 5.49 = 3.41$; SNR multiplied by $5.49^2 = 30.1$.
3. For the Stable Diffusion schedule, $\mathrm{SNR}(T) = 0.0047$ becomes $0.141$ without scaling.
4. SNR = 1 is reached at $t \approx 355$ for scaled latents, but only at $t \approx 816$ (where $\bar\alpha_t \approx 0.032$) without scaling.

5. In code #

NumPy checks the worked example: the parameter and compute count of DiT-XL/2, an adaLN-Zero block that is the identity at initialization, cross-attention shapes and the scale factor. The notebook adds a small latent diffusion model on scikit-learn digits.

import numpy as np

# Counting DiT-XL/2: parameters of one adaLN-Zero block and multiply-adds of the whole model.
d, T, L = 1152, 256, 28
attn = 4 * d * d + 4 * d                 # Q, K, V and output projections with biases
mlp = 8 * d * d + 5 * d                  # d -> 4d -> d with biases
ada = 6 * d * d + 6 * d                  # c -> (shift, scale, gate) for attention and MLP
print(attn + mlp + ada, L * (attn + mlp + ada))           # 23905152 669344256
macs = L * T * (12 * d * d + 2 * T * d)  # per token: linear layers 12 d^2, attention 2 T d
macs += L * 6 * d * d                    # adaLN modulation: once per sample, not per token
macs += T * (16 * d + 32 * d) + 2 * d * d   # patch embedding (2x2x4 -> d), final layer (d -> 2x2x8)
print(f"{macs / 1e9:.2f}")               # 118.62  (DiT paper: 118.64 "Gflops")

# An adaLN-Zero block (one head for brevity) is the identity at initialization.
rng = np.random.default_rng(0)
ln = lambda x: (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + 1e-6)
silu = lambda x: x / (1 + np.exp(-x))
def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)
def block(x, c, P):
    b1, g1, a1, b2, g2, a2 = np.split(silu(c) @ P["Wada"] + P["bada"], 6)
    h = ln(x) * (1 + g1) + b1
    q, k, v = h @ P["Wq"], h @ P["Wk"], h @ P["Wv"]
    x = x + a1 * (softmax(q @ k.T / np.sqrt(x.shape[1])) @ v) @ P["Wo"]
    h = ln(x) * (1 + g2) + b2
    return x + a2 * (np.maximum(h @ P["W1"], 0) @ P["W2"])
k = 16
shapes = dict(Wq=(k, k), Wk=(k, k), Wv=(k, k), Wo=(k, k), W1=(k, 4 * k), W2=(4 * k, k))
P = {name: rng.normal(size=s) / np.sqrt(s[0]) for name, s in shapes.items()}
P["Wada"], P["bada"] = np.zeros((k, 6 * k)), np.zeros(6 * k)       # the "Zero"
x, c = rng.normal(size=(10, k)), rng.normal(size=k)
print(np.array_equal(block(x, c, P), x))                           # True

Cross-attention at the $64 \times 64$ level of Stable Diffusion v1, and the scale factor:

import numpy as np

rng = np.random.default_rng(0)
def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)
N, d, M, d_txt = 64 * 64, 320, 77, 768
img = rng.normal(size=(N, d)).astype(np.float32)          # latent tokens
txt = rng.normal(size=(M, d_txt)).astype(np.float32)      # CLIP text tokens
Wq = rng.normal(size=(d, d)).astype(np.float32) / np.sqrt(d)
Wk = rng.normal(size=(d_txt, d)).astype(np.float32) / np.sqrt(d_txt)
Wv = rng.normal(size=(d_txt, d)).astype(np.float32) / np.sqrt(d_txt)
Q, K, V = img @ Wq, txt @ Wk, txt @ Wv                    # queries from the image, keys/values from the text
A = softmax(Q @ K.T / np.sqrt(d))
out = A @ V
print(Q.shape, K.shape, A.shape, out.shape)          # (4096, 320) (77, 320) (4096, 77) (4096, 320)
print(np.allclose(A.sum(1), 1, atol=1e-5))            # True
print(f"{N * M * d / 1e9:.2f} {N * N * d / 1e9:.2f}")  # 0.10 5.37: G multiply-adds of Q K^T, cross vs self

betas = np.linspace(0.00085**0.5, 0.012**0.5, 1000) ** 2   # Stable Diffusion "scaled linear"
ab = np.cumprod(1 - betas)
snr = ab / (1 - ab)
sigma_z = 1 / 0.18215
print(f"{sigma_z:.2f} {2 * np.log(sigma_z):.2f}")      # 5.49 3.41: log-SNR shift without rescaling
print(f"{snr[-1]:.4f} {sigma_z**2 * snr[-1]:.3f}")     # 0.0047 0.141: SNR at t = T, scaled vs unscaled
print(f"{np.sqrt(ab[-1]) * sigma_z:.3f}")              # 0.375: std of the signal left in z_T (unscaled)
print(np.argmin(np.abs(snr - 1)) + 1, np.argmin(np.abs(sigma_z**2 * snr - 1)) + 1)   # 355 816

6. Where it shows up #

• Text-to-image. Stable Diffusion 1.x and 2.x: KL-f8 autoencoder ($c = 4$), U-Net with cross-attention to a frozen text encoder. SDXL: a larger U-Net, two text encoders, scale factor $0.13025$. Stable Diffusion 3: MM-DiT, $c = 16$, rectified flow. Many later image and video models use DiT-style transformers with adaLN modulation.
• Few-step models (lesson 12). Latent consistency models, SDXL-Turbo and InstaFlow all distill latent diffusion models. With 1–4 steps, the decoder can become a noticeable share of the total time.
• Weather and climate. LDCast (Leinonen et al. 2023) runs latent diffusion for precipitation nowcasting. Latent models cut the cost of large ensembles in the same way as for images.
• Scientific fields with conditions. The three entry points map naturally. Maps on the same grid (topography, land–sea mask, a coarse simulation to downscale) go in by concatenation. Scalars (scenario parameters, physical constants) go through an embedding added to the timestep embedding and used by adaLN. A variable set of station observations becomes tokens for cross-attention or joint attention.
• Training practice. Report the autoencoder's own reconstruction metrics on the quantities you care about (for scientific fields: extremes, spectra, peak values). They bound the final quality.

7. Common mistakes #

• Forgetting the scale (and the shift), applying it twice, or not undoing it before decoding. Every noise level is then wrong by a factor $\sigma_z^2$ in SNR.
• Judging the diffusion model before the autoencoder. Check reconstructions first. No sampler can beat the decoder.
• Too large an $f$ for small structures. Text and faces in images; thin fronts, peaks and high frequencies in scientific fields.
• The wrong entry point for a condition. A map fed through cross-attention loses exact alignment and costs $NM$; a token sequence pooled into adaLN loses its detail (the reason SD3 keeps the sequence).
• adaLN without the zero initialization. DiT found it clearly worse (FID 25.2 against 19.5).
• Attention-cost slips. There is one $T \times T$ attention matrix per head and per layer; halving $p$ gives 4 times the tokens and 16 times the entries of that matrix; DiT's "Gflops" are multiply-adds.
• Retraining the encoder after the diffusion model. The latent distribution changes, so the diffusion model no longer matches. Stable Diffusion's improved autoencoders fine-tuned only the decoder for this reason.

8. Check yourself #

1. A weather data set on a $256 \times 256$ grid has 5 variables. The autoencoder has $f = 8$ and $c = 16$. Give the latent shape, the compression factor and the number of DiT tokens with $p = 2$.
   Answer

   $32 \times 32 \times 16 = 16{,}384$ values against $256 \cdot 256 \cdot 5 = 327{,}680$: $f^2C/c = 64 \cdot 5/16 = 20$ times fewer. With $p = 2$: $16^2 = 256$ tokens, each from a $2 \times 2 \times 16 = 64$-value patch.
2. Unscaled latents have standard deviation 2. What is the log-SNR shift, and what is the effective SNR at a step where unit-variance data would have SNR 1?
   Answer

   Shift $2\ln 2 = 1.386$; effective SNR $= 2^2 \cdot 1 = 4$.
3. Why does adaLN-Zero start every block as the identity, and why does that help?
   Answer

   The modulation MLP's last layer is zero, so the gates $\alpha_1 = \alpha_2 = 0$ and both residual branches add nothing. A deep residual network that starts as the identity has well-behaved signals and gradients; each block then learns a correction, first through the gates. DiT found it clearly better than plain adaLN.
4. Compare cross-attention and MM-DiT joint attention: who attends to whom, which weights are shared, and the cost.
   Answer

   Cross-attention: queries from the $N$ image tokens, keys and values from the $M$ text tokens; the text is not updated; cost $\propto NMd$. Joint attention: one attention over the concatenated $N + M$ tokens, so text and image attend to each other; each modality has its own projections, MLPs and norms; cost $\propto (N + M)^2 d$.
5. Compare the $N^2 d$ term of self-attention at a $32 \times 32$ level with $d = 640$ and at a $64 \times 64$ level with $d = 320$.
   Answer

   $1024^2 \cdot 640 = 0.67$G against $4096^2 \cdot 320 = 5.37$G multiply-adds: the high-resolution level costs 8 times more, although its width is half.
6. A model of climate fields is conditioned on a land–sea mask (same grid), a $\mathrm{CO_2}$ scenario value, and a variable number of station records. How should each condition enter?
   Answer

   The mask is a map on the same grid: concatenate it to the input channels (exact alignment, cheap). The scenario value is a scalar: embed it, add it to the timestep embedding, and use adaLN. The station records are a variable-length set: turn them into tokens for cross-attention or joint attention.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What are the two stages of latent diffusion?

   Show answer

2. What is the compression factor in values of an autoencoder with factor $f$, $C$ image channels and $c$ latent channels?

   Show answer

3. How many DiT tokens does a $512 \times 512$ image give with $f = 8$ and patch size $p = 2$?

   Show answer

4. Latents have standard deviation $\sigma_z$. How does this change the SNR of a VP schedule?

   Show answer

5. Why does Stable Diffusion multiply its latents by $0.18215$?

   Show answer

6. In cross-attention, where do $Q$, $K$ and $V$ come from?

   Show answer

7. What does adaLN-Zero do at initialization, and why?

   Show answer

8. What does "patchify" do in a DiT?

   Show answer

9. How does MM-DiT's joint attention differ from cross-attention?

   Show answer

10. By what factor does the self-attention term grow when the image side doubles at fixed $f$ and $p$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why latent diffusion?" — Most pixel bits are perceptual detail; an autoencoder removes them once, diffusion runs on a latent with 48 times fewer values (64 times fewer positions for $f = 8$), and the decoder is called once. The decoder bounds the quality.
• "What is 0.18215?" — $1/\hat\sigma$ of the KL-f8 latents. It makes them unit-variance so that the noise schedule's SNR is as designed.
• "How does text enter Stable Diffusion?" — Cross-attention: queries from the U-Net's latent tokens, keys and values from frozen CLIP text tokens. The empty prompt is the CFG null condition.
• "What is adaLN-Zero?" — Shift, scale and gate regressed from the timestep and class embedding; the gate starts at zero, so each block starts as the identity.
• "Estimate the attention cost of a DiT on $512 \times 512$ images with $p = 2$." — Latent $64 \times 64$, so $32^2 = 1024$ tokens; about $1024^2 \cdot d$ multiply-adds for each of the two attention products, per layer (summed over the heads, total width $d$).
• "What does MM-DiT change?" — Joint attention over text and image tokens with separate weights per modality, instead of one-way cross-attention.

10. Practice #

The notebook for this lesson has five exercises, each with a check cell and a folded solution:

1. A tiny latent diffusion model: a PCA "autoencoder" on scikit-learn digits, a Gaussian mixture fitted to the latents, and exact-denoiser DDIM sampling in latent space; decode and look at the digits.
2. Measure the latent standard deviation and compute the effective SNR and the step where SNR = 1, with and without the scale factor.
3. Write a token and FLOP calculator, and check it against DiT-XL/2 and Stable Diffusion v1.
4. Write single-head cross-attention in NumPy; check the shapes, that rows sum to 1, and a padding mask.
5. Write an adaLN-Zero block and check that it is the identity at initialization and that only the gates get a gradient at first.

11. Watch and read #

How does Stable Diffusion work? – Latent Diffusion Models EXPLAINEDAI Coffee Break with Letitia · 13 min

A clear tour of the two stages, the autoencoder and cross-attention conditioning in Stable Diffusion.

Diffusion Transformers (ViT, DiT, MMDiT)Julia Turc · 18 min

From ViT patches to DiT's adaLN-Zero blocks and the joint attention of MM-DiT.

How AI connects text and images3Blue1Brown · 2 min

A short visual idea of the shared text–image embedding space behind CLIP text encoders.

• Rombach, Blattmann, Lorenz, Esser and Ommer, High-Resolution Image Synthesis with Latent Diffusion Models (CVPR 2022): section 3.1 (the autoencoder), sections 3.2–3.3 (latent diffusion and cross-attention), section 4.1 (the $f$ trade-off), appendix G (the KL weight and the rescaling).
• Peebles and Xie, Scalable Diffusion Models with Transformers (ICCV 2023): section 3.2 (patchify, the four block designs, adaLN-Zero), Table 1 (configurations), Figure 5 (block ablation), Table 4 (parameters and Gflops).
• Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (Stable Diffusion 3, 2024): section 4 (MM-DiT), section 5 (latent channels, comparison with CrossDiT and UViT, QK-normalization, resolution-dependent shift).
• The DiT reference code: facebookresearch/DiT (models.py: DiTBlock, FinalLayer, modulate, and the zero initialization in initialize_weights).
• Leinonen et al., Latent diffusion models for generative precipitation nowcasting with accurate uncertainty quantification (LDCast, 2023).