This course has one goal: after finishing it, you can walk to a whiteboard and rebuild a Vision Transformer from an empty file — in PyTorch or in raw NumPy — while explaining why every line is there. That combination (implementation fluency + design reasoning) is exactly what ML research interviews probe.

How to use it

Work through the modules in order. You said you know deep learning well but are rusty on attention, so modules 02–03 rebuild attention from first principles before anything vision-specific appears. Every module ends with a short quiz; your score and completion state are saved in this browser, so you can leave and come back. The interactive panels aren't decoration — each one exists because the concept it shows is a common interview probe (softmax temperature, quadratic cost, patch-size trade-offs, parameter counting).

Two rules that make this stick:

• Type the code yourself. Reading module 05 is not the same as reproducing it. After each implementation module, close the page and rewrite the module from memory. The copy button on code blocks is for checking your version, not for skipping the typing.
• Say the shapes out loud. Interviewers watch whether you track tensor shapes. Every code block here annotates shapes; module 04's shape-tracer drills them.

The route

01 motivates the whole enterprise (inductive bias, and why "just use a transformer" was a radical bet). 02–03 build scaled dot-product attention, multi-head attention, and the pre-norm transformer block. 04 assembles ViT itself — patch embedding, the [CLS] token, position embeddings — with a full parameter-count walkthrough. 05 is the complete PyTorch implementation; 06 re-derives it in NumPy including backprop through softmax, LayerNorm, and attention. 07 covers why ViTs are data-hungry and how they're actually trained. 08 collects the "interesting details" that separate good candidates from great ones. 09 traces DeiT → Swin → MAE → DINO. 10 is pure interview practice.

By 2020, a decade of computer vision rested on two assumptions hard-wired into CNNs: locality (nearby pixels are related, so use small kernels) and translation equivariance (a cat is a cat wherever it sits, so share weights across positions). These are inductive biases — prior knowledge the architecture enforces so the model doesn't need to learn it from data.

The ViT paper (Dosovitskiy et al., "An Image is Worth 16×16 Words", ICLR 2021) asked: what if we remove almost all of that? Cut the image into patches, treat the patches as a sequence of tokens, and feed them to a plain transformer encoder — nearly unchanged from the NLP transformer of Attention Is All You Need (2017). A transformer layer has no notion that patch 15 sits next to patch 16. Every token can attend to every other token from layer one; where things are is told to the model only through a learned position embedding.

The paper's central empirical result confirms the trade exactly. Trained only on ImageNet-1k (1.3M images), ViT loses to comparable ResNets — it overfits, because nothing stops it. Trained on ImageNet-21k (14M) it pulls even. Trained on JFT-300M (303M images, Google-internal) it wins decisively, and it keeps improving where CNNs saturate. And crucially, it was cheaper to train per unit of accuracy at scale, because dense matmuls over tokens map onto accelerators better than convolution pipelines.

What ViT keeps, and the one place convolution survives

ViT is not "zero inductive bias." Three vision priors survive: the patchification itself assumes pixels arrive in a 2D grid and that a 16×16 tile is a sensible atomic unit; the position embeddings are learned per-location, so the model can rediscover 2D structure (and does — see module 08); and the whole model is applied to a fixed resolution. Also — a fact interviewers enjoy — the patch projection in every practical implementation is a convolution: a Conv2d with kernel = stride = patch size, which is mathematically identical to "slice and apply a shared Linear." Weight sharing across patches is a conv-style bias ViT quietly retains.

Forget images for two modules. You have a sequence of N vectors x₁ … x_N, each of dimension D, stacked into a matrix X of shape (N, D). We want each token's new representation to be built from the tokens that are relevant to it — where relevance is learned, not hard-coded.

Step 1 — three roles, three projections

Each token plays three roles at once, produced by three learned linear maps:

The retrieval metaphor is exact and worth saying in an interview: a query is what this token is looking for; a key is how a token advertises what it contains; a value is the content it hands over if selected. Why three separate matrices instead of comparing raw x vectors? Because relevance is asymmetric and task-dependent: "what a verb wants" ≠ "what a verb offers." Separate projections let the model place queries and keys in a space where dot product = usefulness, independent of what content gets transmitted.

Step 2 — similarity by dot product

Step 3 — softmax rows into mixing weights

Each row of A is a probability distribution: non-negative, sums to 1. Row i of the output is a weighted average of all value vectors, weighted by how relevant each token is to token i. That's the whole mechanism. Note what it costs: S is an N×N matrix, so compute and memory are O(N²·d) — the fact that drives half the follow-up research in module 09.

Self-attention, seen whole

"Self" just means Q, K, V all come from the same sequence. The full einsum chain, with shapes:

Multi-head attention: parallel subspaces, same price

A single softmax produces a single mixing pattern per token — but a patch may simultaneously need "the other patches of this object," "the texture context," and "the global gist." So we split the model dimension D into H heads of size d_h = D/H (ViT-Base: D = 768, H = 12, d_h = 64). Each head gets its own Q, K, V projections into its 64-dim subspace, runs attention independently, and the H outputs are concatenated back to D and mixed by a final linear W_O.

In code, nobody loops over heads. The trick — worth performing smoothly on a whiteboard — is one fused projection to 3D followed by a reshape so the head axis sits beside the batch axis, making every head's attention a single batched matmul:

The MLP: where most of the parameters live

After attention mixes information across tokens, a two-layer MLP processes each token independently: expand D → 4D, apply GELU, project 4D → D. The 4× ratio is inherited from the original transformer and survives because it keeps working. Two-thirds of a ViT block's parameters sit here, not in attention — a fact that surprises people and therefore gets asked.

Pre-norm and the residual stream

The block wires these together with two residual connections and LayerNorms placed before each sublayer (pre-norm) — a deliberate departure from the 2017 post-norm transformer:

Patch embedding: images become sentences

Take an image (3, 224, 224), slice it into a 14×14 grid of 16×16 patches, flatten each patch to a 768-vector (16·16·3), and multiply by a learned (768, D) matrix. The result: 196 tokens of dimension D. A "word" is a 16×16 tile of pixels — hence the paper's title.

Every real implementation does the slicing and projection in one op: nn.Conv2d(3, D, kernel_size=P, stride=P). With kernel = stride, windows never overlap, so the convolution is "shared Linear applied to each patch" — same math, one fused kernel. Being able to state this equivalence precisely is a small but reliable interview win.

The [CLS] token: a learned blank slate for the answer

Borrowed from BERT: prepend one learned D-dim vector to the 196 patch tokens. It participates in every attention layer like any token, and at the end the classification head reads only its final state. Sequence length becomes 197.

Position embeddings: telling a set it's a grid

Attention is permutation-equivariant — shuffle the 196 patches and the outputs shuffle identically, so ViT would classify a jigsaw-scrambled image the same as the original. The fix: add a learned vector p_i to token i before the first block. One (197, D) table, trained end-to-end.

The whole machine

Build order matters: leaf modules first, assembly last. That's also the order to use at a whiteboard, because each piece is independently testable — after writing Attention, you can already assert its output shape equals its input shape.

1 · Patch embedding

Line by line. flatten(2) merges dims 2 and 3 (the 14×14 grid) into 196 — row-major, so token order is "reading order." The transpose exists purely for convention: transformers everywhere expect (batch, sequence, features). No parameters beyond the conv's D·3·P² + D.

2 · Multi-head self-attention

Why the fused qkv? Three separate Linears would produce identical math; one (D, 3D) matmul is simply faster (one kernel launch, better GEMM shape) and is the idiom you'll see in timm, nanoGPT, and every serious codebase. Why scale before softmax, not after the matmul chain? It must hit the logits; scaling attention outputs later wouldn't fix softmax saturation. Note also transpose(1, 2) before the final reshape — forgetting it silently interleaves head outputs across tokens, a classic bug that still type-checks. (In production you'd let F.scaled_dot_product_attention pick a FlashAttention kernel — see module 08 — but interviews want this explicit version.)

3 · MLP and the block

The two comments in Block.forward are the sentence to say in interviews: attention is the only cross-token operation in the entire network; everything else is per-token. That factorization — communicate, then compute — is the transformer.

4 · Assembly

Details that get probed. (1) expand vs repeat: expand creates a broadcast view — no memory copy; gradients from all B images still accumulate into the single shared [CLS] parameter. (2) Position embedding is added, not concatenated — it keeps D fixed and lets position and content share the same subspaces; concatenation would work but grows every weight matrix downstream. (3) Why a final norm? In pre-norm nets the stream itself is never normalized, so its magnitude grows with depth; the last LN tames it before the head reads. (4) trunc_normal_(std=0.02): small init keeps early attention logits small (near-uniform attention at start), which stabilizes the first steps.

5 · Prove it works

The architecture is the one you just built; what's new is the backward pass. Three derivations do all the work — softmax, LayerNorm, and attention — and everything else is the chain rule over matmuls. One convention first: for Y = XW + b, the three gradients every interviewer expects instantly:

Memorize by shape: GWᵀ is the only product with X's shape; XᵀG the only one with W's. Shapes are the error-correcting code of backprop.

Derivation 1 — softmax

For p = softmax(z): ∂pᵢ/∂zⱼ = pᵢ(δᵢⱼ − pⱼ). Contract that Jacobian with the upstream gradient g and it collapses to a one-liner:

Note the max-subtraction: softmax is invariant to shifting its input, and shifting by the max means the largest exponent is e⁰ = 1 — no overflow ever. Interviewers ask for exactly this trick.

Derivation 2 — LayerNorm

Forward: x̂ = (x − μ)/√(σ² + ε), y = γ·x̂ + β, with μ, σ² over the feature axis. The subtlety: x affects y through three routes — directly, through μ, and through σ². Collecting all three terms (the full derivation is a standard exercise; do it once on paper) gives a closed form worth memorizing:

Read the return as: "the naive gradient, corrected so the output can't change the mean or the variance the layer just removed." The two subtracted means are exactly the projections onto the directions LayerNorm quotients out.

Derivation 3 — attention backward

Forward chain: S = QKᵀ/√d → A = softmax(S) → O = AV. Reverse each arrow, remembering the product rule at both matmuls:

Two things to internalize. Every transpose in backward is the inverse of one in forward — write the forward reshapes down and mirror them, don't improvise. And the residual connections in the block make backward almost embarrassingly simple: x + f(x) means the incoming gradient is copied to both paths and summed:

There is the vanishing-gradient story in executable form: whatever the branches contribute, the identity term guarantees the gradient that leaves this block is at least the gradient that entered it, unmodified. That's the sentence to say when asked "why do residuals help?"

Patchify without a conv, and the check that makes it true

The optimizer: AdamW, weight decay done right

Every serious ViT trains with AdamW. Two reasons it beats SGD+momentum here (while ResNets are happy with SGD): transformer gradients are far less isotropic — embedding tables, LayerNorm gains, and attention projections have wildly different curvature, and Adam's per-parameter scaling absorbs that; and the "W" matters: classic Adam entangles L2 regularization with the adaptive scaling (the penalty gets divided by √v̂, so high-variance weights are barely regularized). AdamW decouples decay — w ← w − lr·(update + λ·w) — restoring true shrinkage. ViT recipes use unusually large decay (0.1–0.3) as an extra regularizer, and exempt biases, LayerNorm parameters, [CLS], and position embeddings from it (decaying a bias or a norm gain toward zero fights their job, not overfitting).

The schedule: warmup, then cosine

Regularization scales opposite to data

The paper's key recipe table reads like a see-saw. Pre-training on JFT-300M: almost no regularization (no dropout, light decay) — 300M images regularize by existing. Training on ImageNet-1k alone: everything at once — RandAugment, Mixup, CutMix, random erasing, label smoothing 0.1, stochastic depth, dropout, heavy decay. The DeiT paper (module 09) is essentially the discovery of the exact augmentation cocktail that replaces 300M images.

Fine-tuning and the resolution trick

Standard practice: pre-train at 224, fine-tune at 384+. Higher resolution with the same patch size means more tokens (224→384: 196→576 patches) — and a position-embedding table of the wrong length. The fix is elegant and a guaranteed interview question: reshape the 196 learned embeddings back to their 14×14 grid, bicubically interpolate to 24×24, flatten to 576. This works precisely because (module 08) the learned table encodes a smooth 2D layout — smooth signals interpolate well. The [CLS] embedding is carried over unchanged, and a brief fine-tune lets the model adapt.

The reference recipes, condensed

1 · Position embeddings secretly learn the grid

Compute cosine similarity between the learned embedding of each patch position and every other. The result (ViT paper, Fig 7 center): each position is most similar to its 2D neighbors, with visible row and column bands — although training never told the model the tokens came from a grid. The removed inductive bias was re-derived from data. This is also, as module 07 showed, why bicubic interpolation of the table is sound.

2 · Attention distance: CNN-like from the bottom, global from the start

Define a head's "attention distance" as the average image-space distance between a query patch and the patches it attends to, weighted by attention — a receptive-field size for attention. The paper's measurement: in early layers, heads split — some are local (small distance, conv-like), others are already global. By mid-depth, nearly all heads are global. So ViT learns a local-to-global processing hierarchy where CNNs have it imposed — but keeps a few global shortcuts from layer one, which a CNN cannot have. With small datasets the local heads fail to emerge, which is one lens on why ViT under-performs there.

3 · O(N²) in practice, and FlashAttention

For ViT-B/16 at 224², N = 197 — the attention matrix is 197×197 ≈ 39k entries per head: trivial. The quadratic only bites at high resolution (1024² with patch 16 → N = 4096 → 16.8M entries per head per layer). Two responses coexist: change the architecture (Swin's windows, module 09) or change the kernel — FlashAttention computes exact softmax attention without ever materializing the N×N matrix in slow GPU memory, tiling the computation through on-chip SRAM with the online-softmax trick (a running max and running sum let you renormalize incrementally). Memory drops from O(N²) to O(N); the math is unchanged. The one-liner: "FlashAttention is an IO optimization, not an approximation."

4 · [CLS] vs global average pooling

Both work; the paper's appendix shows equal accuracy with separately tuned LRs — a nice example of hyperparameter coupling masquerading as an architecture effect. Practical modern note: for dense tasks (segmentation, depth) you use the patch tokens anyway, and models like DINOv2 expose both a [CLS] summary and patch-level features.

5 · Registers: the attention-sink discovery

"Vision Transformers Need Registers" (Darcet et al., 2023/24): large ViTs spontaneously repurpose a few low-information background patches as scratch storage — those tokens get abnormally high attention norms and pollute attention maps with spiky artifacts. The fix is almost comic: append a few extra learned tokens (registers) with no input and no output role, giving the model somewhere to put its global scratch data. Artifacts vanish, dense-task features improve. Lesson to quote: the [CLS] idea generalizes — transformers benefit from explicit "memory" tokens, and if you don't provide them the model improvises them out of your image.

6 · Robustness and what ViTs actually look at

ViTs lean less on high-frequency texture than CNNs of their era: they're notably more robust to patch-level occlusion and adversarial patches, and attention lets a single layer route around a masked region. But the cliché "CNNs = texture, ViTs = shape" oversimplifies — training recipe and augmentation move this dial more than architecture does (heavily-augmented CNNs also become shape-biased). Saying that nuance is exactly the kind of calibrated claim interviewers reward.

7 · Small numbers worth having memorized

Also: head_dim is 64 in every variant (D/H = 64 throughout — a de-facto constant across transformers, big and small); sequence length 197 at 224/patch-16; ViT-L/16 fine-tuned from JFT hit 87.76% ImageNet top-1 in the original paper, and ViT-H/14 88.55%.

DeiT (2021) — "must it really take 300M images?"

Data-efficient Image Transformers made ViT trainable on ImageNet-1k alone, in ~3 days on one 8-GPU node — no JFT. Two ingredients. First, the aggressive augmentation/regularization cocktail from module 07 (RandAugment, Mixup, CutMix, erasing, stochastic depth, repeated augmentation) — supplying synthetically the sample diversity that JFT supplied literally. Second, distillation through a token: alongside [CLS], append a learned distillation token whose head is trained to match a RegNet CNN teacher's predictions (hard-label distillation worked best). The CNN's inductive biases leak into the ViT through the distillation loss — the student gets convolutional priors without convolutional architecture. DeiT-B reached ~83.4% ImageNet top-1 with distillation; the "data hunger" objection to ViT died here.

Swin (2021) — "how do we afford dense prediction?"

Detection and segmentation need high-resolution feature maps, where global attention's O(N²) is unpayable (module 08's arithmetic). Swin restores the two CNN structural ideas ViT discarded, in transformer form. Windowed attention: compute self-attention only within 7×7-patch local windows — cost becomes linear in image size. Shifted windows: alternate layers offset the window grid by half a window, so information crosses window borders every other layer (otherwise windows would be permanently isolated islands). Hierarchy: stages of patch-merging (concatenate each 2×2 neighborhood of tokens, project down) halve resolution and double channels — a feature pyramid, /4 → /8 → /16 → /32, that drops straight into detection frameworks like a ResNet backbone. Swin also swaps learned absolute positions for relative position biases added to attention logits. Trade to articulate: Swin re-buys CNN-like efficiency by re-imposing locality — a deliberate step back from ViT's purity, and the price is that plain ViT scales better with data while Swin wins at dense tasks per FLOP.

MAE (2022) — "what's the BERT moment for vision?"

Masked Autoencoders answer: mask 75% of patches (BERT masks 15% of words — images are far more redundant, so the task must be made much harder to force semantic learning). The asymmetric design is the engineering jewel: the encoder (a plain ViT) sees only the 25% visible patches — no mask tokens — so pre-training compute drops ~4×; a lightweight decoder receives the encoded patches plus shared learnable [MASK] tokens with position embeddings and reconstructs raw pixels (per-patch-normalized, MSE on masked patches only). After pre-training, throw the decoder away. Self-supervised ViT-Huge fine-tuned to 87.8% ImageNet — proof that ViTs can pre-train from unlabeled pixels alone, and that the "needs labels at scale" era was over.

DINO (2021) & DINOv2 (2023) — "can features be useful without any fine-tuning?"

DINO is self-distillation: a student ViT sees aggressive local+global crops, a teacher — an exponential moving average of the student's own weights — sees mild global crops, and the student matches the teacher's soft output distribution (cross-entropy after centering + sharpening, the two tricks that prevent representational collapse). No labels anywhere. The astonishing emergent result: the [CLS] attention maps of the final layer segment objects — clean foreground masks nobody asked for — and k-NN classification on frozen features works nearly as well as linear probing. DINOv2 scaled the recipe (curated 142M-image dataset, ViT-g, plus an iBOT-style masked-token loss and engineering like KoLeo regularization) into general-purpose frozen backbones whose features transfer to depth, segmentation, and retrieval without fine-tuning — the closest vision has to a foundation-model feature extractor, and (with registers added) the standard baseline today.

Rapid-fire cards

Coding drills — do these cold, in order of difficulty

• Drill 1. scaled_dot_product_attention(q, k, v) for shape (B, H, N, d), NumPy or PyTorch, with stable softmax. Target: under 4 minutes, no reference.
• Drill 2. Full Attention module including the qkv-fusion reshape dance. The permute indices must come out right first try — narrate shapes as you type.
• Drill 3. patchify as pure reshape/transpose (no conv, no loop), plus its exact inverse unpatchify. This is also MAE-interview territory.
• Drill 4. Entire ViT forward in PyTorch from an empty file, then verify 86,567,656 on the B/16 config. Target: under 25 minutes.
• Drill 5. NumPy LayerNorm forward + backward, gradient-checked against finite differences in float64.
• Drill 6. Position-embedding interpolation for a resolution change (module 07's snippet) — from memory.
• Drill 7 (boss level). NumPy attention backward, checked against autograd. When this matches to ~1e-16, you own this architecture.

Discussion questions to rehearse out loud

• Design a ViT for 1024×1024 medical images on one A100 — walk through patch size, windowing vs FlashAttention, pos-embed strategy, and what you'd pre-train on given 50k labeled + 5M unlabeled scans. (Good answers combine modules 04, 07, 08, 09.)
• Your ViT trains fine but attention maps are noisy spikes on background patches — what's happening and what are two fixes? (Registers; or train smaller/shorter — module 08.)
• Argue both sides: "Swin's return to locality was the right call" vs "plain ViT + scale was the right call." (The 2023+ evidence — DINOv2, MAE-scaled plain ViTs, FlashAttention erasing much of the efficiency gap — favors plain ViT for representation learning; Swin persists in latency-bound dense pipelines.)
• Why did BERT-style masked pre-training take two years longer to work in vision than in NLP? (Redundancy → needed 75% masking + pixel targets; also tokenization — words are discrete and semantic, patches aren't, which is why BEiT tried discrete visual tokens first and MAE showed raw pixels suffice with the right ratio.)

Where to go next

Primary sources, in reading order: Attention Is All You Need (Vaswani 2017) · An Image is Worth 16×16 Words (Dosovitskiy 2020) · Training data-efficient image transformers (Touvron 2021) · Swin Transformer (Liu 2021) · Masked Autoencoders Are Scalable Vision Learners (He 2021) · Emerging Properties in Self-Supervised Vision Transformers (Caron 2021) · DINOv2 (Oquab 2023) · Vision Transformers Need Registers (Darcet 2023). For code archaeology, read timm's vision_transformer.py — the field's de-facto reference implementation.

You already know CNNs and transformers. The new material is the learning problem: how to turn examples of real images into a model that can create new ones. We will learn one concrete version first—an image-space denoising diffusion probabilistic model, or DDPM—and then change one component at a time.

By the end, you should be able to explain diffusion in two minutes, implement training and sampling from an empty file, diagnose a failing run, and reason about guidance, latent diffusion, and ControlNet. You do not need stochastic calculus to reach that point.

Use the course actively

Each lesson has a concrete outcome and an exercise. Try the exercise before opening the hint or solution. Checkpoints test distinctions that often cause implementation bugs. Mark a lesson complete when you can explain it without reading. Progress and quiz answers are stored in this browser when local storage is available; the downloaded file contains the course, not your personal progress.

First pass: read 01–06 and run the lab. Second pass: add 07–10 and implement one extension. Interview pass: use 11–15 to connect the math, trade-offs, and failure modes. Time estimates beside exercises are practice budgets, not promises about training time.

One notation, everywhere

What you will build

The lab trains on freshly generated rings and squares. This makes the data distribution inspectable and removes setup friction. The model, loss, reverse process, exponential moving average, validation, and checkpointing are real. The small dataset is a teaching choice; attractive samples from it do not establish performance on photographs.

The entire course, including runnable code and interactive explanations, works offline in this file. External reading links and installing PyTorch require an internet connection. The vision transformer course remains available in the course selector.

Suppose you have many images of handwritten digits. Training a network to map arbitrary noise directly to a digit leaves a difficult question: which digit is the correct target for this particular noise? Diffusion gives us a much easier supervised task. Take a real digit, add noise that you generated yourself, and ask a network to predict that noise.

At low noise, the network can use local edges. At high noise, it must use what it has learned about plausible images. At generation time, we start with noise and repeatedly use those predictions to make a slightly cleaner image. The same network is reused at every noise level.

Corrupting an image in one line

There are only two coefficients: how much signal remains and how much noise we add. If Aₜ = 1, we have the original image. If Aₜ is near zero, we have almost pure noise. A noise schedule is the table that determines Aₜ at each time.

Why square roots? Scaling a random variable by a multiplies its variance by a². The coefficients therefore specify a signal variance fraction Aₜ and a noise variance 1−Aₜ. This does not mean every real dataset or every corrupted image has variance exactly one; that also depends on the data.

A numerical anchor: when Aₜ = 0.64, the signal multiplier is 0.8 and the noise multiplier is 0.6. For one pixel with x₀ = 0.5 and ε = −1, the corrupted value is 0.8×0.5 + 0.6×(−1) = −0.2. Negative values are normal. Do not clip xₜ to the original image range.

Why a “process” if we can jump directly?

The original construction adds a small amount of noise repeatedly:

Independent Gaussian noises combine into another Gaussian. After t steps, the original signal is multiplied by √(α₁α₂…αₜ), and the accumulated noise variance is 1−α₁α₂…αₜ. That gives the one-line formula above. It is the distribution of xₜ that agrees; separately drawn noise does not reproduce one particular simulated trajectory.

During training, use the direct formula. There is no need to run t corruption steps or store a chain. For each batch of B images, draw B times and B noise tensors. The computational cost of the training step is essentially one denoiser forward and backward pass, whether the schedule has 200 or 1,000 entries.

Why use this kind of generator?

Diffusion offers a straightforward supervised training objective and flexible conditioning, with an expensive iterative sampling loop. Other model families make different trade-offs. These are broad tendencies, not a ranking that holds for every modern system.

A schedule is part of the model contract

A linear β schedule is simple, but its behavior depends on T. Copying the same endpoints while reducing T can leave a lot of image signal at the terminal step. The lab uses a cosine-shaped cumulative signal schedule, converted into valid β values. Later sampling must use the schedule the model was trained with; you may choose fewer sample times using a suitable update.

Foundation: DDPM, forward-process formulation; Improved DDPM, cosine schedule.

We created xₜ from x₀ and a known ε. Feed the network xₜ and t, and minimize mean squared error between its prediction and ε. The clean image is used to construct the example; it is not an input to this unconditional denoiser.

That expectation just describes repeated random training batches: choose data, choose a noise level, choose noise, compute the error. The mean in the lab averages over images, channels, and pixels. A single training example needs only one randomly selected time; across updates, the model sees the whole range.

i = torch.randint(T, (B,), device=device)    # different level per image
noise = torch.randn_like(x0)                # independent Gaussian entries
A = alpha_bar[i][:, None, None, None]       # [B,1,1,1]
xt = A.sqrt() * x0 + (1 - A).sqrt() * noise
noise_hat = model(xt, i)                    # [B,C,H,W]
loss = F.mse_loss(noise_hat, noise)

The time input is necessary: a faint shape at a low noise level and a faint shape at a high noise level call for different corrections. Think of t as telling the network how much to trust the input.

Can it really recover the exact random noise?

Not in general. Several clean images and noise realizations can produce the same noisy observation. Under squared error, the best possible predictor is the conditional average E[ε | xₜ,t]. It uses learned image structure to estimate which part of the observation is noise.

This distinction explains two apparent paradoxes. First, nonzero loss can remain even for a very good model. Second, a one-step clean-image estimate at a high noise level can be blurry or unrecognizable. It averages over possibilities. Generation works by following a sequence of intermediate states, instead of treating that uncertain average as a final image.

Turn a noise prediction into an image estimate

Rearrange the corruption formula:

At Aₜ = 0.64, the earlier pixel had xₜ = −0.2. If the network predicts ε̂ = −1, then x̂₀ = (−0.2 + 0.6)/0.8 = 0.5, exactly the clean value. A prediction error gets amplified when Aₜ is tiny because we divide by √Aₜ. This is why you should not judge the entire model by clean-image reconstructions at its noisiest time.

Three targets, one corruption equation

You will encounter networks that predict ε, x₀, or a quantity called v. They can represent the same information at interior noise levels, but their losses emphasize noise levels differently.

Here a = √Aₜ and σ = √(1−Aₜ), so a²+σ² = 1. For v-prediction, ε̂ = σ xₜ + a v̂ as well. This diffusion v is not automatically the same target as the straight-path velocity used in flow matching in lesson 12.

A denoiser consumes an image-shaped tensor and produces an image-shaped tensor. A U-Net is a convenient choice: downsampling gives a broad view of the image, and skip connections bring fine spatial information back into the upsampling path. Its output is a noise prediction, so it has the same number of channels as the input and no sigmoid or softmax.

Turn t into something a block can use

We turn the scalar time into sine and cosine features at several frequencies, then pass those features through a small MLP. Every residual block projects that shared embedding into its own channel width and adds it to the feature map. The addition broadcasts over height and width.

emb = time_mlp(sinusoidal_features(i))  # [B,D]
h = conv1(F.silu(norm1(x)))             # [B,C,H,W]
h = h + time_projection(emb)[:, :, None, None]
h = conv2(F.silu(norm2(h)))
return shortcut(x) + h

This is an explanatory block sketch; the executable implementation is in lesson 05. A more expressive variant predicts a scale and a shift and applies (1 + scale) * norm(h) + shift. The important part is that time affects the computation throughout the network, not just the final layer.

The lab's shape contract

Design choices you should be able to defend

• GroupNorm: does not pool statistics across images at different noise levels and behaves consistently at small batch sizes. It is a convenient choice, not a theorem that BatchNorm cannot work.
• SiLU: a smooth nonlinearity with a simple implementation. Its exact choice is less important than correct conditioning and a sound training loop.
• Nearest-neighbor upsampling followed by a convolution: makes spatial sizes easy to track. Other upsampling layers are possible.
• No attention in this first model: the 32×32 shape task is enough to learn the diffusion mechanics with a small convolutional U-Net. Larger image generators often add attention for long-range interactions and text conditioning.

Before training, pass a batch through the network, assert that input and output shapes match, compute a scalar loss, and verify finite gradients. A batch-size-one test is useful because accidental dimension squeezing often appears there.

We have a trained denoiser and a noisy image. How do we make it a little cleaner? First estimate x₀ from the predicted noise. Then use xₜ and that estimate to construct the next state. Repeat from T down to 1.

There is one helpful mathematical fact: if we knew the clean image x₀, the distribution of the previous noisy state given xₜ would be a Gaussian with known coefficients. We cannot use the true x₀ when generating, but we can insert the network's estimate into the mean.

Do not memorize four lines before understanding their job. The first two weights mix the predicted clean image with the current state to get a less noisy state. The variance tells us how much fresh uncertainty to add. The derivation can wait until lesson 11.

The DDPM sampling loop

1. Draw x_T from standard Gaussian noise. This matches the training endpoint approximately when A_T is sufficiently close to zero.
2. Predict ε̂ = εθ(xₜ,t), then compute x̂₀.
3. Set the reverse mean to c₀,ₜ x̂₀ + cₜ,ₜ xₜ.
4. For t > 1, sample xₜ₋₁ = mean + √β̃ₜ z with fresh Gaussian z. At t = 1, return the mean.

The model in the lab uses fixed posterior variance β̃ₜ. Learning a reverse variance is a different design choice. The formula above is the exact forward posterior variance conditioned on x₀, not a claim that the true reverse distribution given only xₜ is always that Gaussian.

Three distinctions that prevent most sampler bugs

• Predicted x₀ versus next xₜ₋₁: x̂₀ aims at the clean endpoint. The sampler usually moves only partway there. Replacing xₜ by x̂₀ on every iteration is not DDPM.
• βₜ versus β̃ₜ: the forward noise variance and the posterior variance are different. Use the variance your sampler was designed for.
• Standard deviation versus variance: multiply fresh noise by √β̃ₜ, not β̃ₜ. Suppress fresh noise on the final iteration.

For pixel images normalized to [−1,1], clipping x̂₀ into that range can stabilize sampling; the lab does so. This is a heuristic and changes the exact un-clipped update. Do not blindly clip latent tensors to [−1,1]; their representation has a different range.

Save the code below as diffusion_lab.py, or use the download button. It requires Python 3.10+ and PyTorch 2.x. It uses no pretrained model, diffusion library, or downloaded dataset. If you need PyTorch, follow its official installation selector for your operating system and accelerator.

"""A small pixel-space DDPM. Requires Python 3.10+ and PyTorch 2.x.
Run: python diffusion_lab.py --steps 3000 --out run
Smoke check: python diffusion_lab.py --steps 2 --batch 4 --T 20 --out smoke
The default data are generated locally; no dataset download is needed.
"""
# PART 1: imports, data, and the denoiser
import argparse
import copy
import math
from pathlib import Path
import torch
from torch import nn
import torch.nn.functional as F


def shapes(batch, rng):
    """Fresh 32x32 rings and filled squares in [-1, 1], shape [B,1,32,32]."""
    axis = torch.linspace(-1, 1, 32)
    yy, xx = torch.meshgrid(axis, axis, indexing="ij")
    center = torch.rand(batch, 2, 1, 1, generator=rng) * 0.8 - 0.4
    size = torch.rand(batch, 1, 1, generator=rng) * 0.20 + 0.22
    dx, dy = xx - center[:, 0], yy - center[:, 1]
    radius = (dx.square() + dy.square()).sqrt()
    ring = torch.sigmoid((0.055 - (radius - size).abs()) * 80)
    square = torch.sigmoid((size - torch.maximum(dx.abs(), dy.abs())) * 80)
    kind = torch.randint(0, 2, (batch, 1, 1), generator=rng).bool()
    return (torch.where(kind, ring, square) * 2 - 1).unsqueeze(1)


class TimeEmbedding(nn.Module):
    def __init__(self, width=64):
        super().__init__()
        freq = torch.exp(-math.log(10000) * torch.arange(width // 2) / (width // 2 - 1))
        self.register_buffer("freq", freq)
        self.mlp = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width))

    def forward(self, t):
        phase = t.float()[:, None] * self.freq[None]
        return self.mlp(torch.cat([phase.sin(), phase.cos()], dim=-1))


class ResBlock(nn.Module):
    def __init__(self, cin, cout, time_dim=64):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, cin)
        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.time = nn.Linear(time_dim, cout)
        self.norm2 = nn.GroupNorm(8, cout)
        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
        self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity()

    def forward(self, x, emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.time(F.silu(emb))[:, :, None, None]
        h = self.conv2(F.silu(self.norm2(h)))
        return self.skip(x) + h


class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.time = TimeEmbedding()
        self.stem = nn.Conv2d(1, 32, 3, padding=1)
        self.enc1 = ResBlock(32, 32)
        self.down1 = nn.Conv2d(32, 64, 4, stride=2, padding=1)
        self.enc2 = ResBlock(64, 64)
        self.down2 = nn.Conv2d(64, 128, 4, stride=2, padding=1)
        self.mid = ResBlock(128, 128)
        self.up2 = nn.Conv2d(128, 64, 3, padding=1)
        self.dec2 = ResBlock(128, 64)
        self.up1 = nn.Conv2d(64, 32, 3, padding=1)
        self.dec1 = ResBlock(64, 32)
        self.out = nn.Sequential(nn.GroupNorm(8, 32), nn.SiLU(), nn.Conv2d(32, 1, 3, padding=1))

    def forward(self, x, t):
        emb = self.time(t)                            # [B,64]
        s1 = self.enc1(self.stem(x), emb)             # [B,32,32,32]
        s2 = self.enc2(self.down1(s1), emb)           # [B,64,16,16]
        h = self.mid(self.down2(s2), emb)            # [B,128,8,8]
        h = self.up2(F.interpolate(h, scale_factor=2, mode="nearest"))
        h = self.dec2(torch.cat([h, s2], dim=1), emb)
        h = self.up1(F.interpolate(h, scale_factor=2, mode="nearest"))
        h = self.dec1(torch.cat([h, s1], dim=1), emb)
        return self.out(h)                           # [B,1,32,32], no sigmoid

# PART 2: schedules and one training loss
class Diffusion(nn.Module):
    def __init__(self, T=200):
        super().__init__()
        if T < 2:
            raise ValueError("T must be at least 2")
        self.T = T
        # Float64 for construction, float32 buffers for the training arithmetic.
        u = torch.linspace(0, 1, T + 1, dtype=torch.float64)
        f = torch.cos((u + 0.008) / 1.008 * math.pi / 2).square()
        target_ab = f / f[0]
        beta = (1 - target_ab[1:] / target_ab[:-1]).clamp(max=0.999)
        alpha = 1 - beta
        ab = alpha.cumprod(0)                        # index i means math time t=i+1
        prev = torch.cat([torch.ones(1, dtype=ab.dtype), ab[:-1]])
        values = dict(beta=beta, alpha=alpha, ab=ab,
                      posterior_var=beta * (1 - prev) / (1 - ab),
                      c0=beta * prev.sqrt() / (1 - ab),
                      ct=alpha.sqrt() * (1 - prev) / (1 - ab))
        for name, value in values.items():
            self.register_buffer(name, value.float())

    @staticmethod
    def at(values, i):
        return values[i][:, None, None, None]         # [B,1,1,1]

    def q_sample(self, x0, i, eps):
        ab = self.at(self.ab, i)
        return ab.sqrt() * x0 + (1 - ab).sqrt() * eps


def loss_batch(model, diffusion, batch, rng, device):
    x0 = shapes(batch, rng).to(device)
    i = torch.randint(diffusion.T, (batch,), generator=rng).to(device)
    eps = torch.randn(x0.shape, generator=rng).to(device)
    xt = diffusion.q_sample(x0, i, eps)
    pred = model(xt, i)
    assert pred.shape == eps.shape
    return F.mse_loss(pred.float(), eps.float())

# PART 3: two samplers; both reuse the same trained denoiser
@torch.no_grad()
def sample(model, diffusion, n=16, method="ddim", steps=50, seed=123):
    model.eval()
    device = next(model.parameters()).device
    rng = torch.Generator().manual_seed(seed)
    x = torch.randn(n, 1, 32, 32, generator=rng).to(device)
    if method == "ddpm":
        indices = list(range(diffusion.T - 1, -1, -1))
    elif method == "ddim":
        steps = min(max(2, steps), diffusion.T)
        indices = torch.linspace(diffusion.T - 1, 0, steps).round().long().tolist()
    else:
        raise ValueError("method must be ddpm or ddim")
    for k, i in enumerate(indices):
        t = torch.full((n,), i, device=device, dtype=torch.long)
        eps = model(x, t)
        ab = diffusion.ab[i]
        x0 = ((x - (1 - ab).sqrt() * eps) / ab.sqrt()).clamp(-1, 1)
        if method == "ddpm":
            mean = diffusion.c0[i] * x0 + diffusion.ct[i] * x
            if i > 0:
                noise = torch.randn(x.shape, generator=rng).to(device)
                x = mean + diffusion.posterior_var[i].sqrt() * noise
            else:
                x = mean                            # final step adds no noise
        else:
            # Recompute eps after clipping so x = sqrt(ab)*x0 + sqrt(1-ab)*eps.
            eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
            j = indices[k + 1] if k + 1 < len(indices) else -1
            ab_next = diffusion.ab[j] if j >= 0 else x.new_tensor(1.0)
            x = ab_next.sqrt() * x0 + (1 - ab_next).sqrt() * eps
    return x


def save_grid(x, path):
    """Portable grayscale PGM: no image library required. Open in an image viewer."""
    x = ((x.detach().cpu().clamp(-1, 1) + 1) * 127.5).round().to(torch.uint8)
    n, _, h, w = x.shape
    cols = min(4, n)
    rows = (n + cols - 1) // cols
    grid = torch.full((rows * h, cols * w), 255, dtype=torch.uint8)
    for k in range(n):
        r, c = divmod(k, cols)
        grid[r*h:(r+1)*h, c*w:(c+1)*w] = x[k, 0]
    Path(path).write_bytes(f"P5\n{cols*w} {rows*h}\n255\n".encode() + bytes(grid.flatten().tolist()))

# PART 4: validation, EMA, checkpointing, and the executable entry point
@torch.no_grad()
def validation(model, diffusion, batch, device):
    model.eval()
    rng = torch.Generator().manual_seed(9001)          # identical held-out draws each time
    return sum(loss_batch(model, diffusion, batch, rng, device).item() for _ in range(8)) / 8


@torch.no_grad()
def update_ema(ema, model, decay):
    for avg, current in zip(ema.parameters(), model.parameters()):
        avg.lerp_(current, 1 - decay)
    for avg, current in zip(ema.buffers(), model.buffers()):
        avg.copy_(current)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=3000, help="total steps, including before resume")
    parser.add_argument("--batch", type=int, default=64)
    parser.add_argument("--T", type=int, default=200)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--out", default="run")
    parser.add_argument("--resume", default=None)
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"])
    args = parser.parse_args()
    if args.steps < 1 or args.batch < 1:
        parser.error("steps and batch must be positive")
    device = args.device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    torch.manual_seed(7)
    torch.set_num_threads(4)
    rng = torch.Generator().manual_seed(42)
    model = TinyUNet().to(device)
    ema = copy.deepcopy(model).requires_grad_(False)
    diffusion = Diffusion(args.T).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.0)
    start = 0
    if args.resume:
        state = torch.load(args.resume, map_location="cpu", weights_only=True)
        if state["T"] != args.T or state["batch"] != args.batch:
            raise ValueError("Resume with the checkpoint's T and batch size")
        model.load_state_dict(state["model"])
        ema.load_state_dict(state["ema"])
        diffusion.load_state_dict(state["diffusion"])
        optimizer.load_state_dict(state["optimizer"])  # includes saved learning rate
        rng.set_state(state["rng"])
        start = state["step"]
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    save_grid(shapes(16, torch.Generator().manual_seed(2026)), out / "data.pgm")
    print(f"device={device}; parameters={sum(p.numel() for p in model.parameters()):,}; terminal alpha_bar={diffusion.ab[-1].item():.3g}")
    if start >= args.steps:
        print("Checkpoint already reached requested total steps; evaluating it.")
    for step in range(start + 1, args.steps + 1):
        model.train()
        optimizer.zero_grad(set_to_none=True)
        loss = loss_batch(model, diffusion, args.batch, rng, device)
        if not torch.isfinite(loss):
            raise RuntimeError("Non-finite loss: inspect the data, schedule, and activations")
        loss.backward()
        grad_norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if not torch.isfinite(grad_norm):
            raise RuntimeError("Non-finite gradient")
        optimizer.step()
        decay = min(0.999, (1 + step) / (10 + step))   # faster EMA adaptation early on
        update_ema(ema, model, decay)
        if step == 1 or step % 100 == 0:
            print(f"step={step} train_mse={loss.item():.5f} grad_norm={grad_norm.item():.3f}")
        if step % 500 == 0 or step == args.steps:
            torch.save(dict(step=step, T=args.T, batch=args.batch, model=model.state_dict(),
                            ema=ema.state_dict(), diffusion=diffusion.state_dict(),
                            optimizer=optimizer.state_dict(), rng=rng.get_state()), out / "checkpoint.pt")
    print(f"heldout_ema_mse={validation(ema, diffusion, args.batch, device):.5f}")
    for method in ("ddpm", "ddim"):
        samples = sample(ema, diffusion, n=16, method=method, steps=50)
        save_grid(samples, out / f"samples_{method}.pgm")
    print(f"Saved data, samples, and checkpoint in {out.resolve()}")


if __name__ == "__main__":
    main()

Run in three stages

# 1. Verify the program runs; these two steps do not train a useful generator.
python diffusion_lab.py --steps 2 --batch 4 --T 20 --out smoke

# 2. Train the default 200-level model. Use a smaller batch if memory is limited.
python diffusion_lab.py --steps 3000 --batch 64 --out run

# 3. Continue to 6000 total updates. Keep T and batch size consistent.
python diffusion_lab.py --resume run/checkpoint.pt --steps 6000 --batch 64 --out run

The script selects CUDA, then Apple MPS, then CPU when available; you can override with --device cpu, --device cuda, or --device mps. Runtime depends strongly on hardware. Every 500 updates and at the end, it writes a checkpoint. It evaluates the EMA model and generates sample grids after training.

PGM is a simple grayscale image format; open it in an image viewer such as Preview or convert it with your preferred image tool. The format keeps the Python lab dependent only on PyTorch. Resume restores the saved optimizer learning rate. Random draws for training are separate from validation and sampling, so changing when you inspect samples does not consume training randomness. Exact cross-device reproducibility is not guaranteed.

Why the extra training machinery matters

Exponential moving average (EMA) maintains a slowly changing copy of the weights. It often produces steadier samples than the latest optimization step. The lab starts with a faster-moving average and gradually increases its smoothing. Sampling uses the EMA copy in evaluation mode with gradients disabled.

Validation uses a separate, fixed stream of generated images, times, and noise. This makes comparisons less noisy. It measures held-out denoising error, not perceptual image quality. Both matter. For a real image dataset, make a proper split by source or subject to avoid near-duplicate leakage.

Checkpointing saves enough to continue optimization. Saving only the model weights is sufficient for inference but loses the optimizer moments, EMA, and training-stream position. In a larger system, also store the data sampler state, full experiment configuration, library versions, and mixed-precision scaler if used.

What the executed training run produced

A CPU run with 3,000 updates, batch size 32, T=200, and PyTorch 2.9.1 reached a fixed held-out EMA noise MSE of 0.00861. The grids below contain all 16 samples for seed 123; no samples were selected or removed. Both samplers use the same EMA checkpoint and initial Gaussian noise.

Target data

DDPM · 200 evaluations

DDIM · 50 evaluations

Read the images, not just the loss. The 200-step DDPM grid contains recognizable rings and squares. The 50-step DDIM grid is faster in model evaluations but retains visible background noise in several samples. This is a useful measured limitation: the faster sampler is not quality-equivalent in this run. Compare more DDIM steps or improved training before claiming a speedup at equal quality. These small grids are an implementation demonstration, not a statistical quality benchmark.

A sensible order of checks

1. Inspect the data. Visualize a clean batch. Confirm the intended range, channel order, and spatial dimensions. Normalization should happen exactly once.
2. Check the schedule. Verify 0<β<1, decreasing A, finite posterior coefficients, and a sufficiently small terminal A. Check the final-step identities from lesson 04.
3. Check noising and inversion. Construct xₜ using known ε. Reconstruct x₀ with that same ε, before clipping. At moderate noise levels the error should be close to floating-point precision.
4. Overfit a fixed corruption batch. Fix x₀, t, and ε, and optimize on that one batch. Error should decrease substantially. This tests whether the network and optimizer can fit the exact targets.
5. Overfit a tiny image set with fresh corruption. This is harder: error need not reach zero because the denoising task remains ambiguous. Inspect reconstructions at low, medium, and high noise.
6. Generate from pure noise. Use fixed initial seeds, first with the full DDPM sampler. Only introduce faster sampling after the baseline behaves sensibly.

Match symptoms to tests

import torch
from diffusion_lab import Diffusion, TinyUNet

d = Diffusion(T=200)
assert (d.ab[1:] < d.ab[:-1]).all()
assert d.ab[-1] < 1e-5
assert d.posterior_var[0] == 0
assert torch.allclose(d.c0[0], torch.tensor(1.0))
assert d.ct[0] == 0

x0 = torch.randn(3, 1, 32, 32)
eps = torch.randn_like(x0)
i = torch.tensor([20, 50, 100])
A = d.at(d.ab, i)
xt = d.q_sample(x0, i, eps)
reconstructed = (xt - (1 - A).sqrt() * eps) / A.sqrt()
assert torch.allclose(reconstructed, x0, atol=1e-6)

model = TinyUNet()
pred = model(xt, i)
assert pred.shape == x0.shape
(pred - eps).square().mean().backward()
assert all(p.grad is not None and torch.isfinite(p.grad).all()
           for p in model.parameters())
print("Schedule, endpoint, inversion, shapes, and gradients pass.")

Evaluate by noise level

A single average hides where a model struggles. Divide times into low-, middle-, and high-noise bands and report MSE separately. For reconstructions, convert predictions to x̂₀ and compare at fixed times. For generation, use the same seed set and sampler across checkpoints. Keep random sample grids, not just the prettiest images.

For real-image research, FID compares the mean and covariance of generated and real features from a chosen network. It depends on the feature extractor, preprocessing, sample count, and reference set. It is not a direct test of prompt following, diversity in every subgroup, or absence of memorization. Pair it with task-specific measurements and human inspection. For ControlNet, assess both image quality and agreement with the supplied structure.

Training with T = 200 does not force us to evaluate the denoiser 200 times for every image. The network has learned to make predictions at many noise levels. A sampler can choose a smaller decreasing sequence of those levels and use an update that connects them.

Do not simply skip entries in the adjacent-step DDPM loop. Its coefficients describe t→t−1 under the original schedule. A jump from t to s<t needs a matching transition.

Deterministic DDIM in two lines

At the current time t, predict ε̂ and reconstruct x̂₀. Then form a state at the next chosen time s using less noise:

Because s is less noisy, Aₛ is larger. The update preserves the current predicted image/noise decomposition while changing its proportions. At the next iteration the network makes a new prediction, so the whole trajectory does not rely on one initial estimate.

This is DDIM with η = 0: no fresh noise is injected during the trajectory. Different initial noise seeds still give different outputs. With a fixed initial noise, fixed conditions, and deterministic numerical execution, the path is deterministic.

What faster sampling costs

Fewer evaluations usually reduce latency, but large jumps ask more of an imperfect denoiser and increase discretization error. Compare 10, 25, 50, and all training levels with the same seeds. There is no universal best count. More steps are not guaranteed to improve every model under every solver and guidance setting.

DDPM, DDIM, and ODE solvers are not synonyms

DDPM is stochastic ancestral sampling. Deterministic DDIM has a connection to a continuous ordinary differential equation (ODE). Other methods, such as DPM-Solver, exploit the structure of diffusion ODEs to take efficient numerical steps. Some solvers evaluate the denoiser multiple times per step, so count model evaluations as well as steps. Lesson 12 separates changing a solver from training a model to generate in very few steps.

Read next: DDIM and DPM-Solver.

An unconditional model learns a distribution over all training images. A conditional model also sees a label, text prompt, or another signal and learns the distribution associated with that condition. The corruption equation and noise target need not change: the prediction becomes εθ(xₜ,t,c).

How the condition reaches the model

For a class label, add a learned embedding to the time embedding. For text, a text encoder produces token embeddings. In a U-Net with cross-attention, image features supply the queries and text embeddings supply keys and values. You already know attention; the new point is that spatial image locations can retrieve different pieces of the prompt.

Training needs image/condition pairs. Giving an unconditional model a prompt at inference does not create text understanding. The architecture needs an input path for the condition, and training must teach it how to use that path.

Classifier-free guidance (CFG)

During training, randomly remove the condition for some examples, substituting a learned null label or an encoded empty prompt. The same network then learns both conditional and condition-absent predictions. During sampling, combine them:

In this convention, w = 0 is unconditional and w = 1 is the ordinary conditional prediction. Values above 1 extrapolate in the conditional direction. Stronger guidance often improves adherence at a cost to diversity and can introduce saturation or artifacts. It is a sampling control, not a guarantee that larger is better.

Different papers write the same expression with a shifted scale: (1+s)ε̂_cond − s ε̂_uncond. Their s is w−1 here. State your convention before discussing a number.

A small implementation sketch

# Training: label K is a dedicated null entry, not a real class.
drop = torch.rand(B, device=device) < 0.1
used_label = torch.where(drop, K, label)
noise_hat = model(xt, i, used_label)
loss = F.mse_loss(noise_hat, noise)

# Sampling: same xt and time for both branches.
xt2 = torch.cat([xt, xt], dim=0)
t2 = torch.cat([i, i], dim=0)
y2 = torch.cat([torch.full_like(label, K), label], dim=0)
uncond, cond = model(xt2, t2, y2).chunk(2)
noise_hat = uncond + w * (cond - uncond)

This is an extension sketch, not a call signature supported by the unmodified lab. Add the class embedding from exercise 03 first. A 10% condition-drop probability is a reasonable starting experiment for the shape task; tune it rather than treating it as universal.

What about classifier guidance?

Classifier guidance uses the gradient of a separate noisy-image classifier, ∇_xₜ log p(c|xₜ,t), to steer the diffusion score. That classifier must work at the relevant noise levels; a clean-image classifier is not automatically suitable. CFG obtains a related steering signal from the two denoiser predictions and needs no separate classifier.

Source: Classifier-Free Diffusion Guidance.

Why denoise a compressed image?

Running a large denoiser on every pixel of a high-resolution image is expensive. A latent diffusion model first encodes the image into a smaller spatial tensor z₀. Diffusion happens there; a decoder converts the final clean latent back into pixels.

For example, reducing 512×512 pixels to a 64×64 latent grid reduces the number of spatial positions by 64×. It does not promise a 64× end-to-end speedup: channel widths, attention, the autoencoder, and other components still cost compute.

The encoder and decoder are commonly pretrained and frozen while the denoiser is trained. Autoencoder reconstruction puts a ceiling on fidelity: details it discards cannot be recovered merely by training a better denoiser on its latents.

Latent scaling is not cosmetic

Use the exact encoder preprocessing, latent scale or shift, and decoder convention associated with the chosen autoencoder and checkpoint. If you change the latent magnitude, you change its effective signal-to-noise ratio under the same noise schedule. Do not copy a scaling constant from an unrelated model family.

Text-to-image training adds encoded captions as a condition. If the text encoder is frozen, its outputs can sometimes be cached, but caching reduces flexibility for caption changes and augmentations. Likewise, caching latents locks in the image preprocessing and any chosen crop. Keep those trade-offs explicit.

What changes with a diffusion transformer?

A diffusion transformer, or DiT, uses a transformer as the denoiser. Patchify the noisy latent, process tokens with time and condition information, and project back to an image-shaped prediction. Unlike a classification ViT, it needs an output for every patch position, not one class vector.

The original DiT design uses adaptive LayerNorm conditioning and initially zeroed modulation/output pathways. Other diffusion transformers use cross-attention or joint image/text attention. “Diffusion” describes the generative formulation; “transformer” describes the network that predicts the required quantity. They are different choices.

Image editing is a change in the starting problem

Image-to-image: encode or normalize a source image, corrupt it to a chosen starting noise level, then denoise under the desired condition. More initial noise generally permits larger changes and preserves less of the source. “Strength” is a UI mapping onto this choice and depends on the pipeline.

Inpainting: supply a mask and preserve known regions while generating unknown regions. A simple sampler can reinsert appropriately noised known pixels at every step, although boundaries can be inconsistent. Dedicated inpainting models are trained to use masks and visible image content. A mask convention—whether 1 means keep or regenerate—must be stated explicitly.

Primary reading: Latent Diffusion Models and DiT.

A prompt can ask for “a person raising their left arm,” but it does not specify exact joint coordinates. ControlNet adds an aligned image-like condition—such as a pose skeleton, edge map, depth map, or segmentation map—to a pretrained image generator. The model can vary appearance while following that structure.

The architecture, without the mystery

In the original U-Net-based ControlNet, the pretrained denoiser remains frozen. A trainable copy of its encoder and middle block receives the noisy latent, time, text, and an encoded spatial condition. Its multiscale features pass through zero-initialized 1×1 convolutions. The resulting residuals are added into the frozen U-Net's skip and middle features, influencing the decoder.

At initialization those residuals are zero, so attaching the branch leaves the base model's output unchanged. During training the branch learns useful corrections. A condition encoder reduces the control image to an appropriate feature resolution. This is more than concatenating an edge map to an arbitrary frozen model's input.

Architecture source: Adding Conditional Control to Text-to-Image Diffusion Models. The diagram is a simplified explanation of its U-Net design; control mechanisms in other backbones may differ.

Why zero convolutions can learn

Consider one output bridge y = Wh+b with W = 0 and b = 0. Its output is zero, but the gradient with respect to W contains the incoming loss gradient multiplied by h. If h is nonzero and the loss is sensitive to this output, W can receive a nonzero update on the first step.

Initially the gradient through this bridge into h is Wᵀ times the incoming gradient, so it is zero. After the bridge weights move away from zero, gradients can reach the earlier branch. Zero-initializing a bridge is different from zero-initializing every weight in a network. The copied features give the bridge something useful to learn from.

import torch
from torch import nn
h = torch.ones(2, 4, 8, 8, requires_grad=True)
bridge = nn.Conv2d(4, 4, 1)
nn.init.zeros_(bridge.weight)
nn.init.zeros_(bridge.bias)
loss = (bridge(h) - 1).square().mean()
loss.backward()
assert bridge.weight.grad.abs().sum() > 0
assert bridge.bias.grad.abs().sum() > 0
assert h.grad.abs().sum() == 0

A training example has three aligned parts

Prepare a target image, a spatial condition, and an optional caption. Derive an edge/depth/pose condition from the training image when appropriate, but make sure inference conditions use compatible conventions. Jointly transform image and control when cropping or flipping. Independently cropping a pose map and its target teaches contradictory supervision.

1. Freeze the pretrained autoencoder, text encoder, and base denoiser weights. Initialize the control branch from the matching base encoder and middle block; initialize its output bridges to zero.
2. Encode the target image into correctly scaled z₀. Encode the caption. Keep the control map in the format and range expected by its condition encoder; it is not simply another VAE latent unless the architecture says so.
3. Choose t and ε, then create zₜ using the base model's schedule.
4. Run the control branch to produce multiscale residuals. Pass those residuals into the frozen base U-Net.
5. Compute the base model's expected target loss—noise MSE for an ε-predicting base, or the matching alternative. Update only the control branch and bridges.

# Architecture-level sketch: these interfaces depend on the chosen backbone.
base.requires_grad_(False)
control.requires_grad_(True)
with torch.no_grad():
    z0 = encode_and_scale(images)
    text = text_encoder(captions)
zt = add_noise(z0, noise, t)
residuals = control(zt, t, text, control_maps)
prediction = base(zt, t, text, control_residuals=residuals)
loss = F.mse_loss(prediction.float(), noise.float())
loss.backward()  # gradients pass through the frozen decoder into residuals
control_optimizer.step()

Control strength is different from text guidance

A control scale multiplies the branch residuals. CFG combines conditional and condition-absent denoiser outputs. They act at different places. Holding the initial noise fixed, sweep control scale while leaving CFG unchanged, then do the reverse. Strong control can preserve geometry at the expense of texture or flexibility; weak control can let the model ignore the supplied structure.

You can also enable control for only part of the denoising trajectory. Early high-noise steps often have a large influence on layout, but the best schedule depends on the task and control type. With multiple ControlNets, sum compatible residuals with separate scales and inspect conflicts rather than assuming every condition can be satisfied simultaneously.

Be explicit about CFG routing: one implementation may feed spatial control into both the prompt-present and prompt-absent branches; another may apply it only to the prompt-present branch. In the first case, the “unconditional” branch is text-unconditional but still spatially conditioned. The same numeric scales can behave differently across these choices.

ControlNet versus LoRA

LoRA parameterizes a weight update using low-rank matrices, reducing the number of trainable parameters. ControlNet adds a pathway for a spatial input. They solve different design problems and can be combined. Fine-tuning appearance with a LoRA does not, by itself, teach a model to interpret an arbitrary pose or depth map.

What to test before a large run

• Initialization equivalence: in evaluation mode, the newly attached zero-bridge branch produces the same output as the base model for identical inputs.
• Gradient flow: base parameter gradients stay absent; bridge gradients are present; earlier branch gradients emerge as the bridges learn.
• Alignment: overlay control maps on target images after all preprocessing. Check orientation, cropping, resizing, depth convention, and skeleton format.
• A tiny controlled overfit: on a small paired set, inspect whether generated structures move when the control changes at a fixed seed. Also test held-out layouts.
• Quality and adherence: measure both. Edge similarity, pose keypoint distance, or depth agreement can assess structure, but each depends on the estimator used and is not a complete quality metric.

For a real pretrained model, use the authors' implementation or inspect the matching Diffusers ControlNet training example. Architectural interfaces, preprocessing, and memory needs must match your selected checkpoint.

You can build the lab without this lesson. Read it when you want to explain why noise regression is connected to a probabilistic model, or when an interviewer asks for a derivation. We keep the main argument short and put the algebra behind expandable panels.

Why does a noise loss train a generator?

The generative model starts with a simple prior p(x_T) and multiplies learned reverse transitions pθ(xₜ₋₁|xₜ). To assign a probability to x₀, it must account for all possible intermediate states. That sum or integral is impractical. We instead use the known forward corruption process q as a way to construct a tractable lower bound on log pθ(x₀), called the evidence lower bound, or ELBO.

Minimizing the negative bound separates into a terminal prior term, a reconstruction term at the data boundary, and terms that make learned reverse transitions match known forward posteriors. With fixed Gaussian reverse variances, those intermediate terms reduce to weighted squared errors between means. Expressing the means using a predicted ε turns them into weighted noise-prediction errors.

The commonly used simple noise MSE drops those particular ELBO weights. It is motivated by the bound but is not, in general, numerically equal to the exact negative ELBO. A training MSE value is not an image log-likelihood.

The score: a direction, not a quality rating

The score of a noisy image distribution is s(xₜ,t) = ∇_xₜ log pₜ(xₜ). It points toward increasing density locally. It is not a scalar confidence score and not a gradient with respect to network weights.

For the conditional corruption Gaussian, ∇_xₜ log q(xₜ|x₀) = −ε/√(1−Aₜ). Average this over possible x₀ given xₜ and use the optimal MSE predictor:

This is the bridge between denoising and score estimation. The clean-data distribution can be complicated or concentrated on a thin subset of image space; adding Gaussian noise smooths it. Learning directions at several noise levels is more useful for generation than trying to jump directly to one final clean-image average.

Foundations: DDPM and Score-Based Generative Modeling through SDEs.

Flow matching with a straight path

Here we deliberately introduce a new time symbol u to make its direction explicit: u = 0 is noise, u = 1 is data. Draw independent z∼N(0,I), data x₀, and u∼Uniform(0,1). Interpolate and predict the interpolation velocity:

At generation time start with z and integrate dx/du = vθ(x,u) from 0 to 1. One Euler update is x ← x + Δu·vθ(x,u). The network learns an average velocity for the possible pairs passing through each x,u; that is why independent straight training pairs need not yield straight model trajectories.

Do not confuse this target x₀−z with diffusion v = aε−σx₀ from lesson 02. Some implementations reverse time or use another path, changing signs and formulas. Before borrowing code, write the two endpoints and differentiate the path.

How very few-step models become possible

Better solvers use the same learned field more accurately per evaluation. Distillation trains a student to reproduce a teacher's larger movement in fewer steps. For example, progressive distillation can train one student step to match two teacher steps, then repeat. Consistency models learn mappings that agree on an endpoint for states along a shared trajectory; they can be trained by distillation or through suitable direct training objectives.

These choices change compute and error in different places. A distilled checkpoint may expect particular noise levels and guidance settings. A generic solver may not deliver the advertised result if those assumptions are ignored. Count teacher-generation cost, student-training cost, and inference cost separately.

Where advanced training refinements fit

Noise-level sampling and loss weighting decide which tasks receive more optimization effort. Preconditioning rescales model inputs, outputs, and targets to make those tasks numerically better behaved. Learned variance changes the reverse distribution. None should be introduced as a magic fix before the baseline noising and sampler tests pass.

For example, with SNR = A/(1−A), one Min-SNR formulation weights ε MSE by min(SNR,γ)/SNR. The corresponding v-prediction weighting is different. Use the formula for your parameterization, and handle endpoint limits rather than dividing by zero. This is optional refinement, not part of the lab's first-run recipe.

Primary reading: Flow Matching, Progressive Distillation, Consistency Models, EDM, and Min-SNR weighting.

These are practical starting points, not universal optimal settings. Keep a record of data preprocessing, prediction target, schedule, optimizer, EMA, sampler, and random seeds. Change one major component at a time, and compare against a saved baseline.

Recipe A · Learn the mechanics on the shape lab

Setup: 32×32 grayscale, 823,841-parameter U-Net, cosine schedule with T=200, ε prediction, uniform times, AdamW at 2×10⁻⁴ with zero weight decay, batch 32–64, float32, gradient clipping at norm 1, and EMA. Begin with a two-step smoke run, then try 3,000 updates. CPU is viable for experimentation; accelerator training is preferable for iteration speed.

Success: finite gradients; strong reduction on a fixed-corruption overfit test; improving held-out noise error; recognizable and varied uncurated samples. Compare the full DDPM sampler with 25–50 DDIM steps only after inspecting the baseline. A step count alone is not evidence of convergence.

Recipe B · Replace synthetic shapes with MNIST or small images

Setup: keep ε prediction and the model first. For MNIST, pad 28×28 images to 32×32, turn pixels into [0,1], then map to [−1,1]. Use a shuffled DataLoader and a fixed held-out split. Start with batch 64, learning rate 2×10⁻⁴, and T=200 or 1,000 with a freshly constructed cosine schedule. Evaluate at regular update counts and stop based on validation and sample behavior rather than a copied epoch number.

# Optional dataset replacement: requires a compatible torchvision installation.
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
transform = transforms.Compose([
    transforms.Pad(2),           # 28x28 -> 32x32
    transforms.ToTensor(),      # [0,255] -> [0,1]
    transforms.Normalize((0.5,), (0.5,)),
])
train_data = datasets.MNIST("data", train=True, download=True, transform=transform)
loader = DataLoader(train_data, batch_size=64, shuffle=True, num_workers=0)
# Replace shapes(...) in the training path with x0 from the loader.
# Keep the held-out validation loader separate; labels are unused unless conditioning.

Next extension: for 32×32 RGB images, change input and output channels to 3 and inspect whether the small model has enough capacity. Consider a wider U-Net and low-resolution attention. Expect a much harder task than synthetic shapes. Avoid geometrically or semantically invalid augmentation, such as arbitrary flips for digits.

Recipe C · Fine-tune a pretrained latent model

Setup: choose one documented checkpoint family and retain its autoencoder preprocessing, latent scaling, text encoding, prediction target, and scheduler contract. Freeze the autoencoder and usually the text encoder. Use image/caption pairs and a held-out prompt set. For a modest adaptation, a denoiser learning-rate sweep around 10⁻⁶–10⁻⁵ is a conservative initial experiment; adapter-only training often needs a different range.

Memory: mixed precision, gradient accumulation, and activation checkpointing are practical tools. Keep numerically sensitive losses and schedule calculations in float32. CUDA float16 commonly needs loss scaling; bfloat16 has a wider exponent range where supported. Effective batch size is microbatch×accumulation×number of workers participating in distributed data-parallel training. Normalize accumulated losses consistently.

Success: improvement on held-out target prompts with reasonable diversity, and no unacceptable loss of baseline capabilities. Evaluate checkpoints on fixed seeds and both target-domain and general prompts. Overfitting can improve training examples while damaging broader behavior.

Recipe D · Train spatial control

Setup: use a ControlNet architecture matched to the base checkpoint; freeze base components; initialize the copied blocks and zero output bridges correctly. Inspect paired image/control overlays before training. Start with a small paired overfit experiment, then expand. A learning-rate sweep around 10⁻⁵ for a copied control branch is a starting experiment, not a hardware-independent prescription.

Conditioning: decide whether and how often to drop text based on whether inference must work with weak or absent prompts. Document whether control enters both CFG branches. Keep control-map processing consistent: edge thresholds, pose drawing, depth direction and normalization, and segmentation label conventions can all change the task.

Success: held-out structure is followed when control strength is nonzero, scale zero recovers the base path, gradients reach the intended trainable parameters, and visual quality remains acceptable. Report separate structure and appearance assessments. Test conflicting prompts and controls to understand the model's limits.

What a useful experiment record contains

Write down a hypothesis, the single main change, compute used, held-out metric, fixed-seed sample grid, failure cases, and whether you keep the change. “It looked better” without fixed seeds or a comparable sampler is not a reliable experiment.

Answer aloud before revealing each response. A strong answer names its assumptions, tracks shapes, explains the mechanism, and proposes a way to test it. A useful self-score is 0 = cannot explain, 1 = definition only, 2 = mechanism correct, 3 = mechanism plus a concrete failure mode or test.

A 45-minute coding mock

1. 0–8 min: specify tensor shapes and an indexing convention; construct a valid schedule.
2. 8–18 min: implement noising, a time-aware denoiser interface, and one optimizer step.
3. 18–30 min: implement fixed-posterior DDPM sampling and the clean endpoint.
4. 30–38 min: write a known-noise inversion check and one finite-gradient check.
5. 38–45 min: explain how to add DDIM and CFG, and diagnose a training/sampling mismatch.

Rubric: correct mechanics 40%, shapes and boundary conditions 25%, explanation 20%, useful tests and diagnosis 15%. You need not type a large production U-Net in 45 minutes; use a small time-aware model and state the architectural contract. Do not claim an implementation works because it merely produces an image-shaped tensor.

The capstone: build, compare, explain

1. Rebuild the unconditional lab. Reproduce noising, training, EMA, checkpointing, and DDPM sampling. Pass the mathematical boundary and inversion checks.
2. Compare samplers. Evaluate full DDPM and DDIM at 10, 25, and 50 steps on the same initial seeds. Record image quality, diversity, and model evaluations.
3. Add class conditioning and CFG. Produce class-by-seed-by-guidance grids. Verify the w=0 and w=1 limits.
4. Add spatial control. Use the small ControlNet adaptation from lesson 10, or a matching pretrained setup. Prove initialization equivalence and gradient flow; evaluate held-out conditions.
5. Write a one-page report. Describe the model contract, two failures you diagnosed, one controlled ablation, and one remaining limitation.

Completion standard: another person should be able to run your code, understand the conventions, resume training, reproduce a fixed-seed comparison on the same setup, and see why you trust the result. If only the training loss is convincing, keep investigating.

The five equations to reproduce from memory

Also be able to reconstruct the DDPM mean and variance with the reference in lesson 04, and write the deterministic DDIM jump in lesson 07. Understanding the boundary cases matters more than memorizing every coefficient.

A small glossary

Primary sources, in a useful reading order

These links are further reading, not prerequisites for following the course. The historical formulations are identified deliberately; this is not a leaderboard of current models.

1. Ho et al. · Denoising Diffusion Probabilistic Models — connect the original objective and sampling algorithm to the lab.
2. Nichol & Dhariwal · Improved DDPM — schedule design and learned reverse variance.
3. Song et al. · DDIM — use fewer sampling steps with a trained diffusion model.
4. Ho & Salimans · Classifier-Free Diffusion Guidance — the conditional/unconditional combination.
5. Rombach et al. · Latent Diffusion Models — compression and conditional image synthesis.
6. Zhang et al. · ControlNet — learn spatial control with a frozen base model.
7. Peebles & Xie · DiT — a transformer denoising backbone.
8. Song et al. · Score-Based Generative Modeling through SDEs — the continuous-time connection.
9. Lipman et al. · Flow Matching — learn velocity fields along probability paths.
10. Lu et al. · DPM-Solver — fast numerical sampling.
11. Salimans & Ho · Progressive Distillation and Song et al. · Consistency Models — learn to use fewer generation steps.
12. Karras et al. · EDM and Hang et al. · Min-SNR — numerical design and noise-level weighting.