Diffusion Models: Hierarchical VAE with Fixed Encoder

#Why This Lesson Sits in the Deep-Learning Track

The generative-AI track has a full lesson on diffusion that walks through the ELBO derivation, classifier-free guidance, and the sampler landscape (DDIM, DPM-Solver++, EDM, LCM). This lesson does something different. It treats diffusion as an architecture story — a particular way of stacking the deep-learning primitives you have already met. The math is in the genai lesson; the architecture is here.

If a colleague asks "is diffusion really just a deep VAE?" the answer they want lives in this lesson. The mechanics of why the U-Net has skip connections, where the time embedding goes, why self-attention shows up at the 16x16 and 8x8 resolutions, what changes when you swap a U-Net for a Diffusion Transformer — those are deep-learning questions. The probability derivations are statistics questions. We are doing the first one.

#The Unifying View: Diffusion is a VAE With T Encoders

In the previous lesson you built a VAE. Encoder maps x to a distribution q(z | x); reparameterize to get a sample z; decoder maps z to a reconstruction; train with the ELBO — reconstruction term plus KL regularizer. One learned encoder, one learned decoder, one stochastic bottleneck.

Now imagine stacking that VAE T times. The first encoder takes x_0 (a clean image) to x_1 (slightly noisy). The second takes x_1 to x_2 (a bit noisier). After T steps you have x_T, which is pure Gaussian noise. The decoder runs in reverse — x_T to x_ to x_ ... back to x_0.

A "hierarchical VAE" with T levels would normally have T learned encoders and T learned decoders. Diffusion makes two ruthless simplifications:

1. The T encoders are fixed, not learned. Each q(x_t | x_) is a hand-chosen Gaussian noise-addition kernel. No parameters. No training. You just write down a noise schedule and apply it.
2. The T decoders are one shared network. A single neural network — typically a U-Net or a Diffusion Transformer — is invoked T times with the timestep t as an extra input. The same parameters do all the work at every noise level.

That is the entire reframe. Diffusion is a hierarchical VAE where the encoder is free and the decoder is shared.

This reframe matters because every deep-learning instinct you already have transfers. The decoder is a CNN-ish (or transformer-ish) architecture with the usual building blocks. The training loop is standard supervised learning — predict-a-target with MSE. The hardware footprint is dominated by the same things that dominate CNN training (memory bandwidth on activations, compute on attention). You are not learning a new paradigm; you are learning where to plug the parts in.

#Forward Process: Fixed Gaussian Encoders

The forward "encoder" at level t takes x_ and produces x_t by adding a sliver of Gaussian noise. The variance is set by a hand-chosen schedule β_t:

q(x_t | x_{t-1}) = N(√(1 - β_t) · x_{t-1},  β_t · I)

The schedule β_t is a hyperparameter, not a learned weight. The original DDPM (Ho 2020) used a linear schedule from β_1 = 10^ to β_T = 0.02 over T = 1000 steps. Modern training uses cosine or sigmoid schedules that keep "signal alive" longer at high noise levels — important for high-resolution images where the linear schedule turns the image to mush too fast. The genai-track lesson lists the trade-offs.

q(x_t | x_0) = N(√ᾱ_t · x_0,  (1 - ᾱ_t) · I)

The forward process has no parameters, so it has no gradients. Pure noise injection.

#Reverse Process: One Shared Learned Decoder

The reverse process has to go the other way — given x_t, produce x_. We parameterize it as Gaussian too:

p_θ(x_{t-1} | x_t) = N(μ_θ(x_t, t),  σ_t² · I)

Two practical simplifications, both due to Ho 2020:

1. The variance is fixed, not learned. Most implementations set σ_t² = β_t (or the closed-form posterior variance β̃_t = (1-ᾱ_)/(1-ᾱ_t) · β_t). Either works. The network only predicts the mean. Nichol & Dhariwal 2021's improved-DDPM did learn the variance for slightly better log-likelihoods, but it costs an extra head and most production code skips it.
2. The mean is parameterized by predicted noise, not by a direct mean head. Instead of outputting μ_θ(x_t, t) directly, the network outputs a noise prediction ε_θ(x_t, t) and the mean is reconstructed by the closed-form formula:

μ_θ(x_t, t) = (1/√α_t) · (x_t - (β_t / √(1 - ᾱ_t)) · ε_θ(x_t, t))

This is just algebra — the same formula you would get by inverting the forward identity x_t = √ᾱ_t · x_0 + √(1 - ᾱ_t) · ε and substituting back. The full ELBO-to-noise-prediction derivation lives in the generative-AI track lesson on diffusion (section "From Variational Bound to L_simple") — refer there if you want the six-step proof. For this lesson, treat the noise parameterization as the practical default: it is what every production codebase actually trains.

#The Denoising U-Net: Architecture Deep-Dive

#Skip connections: the U-Net spine

Why skip connections matter for diffusion specifically:

• Low-frequency content lives in the bottleneck, but high-frequency detail lives in the early-layer activations. A denoising decoder needs both. Without skip connections the decoder has to reconstruct fine detail purely from a small bottleneck representation, which is exactly the VAE blurriness problem.
• The input x_t and the target ε are spatially aligned at full resolution. A pixel of noise in x_t at coordinate (i, j) is the same pixel of ε to predict at (i, j). The skip connections give the network a direct path from input pixels to output pixels — it does not have to relay every pixel through the bottleneck.
• Gradient flow. Skip connections are residual-style: gradients can backprop through the short path during training, avoiding vanishing-gradient pain in a network that is 30-50 layers deep.

If you remove the skip connections, your "U-Net" becomes a plain encoder-decoder, and the samples come out as recognizable shapes drowning in muddy texture — the classic symptom of a model that knows where things are but not what they look like up close.

#Time conditioning via sinusoidal embedding + MLP

1. Sinusoidal time embedding. Map t to a vector of sines and cosines at geometrically-spaced frequencies — the same [sin(t/10000^{2i/d}), cos(t/10000^{2i/d})] formula as Vaswani et al. 2017. Result: a 128- or 256-dimensional dense vector that encodes t smoothly.
2. MLP projection. Pass that sinusoidal vector through a 2-layer MLP (with SiLU/Swish activation) to produce a "time embedding vector" of the right dimensionality to add into the residual blocks.
3. Injection into every block. Inside each residual block, the time embedding is projected to match the channel count and added to the feature map after the first conv-norm-activation. Every block at every resolution gets to read the clock.

#Residual blocks: GroupNorm + SiLU + Conv

Inside each U-Net stage there are 2-3 residual blocks. The modern recipe (since ADM, Dhariwal & Nichol 2021):

• GroupNorm instead of BatchNorm. BatchNorm depends on batch statistics, which behave badly when your batch contains samples at very different noise levels (which diffusion always does). GroupNorm normalizes within channel groups of a single sample — stable across noise levels.
• SiLU activation (also called Swish) — x · sigmoid(x). Smoother than ReLU around zero, slightly better empirical results in diffusion. Same activation transformers settled on.
• Conv 3x3, dropout (optional), conv 3x3, residual add. Two 3x3 convolutions per block with a residual connection — the same shape Kaiming He's ResNet introduced in 2015.
• Time embedding addition between the two convs, after the first norm + activation. The block reads the clock once and propagates it.

#Self-attention at low resolutions

Pure convolution has a finite receptive field. At early layers and full resolution, that is fine — local denoising is a local operation. But once you have downsampled to 16x16 or 8x8, you want the network to reason globally about the image: "the left side is a face, so the right side should also look face-y", "this is a horse, so its legs should match", "this brushstroke is impressionist, propagate that style everywhere".

Why only at low resolution? Self-attention is O(n²) in sequence length. At 64x64 that is 64² = 4096 tokens, with a 4096² = 16M-element attention matrix — too expensive at every layer. At 8x8 it is 64 tokens, 4096-element matrix — cheap. So you keep the local conv operations at high resolution and pay for global reasoning only at the bottleneck.

#Cross-attention for conditioning

Q = W_Q · spatial_features
K = W_K · text_embeddings
V = W_V · text_embeddings
attention_output = softmax(Q K^T / √d) · V

This is the mechanism that lets Stable Diffusion know that "a cat on a skateboard" means certain spatial regions of the latent should be cat-pixels and others should be skateboard-pixels. The cross-attention layers are interleaved with the self-attention layers in every block at every conditioning-eligible resolution.

#Modern architectures: DiT and MM-DiT

Modern production models that use DiT or DiT-like architectures:

• Stable Diffusion 3 / 3.5 (Stability AI 2024) — uses MM-DiT (Multimodal Diffusion Transformer), where image tokens and text tokens flow through joint transformer blocks instead of being injected via cross-attention. Big quality jump over the SD2 U-Net.
• FLUX.1 (Black Forest Labs 2024) — also MM-DiT, currently the strongest open-weights image model. The team is largely the original Stable Diffusion authors.
• Sora (OpenAI 2024) — DiT applied to video. Patches are now spatiotemporal "cubelets" instead of 2D patches; otherwise it is the same recipe.
• Imagen 3 (Google DeepMind 2024) — uses a transformer backbone with text-conditioning baked in at every layer.

The migration matters because it tells you what the deep-learning track owes the transformer track. A modern diffusion model is much closer to "a ViT trained with a noise-prediction objective" than to "a U-Net". The skip connections become learned residual connections inside transformer blocks; the cross-attention is replaced by joint self-attention over [image_tokens ; text_tokens]; the time conditioning is replaced by AdaLN-Zero. But the conceptual story (fixed encoder chain, shared decoder, noise-prediction loss) is unchanged.

#Training: A Standard Supervised Loop in Disguise

Diffusion training is shockingly straightforward — it looks more like supervised classification training than the GAN dance of the previous lesson. The full algorithm:

for each minibatch x_0 from the dataset:
    t   ~ Uniform({1, ..., T})           # one random timestep per sample
    ε   ~ N(0, I)                        # one Gaussian noise per sample
    x_t = √ᾱ_t · x_0 + √(1 - ᾱ_t) · ε   # forward reparameterization
    ε̂   = ε_θ(x_t, t)                    # forward pass through the U-Net
    L   = ||ε - ε̂||²                     # MSE on noise prediction
    backprop, Adam step

Practical training notes worth knowing:

• EMA on weights. Almost all production diffusion models track an exponential moving average of weights with decay 0.999 or 0.9999 and use the EMA weights for sampling. The training weights are noisier; the EMA is what produces good samples. Standard trick from supervised CV, doubly important here.
• Mixed precision. Train in bfloat16 or float16 with a small float32 master copy of weights. Saves 2x memory and runs 2-3x faster on modern GPUs. The MSE loss is well-conditioned enough that this works without numerical issues.
• Gradient checkpointing on the U-Net activations. Lets you fit larger batches at the cost of an extra forward pass.
• Conditioning dropout (10%). During training, with 10% probability, replace the text conditioning with a null token (an "empty prompt" embedding). This makes the same model usable for both conditional and unconditional generation — the prerequisite for classifier-free guidance at sampling time.

#Hands-On 1: Forward Diffusion and a Toy Denoiser on 2D Data

Time to see this run. The 2D toy case is the cleanest demo — you can plot the data distribution at every timestep and watch it morph into N(0, I).

A few things are worth noticing in this code.

• The forward process has zero parameters. Nothing about q_sample is learned. It is pure noise injection on a fixed schedule. Look at the snapshot plots — the clusters smear out into an isotropic Gaussian, deterministically (in distribution) from the schedule.
• One network handles all timesteps. The same (W1, W2, W3) weights produce the noise prediction at t=5 as at t=95. The time embedding tells it which noise level it is looking at.
• Training is a single MSE. No discriminator, no balance, no warmup tricks. Sample (x_0, t, ε), compute x_t, MSE on predicted noise, step.
• Sampling is iterative. The DDPM sampler walks T=100 reverse steps. The model is called 100 times for each generated sample. This is the cost diffusion pays for stable training and stable sampling: many forward passes per sample.

This is a toy. A production diffusion model replaces the 2-layer MLP with a U-Net (or DiT), the 2D Gaussian mixture with a million-image dataset, T=100 with T=1000 (or fewer with modern flow-matching schedules), and the linear schedule with cosine. But the loop is character-for-character identical.

#Sampling: DDPM, DDIM, and the Steps Trade-Off

#DDPM: Stochastic Reverse Process

The original DDPM sampler walks the full Markov chain in reverse, one step at a time, with a small noise injection at every step:

x_T  ~ N(0, I)
for t from T down to 1:
    z      ~ N(0, I) if t > 1 else 0
    x_{t-1} = (1/√α_t) · (x_t - (β_t / √(1-ᾱ_t)) · ε_θ(x_t, t)) + σ_t · z
return x_0

This is faithful to the variational derivation and the theory is clean. It requires roughly T forward passes through the U-Net per sample — for T=1000 that is 1000 network evaluations to make a single image. At 50ms per pass on a fast GPU that is 50 seconds per sample. Hopeless for interactive use.

#DDIM: Deterministic and Step-Skipping

x̂_0     = (x_t - √(1-ᾱ_t) · ε_θ(x_t, t)) / √ᾱ_t            # estimate x_0
x_{t-1} = √ᾱ_{t-1} · x̂_0 + √(1 - ᾱ_{t-1}) · ε_θ(x_t, t)    # rebuild at smaller t

The picture you should walk away with from this playground:

• DDPM and DDIM at 1000 steps look the same. They produce indistinguishable samples — DDIM is deterministic, DDPM is stochastic, but the marginal distributions match.
• DDIM at 50 steps loses almost nothing visible. This is the 20x speedup that made diffusion practical for production. Stable Diffusion's default is 20-50 DDIM-class steps.
• DDIM at 10 steps is where the trade-off bites. The clusters start to bleed into each other; the model does not get enough denoising passes to fully resolve the structure.
• The trained model is the same in every case. The sampler is a runtime choice. You train once with L_simple at T=1000; you sample with whichever scheduler your latency budget allows.

Production samplers (DPM-Solver++, EDM-Heun, etc.) push the trade-off curve further — they hit DDIM-1000 quality at 15-25 steps by solving the underlying probability-flow ODE with higher-order numerics. The genai-track lesson has the full sampler cheat sheet.

#A Predict-First on Schedule Mismatch

#Comparison: Diffusion vs VAE vs GAN

Three paradigms, one table. This is the cleanest summary of where diffusion sits in the generative-model landscape.

The headline difference is the inference cost. Diffusion pays for stable training and sharp samples by running the decoder T times instead of once. That cost is the only structural weakness, and it is what every "distillation" line of research (LCM, SDXL Turbo, Consistency Models, Stable Diffusion Turbo) is trying to close.

#Classifier-Free Guidance (Deep Dive)

#History: From Sohl-Dickstein to FLUX

#Where the Compute Goes

A practical thing worth knowing: diffusion training is comparable in compute cost to large autoregressive language models. SDXL was trained on roughly 5x10^22 FLOPs (around 256 A100s for a month). FLUX.1 [pro] is on the same order. Sora is significantly larger — OpenAI has not published exact numbers, but estimates put it in the same range as GPT-4 training. The cost is dominated by the U-Net (or DiT) forward and backward passes, multiplied by the number of training steps, multiplied by the batch size, multiplied by the spatial size of x_t.

Inference is the more visible cost. Naive DDPM at T=1000 with classifier-free guidance is 2000 U-Net forward passes per image. A 1024x1024 SDXL pass is roughly 10 GFLOPs of compute on the U-Net plus 50 MFLOPs on the VAE decoder. At 2000 passes that is 20 TFLOPs — about 3 seconds on a single H100. With DPM-Solver++ at 25 steps and bf16 you are down to 1.4 seconds. With SDXL Turbo at 1 step you are at 150ms — fast enough for interactive UIs. The compute envelope is the engineering constraint that determines whether a diffusion model can be exposed as a "real-time" feature versus a "wait two seconds" feature.

#Key Takeaways

• Diffusion is a hierarchical VAE with T fixed Gaussian encoders and one shared learned decoder. The forward process injects noise on a fixed schedule; the reverse process predicts that noise with a single neural network applied T times. Everything else is engineering.
• The denoising U-Net is a stack of ResNet blocks with a time embedding threaded through and self-attention at low resolutions. Skip connections carry high-frequency detail past the bottleneck. Cross-attention injects text conditioning. Modern flagship models replace the U-Net with a Diffusion Transformer (DiT) or Multimodal DiT.
• Training is an MSE on predicted noise. No discriminator, no adversarial balance, no posterior collapse. Sample (x_0, t, ε), compute x_t, predict ε, backprop. Embarrassingly parallel and pleasantly stable.
• Sampling cost is the structural weakness. DDPM needs T forward passes per sample; DDIM and modern samplers cut this to 20-50; LCM and Turbo distill down to 1-4 steps with mild quality loss.
• Classifier-free guidance is the single most important inference-time trick. Train one model that handles both conditional and unconditional inputs, then linearly extrapolate the noise prediction toward the conditional direction at sampling time.

#Forward References

• Generative-AI track lesson on diffusion. Full ELBO-to-L_simple derivation, sampler taxonomy, classifier-free guidance with code, and the latent-diffusion / Stable Diffusion family
• Stochastic Calculus & SDEs in Math Foundations — the continuous-time view of diffusion: Song et al. 2021's score-SDE formulation, the Karras EDM framework, why DDPM is the Variance-Preserving SDE and score matching is the Variance-Exploding SDE, and how Brownian motion and the Fokker-Planck equation underpin every modern sampler
• Track-06 lesson on flow matching and rectified flow. The post-diffusion paradigm that FLUX and SD3 actually use for training, with the same MM-DiT architecture but a different training objective
• Track-06 lesson on controllable generation. ControlNet, IP-Adapter, LoRA for steering the diffusion architecture you learned here