Breadth: deep learning, from initialization to diffusion and flow matching

Before this lesson: lesson 07 (AdamW and mixed precision in code) and lesson 08 (attention). Depth: the Deep learning track, especially initialization, normalization and residuals and training dynamics; the Optimization lessons on Adam and AdamW and schedules and clipping; the Generative models track for diffusion and flow matching. Time: about 95 minutes, plus the notebook.

---

Why this matters #

Deep-learning breadth questions come in phone screens, in "ML knowledge" rounds, and as follow-ups in ML-coding rounds ("why did you put the LayerNorm there?"). Research scientist loops often add a deep dive on generative models. Three failures are common:

• a popular but weak explanation: "BatchNorm works because it removes internal covariate shift" (this explanation is disputed);
• a wrong detail: "AdamW is Adam with L2 regularization";
• no number: "diffusion is slow" instead of "about 10–50 network calls with an ODE solver, 1–4 after distillation".

Use the four-part answer of lesson 09: one sentence that defines the thing, the formula or mechanism, one trade-off, one number or example.

1. Predict first #

2. The picture #

2.1 One number per layer #

Start with the simplest deep network: one number per layer, $h_l = a_l h_{l-1}$. The output is $h_L = (a_L \cdots a_1)\,h_0$, and by the chain rule the gradient with respect to $h_0$ is the same product. With $a_l = 0.5$ and $L = 20$ the product is about $10^{-6}$: the gradient vanishes. With $a_l = 2$ it is about $10^{6}$: the gradient explodes. Most tools of this lesson keep the per-layer gains near 1, or let training survive when they are not:

Tool	What it does to the gain
initialization (Xavier, He)	sets the gain to 1 at step 0
normalization (BatchNorm, LayerNorm, RMSNorm)	resets the scale of the activations at every layer and every step
residual connection $h_l = h_{l-1} + F(h_{l-1})$	adds an identity path, so the gain cannot fall below the identity
gradient clipping	caps the size of one step when a gain blows up
warmup	keeps the steps small while the curvature and Adam's statistics settle
Adam	gives each parameter its own step scale, about $1/\sqrt{\text{mean of } g^2}$

Interactive figure: loading…

2.2 Initialization: keep the variance constant #

A layer computes $z = Wh$. If the weights are random with variance $\sigma_w^2$, each output adds up $n$ random terms, so its variance is about $n\,\sigma_w^2$ times the input's second moment. A ReLU then sets half of the values to zero, which halves the second moment. To keep the signal at the same size through many layers, choose $\sigma_w^2 = 2/n$ for ReLU (He initialization); section 3.3 proves it.

Interactive figure: loading…

2.3 An architecture is an inductive bias #

An architecture fixes which inputs interact and which symmetry is built in; this is its inductive bias. A CNN mixes only nearby pixels and shares its weights across positions. An RNN reads one step at a time through a state of fixed size. A transformer lets every token interact with every other token, at a cost that grows like $T^2$ for $T$ tokens, and it knows the order only through positional information. Less built-in bias means more data is needed.

2.4 Three ways to turn noise into data #

• Autoregressive: generate one piece at a time, each conditioned on the pieces before it: $n$ sequential network calls for $n$ tokens.
• Diffusion: corrupt data with Gaussian noise in small steps, learn to predict the noise, and remove it step by step: about 10–1000 calls, each on the whole sample.
• Flow matching: join each noise sample to a data sample by a straight line, learn the average velocity at each point, and integrate an ODE: about 10–50 calls.

With Gaussian noise, diffusion and flow matching are one model family (section 3.7).

Interactive figure: loading…

3. Definitions and main results #

3.1 Optimizers: SGD, momentum, Adam, AdamW #

Notation: $\theta$ are the parameters, $g_t$ the minibatch gradient at step $t$, $\eta$ the learning rate. All operations on vectors are elementwise.

• SGD: $\theta_t = \theta_{t-1} - \eta g_t$.
• Momentum: $m_t = \mu m_{t-1} + g_t$ and $\theta_t = \theta_{t-1} - \eta m_t$, with $\mu \approx 0.9$. It averages the gradient over about $1/(1-\mu)$ steps: it damps the zig-zag across a narrow valley and speeds up the travel along it.
• Adam: $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ and $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$; bias corrections $\hat m_t = m_t/(1-\beta_1^t)$ and $\hat v_t = v_t/(1-\beta_2^t)$; update $\theta_t = \theta_{t-1} - \eta\,\hat m_t/(\sqrt{\hat v_t} + \epsilon)$, where the small constant $\epsilon$ prevents a division by 0. Defaults: $\beta_1 = 0.9$, $\beta_2 = 0.999$ (0.95 in LLM training), $\epsilon = 10^{-8}$. The step of each coordinate is about $\eta$ times a signal-to-noise ratio, so it does not depend on the scale of the gradient. The bias corrections undo the zero start of $m$ and $v$ (lesson 07 proves it).
• AdamW: the Adam step plus a separate decay, $\theta_t \leftarrow \theta_t - \eta\lambda\,\theta_{t-1}$ (PyTorch's form). The weight decay $\lambda$ is typically 0.1 in LLM training. It is applied to weight matrices; biases and normalization gains are usually excluded.

Proposition (L2 = weight decay for SGD, but not for Adam; Loshchilov and Hutter).

1. L2 regularization adds $\tfrac{\lambda'}{2}\|\theta\|^2$ to the loss, so the gradient becomes $g_t + \lambda'\theta_{t-1}$.
2. SGD: $\theta_t = \theta_{t-1} - \eta(g_t + \lambda'\theta_{t-1}) = (1 - \eta\lambda')\,\theta_{t-1} - \eta g_t$. This is weight decay with rate $\eta\lambda'$ per step: the same method.
3. Adam: the term $\lambda'\theta$ enters $m$ and $v$ and is then divided by $\sqrt{\hat v}$. A weight with large gradients has a large $\sqrt{\hat v}$, so its decay is small; a weight with small gradients is decayed strongly. No single $\lambda'$ gives every weight the same decay.
4. AdamW applies the decay outside the adaptive step, so every weight shrinks by the same factor $1 - \eta\lambda$ per step. $\blacksquare$

Worked example 4A puts numbers on step 3.

Which optimizer. SGD with momentum is still strong for CNNs on images. Transformers train much better with Adam-type methods, and AdamW is the default. Matrix optimizers such as Muon (orthogonalized momentum for weight matrices) are a newer option for large models.

3.2 Schedules, clipping, large batches, mixed precision #

• Warmup: raise $\eta$ linearly from about 0 to its peak over the first steps (often a few hundred to a few thousand). At initialization the curvature is high and Adam's $\hat v$ rests on few samples; post-LN transformers diverge without warmup (section 3.4).
• Decay: a cosine schedule decays smoothly to about 10% of the peak at the planned last step. Warmup–stable–decay (WSD) keeps the peak and decays only in the last 10–20% of the steps, so a run can be extended from a checkpoint of the stable phase.
• Global-norm clipping: stack all gradients in one vector $g$ and use $g \cdot \min(1, c/\|g\|_2)$, often with $c = 1.0$. It keeps the direction and caps the step. Log $\|g\|$ before clipping: a rising norm warns of a loss spike.
• Large batches: for SGD, multiply $\eta$ by $k$ when the batch grows by $k$ (the linear scaling rule, Goyal et al., 2017), with warmup. The rule fails above a critical batch size. For Adam a common heuristic is $\eta \propto \sqrt{B}$ for batch size $B$.
• Mixed precision: compute the matrix products in 16 bits; keep fp32 master weights and accumulate sums in fp32. bf16 has the range of fp32 and needs no loss scaling; fp16 has a small range and needs loss scaling (multiply the loss by $S$, divide the gradients by $S$). See lesson 07 for the details.

Interactive figure: loading…

3.3 Initialization from variance propagation #

Theorem. Let $z = Wh$ with $h \in \mathbb R^n$ ($n$ is the fan-in). The entries $W_{ij}$ are independent with mean 0 and variance $\sigma_w^2$, independent of $h$, and every $h_j$ has the same second moment $\mathbb E[h_j^2]$. Then

If $z_i$ is symmetric about 0, then $\mathbb E[\mathrm{ReLU}(z_i)^2] = \tfrac12 \operatorname{Var}(z_i)$.

Proof.

1. $z_i = \sum_j W_{ij} h_j$. Each term has mean $\mathbb E[W_{ij}]\,\mathbb E[h_j] = 0$, so $\mathbb E[z_i] = 0$.
2. For $j \ne k$: $\mathbb E[W_{ij} W_{ik} h_j h_k] = \mathbb E[W_{ij}]\,\mathbb E[W_{ik}]\,\mathbb E[h_j h_k] = 0$. The cross terms vanish.
3. So $\operatorname{Var}(z_i) = \mathbb E[z_i^2] = \sum_j \mathbb E[W_{ij}^2]\,\mathbb E[h_j^2] = n\sigma_w^2\,\mathbb E[h_j^2]$.
4. ReLU keeps $z_i$ when $z_i > 0$ and outputs 0 otherwise. By symmetry, the positive half carries half of $\mathbb E[z_i^2]$: $\mathbb E[\mathrm{ReLU}(z_i)^2] = \tfrac12 \operatorname{Var}(z_i)$. $\blacksquare$

Consequences.

• ReLU: the second moment stays constant from layer to layer when $n\sigma_w^2 \cdot \tfrac12 = 1$, that is $\sigma_w^2 = 2/n_{\text{in}}$ (He or Kaiming initialization).
• tanh or linear units (gain about 1 near 0): $\sigma_w^2 = 1/n_{\text{in}}$ (LeCun). The backward pass gives the same condition with $n_{\text{out}}$; Xavier (Glorot) initialization takes the compromise $\sigma_w^2 = 2/(n_{\text{in}} + n_{\text{out}})$.
• All weights equal to 0 (or to one constant) fails: every unit of a layer gets the same gradient, and the units stay identical. Random initialization breaks this symmetry. Biases start at 0.
• Transformers: GPT-2 uses a fixed standard deviation of 0.02 and scales the output projection of each residual branch by $1/\sqrt{N}$, with $N$ the number of residual layers, so the residual stream does not grow with depth.

3.4 Normalization and residual connections #

For one group of numbers $x_1, \dots, x_n$ with mean $\mu$ and variance $\sigma^2$: $\hat x_i = (x_i - \mu)/\sqrt{\sigma^2 + \epsilon}$ and $y_i = \gamma \hat x_i + \beta$, with a small constant $\epsilon$ (often $10^{-5}$), a learned gain $\gamma$ and a learned shift $\beta$. The norms differ in the group of numbers they use:

Norm	Normalizes over	Train vs eval	Used in
BatchNorm	one channel, over the batch (and space)	batch statistics in training, running averages in eval	CNNs
LayerNorm	the features of one token or sample	the same function	original transformer, BERT, GPT-2
RMSNorm	the features of one token, with no mean and no $\beta$: $y = \gamma \odot x/\sqrt{\tfrac1D \sum_i x_i^2 + \epsilon}$	the same function	T5, Llama and most recent LLMs
GroupNorm	groups of channels of one sample	the same function	small-batch vision, diffusion U-Nets

Here $D$ is the number of features. Why normalization helps. It keeps the activations in a good range at every step, not only at initialization, and it makes the output invariant to the scale of the incoming weights. Santurkar et al. (2018) showed that BatchNorm's benefit is not explained by less "internal covariate shift"; it makes the loss and its gradients smoother, which allows larger learning rates.

Problems of BatchNorm. The output for one sample depends on the rest of the batch. The statistics are noisy with a small batch per GPU; training and eval compute different functions (forgetting model.eval() is a classic bug); autoregressive decoding has no batch; and statistics over the time axis leak future tokens into a causal model. So transformers use LayerNorm or RMSNorm.

Residual connections. $x_{l+1} = x_l + F_l(x_l)$. With $J_k = \partial F_k/\partial x_k$, the Jacobian is

Expanding the product gives the identity $I$ plus other terms, so the gradient reaches every layer directly, however small the $J_k$ are. A plain stack has the product $J_{L-1} \cdots J_l$, which vanishes or explodes. This made networks with more than 100 layers trainable (ResNet).

Pre-LN vs post-LN. Post-LN (the original transformer) is $x_{l+1} = \mathrm{LN}(x_l + F(x_l))$: the norm sits on the residual path. Pre-LN (GPT-2 and later) is $x_{l+1} = x_l + F(\mathrm{LN}(x_l))$, plus one final norm before the output head: the residual path is a clean identity. Xiong et al. (2020) showed that post-LN has large gradients near the output at initialization, so it needs warmup; pre-LN trains stably at depth, even without warmup. The price of pre-LN: the norm of the residual stream grows with depth, so each later block changes it relatively less.

3.5 Regularization in deep networks #

• Dropout: in training, set each activation to 0 with probability $p$ and scale the others by $1/(1-p)$ (inverted dropout), so the expected value does not change; at eval, do nothing. It acts like an average over many sub-networks. Single-epoch LLM pretraining often uses $p = 0$; fine-tuning on small data often uses about 0.1.
• Weight decay (section 3.1).
• Data augmentation encodes symmetries you know (flips, crops, colour changes). It is the strongest regularizer in vision.
• Label smoothing trains on the target $(1-s)\,\text{onehot} + s/K$ for $K$ classes and smoothing $s$ (often 0.1). It stops the logits from growing without bound.
• Early stopping on a validation metric.

See the regularization lesson of the Deep learning track.

3.6 Architectures and their inductive biases #

Architecture	Built-in assumption	Cost per layer ($T$ positions, width $d$)	Weak spot
CNN	locality, translation equivariance (shared kernels)	$O(Tkd^2)$ for kernel size $k$	long-range context needs depth or pooling
RNN / LSTM	sequential order, a state of fixed size	$O(Td^2)$, sequential in $T$	no parallel training over time; gradient paths of length $T$
Transformer	every pair interacts; order only through positional information	$O(T^2 d + T d^2)$, parallel	cost quadratic in $T$; weak bias, needs much data
ViT	a transformer on $P \times P$ patches, $T = HW/P^2$ for an $H \times W$ image	as the transformer	weaker bias than a CNN: needs large-scale pretraining or strong augmentation
GNN	permutation equivariance on a graph	$O(\lvert E\rvert d + \lvert V\rvert d^2)$ for edges $E$, nodes $V$	over-smoothing with depth

Formulas that interviewers ask:

• Conv output size: $\lfloor (n + 2p - k)/s \rfloor + 1$ for input size $n$, padding $p$, kernel $k$, stride $s$. For $n = 32$, $k = 3$, $s = 2$, $p = 1$: $\lfloor 31/2 \rfloor + 1 = 16$.
• Receptive field of $L$ stacked $k \times k$ convolutions with stride 1: $1 + L(k-1)$. Two $3 \times 3$ layers see $5 \times 5$ with 18 weights per channel pair instead of 25, and add one more nonlinearity (the VGG argument).
• Attention: $\mathrm{softmax}(QK^\top/\sqrt{d_k})\,V$. If the entries of a query $q$ and a key $k$ are independent with mean 0 and variance 1, then $q \cdot k$ has variance $d_k$. Dividing by $\sqrt{d_k}$ brings it back to 1, so the softmax does not saturate at initialization.
• LSTM: the cell state is updated by addition, $c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t$ (forget gate $f_t$, input gate $i_t$, candidate $\tilde c_t$). With $f_t$ near 1 the gradient flows along $c$ without repeated products with a weight matrix, so it keeps longer dependencies than a plain RNN.
• Transformer parameters: about $12d^2$ per block: $4d^2$ for $W_Q, W_K, W_V, W_O$ and $8d^2$ for an MLP with hidden size $4d$. GPT-2 small ($L = 12$ blocks, $d = 768$, vocabulary 50,257, context 1,024): $12 \cdot 12 \cdot 768^2 = 84.9$M in the blocks, plus $38.6$M token and $0.8$M position embeddings: about $124$M.

Newer blocks: a mixture of experts sends each token to the top-$k$ of $E$ expert MLPs, so the parameters grow with $E$ and the compute per token with $k$ (mixture of experts); a state-space model (Mamba) replaces attention by a linear recurrence, linear in $T$ but weaker at exact recall from a long context (state-space models).

3.7 Generative models: autoregressive, diffusion, flow matching #

Notation: $x$ is a clean data sample and $\varepsilon \sim \mathcal N(0, I)$ is Gaussian noise of the same shape. Time conventions differ: in DDPM, $t = 0$ is data; in flow matching as written here, $t = 0$ is noise and $t = 1$ is data.

Autoregressive models. The chain rule of probability gives $p_\theta(x) = \prod_{i=1}^{n} p_\theta(x_i \mid x_{<i})$.

• Training: minimize $-\sum_i \log p_\theta(x_i \mid x_{<i})$ with teacher forcing; all positions train in parallel under a causal mask.
• Sampling: $n$ sequential calls; a KV cache makes each call cheap, but the calls cannot run in parallel.
• Likelihood: exact. Natural for discrete tokens; continuous data first needs a tokenizer (for example a VQ autoencoder). At sampling time the model conditions on its own outputs, which it never saw in training, so errors can compound (exposure bias).

VAE and GAN in one paragraph each. A VAE trains an encoder $q(z \mid x)$ and a decoder $p(x \mid z)$ by maximizing the ELBO (a reconstruction term minus a KL term to the prior); sampling is one decoder call; the likelihood is only bounded; samples tend to be blurry. A GAN trains a generator $G(z)$ against a discriminator in a minimax game; sampling is one call; there is no likelihood; training can be unstable, and mode collapse (the generator covers only part of the data) is the typical failure.

Diffusion (DDPM). The forward process is $x_t = \sqrt{\bar\alpha_t}\, x + \sqrt{1 - \bar\alpha_t}\,\varepsilon$ for $t = 1, \dots, T$, where $\bar\alpha_t$ falls from about 1 (clean) to about 0 (pure noise).

• Training: sample $t$ and $\varepsilon$, form $x_t$, and minimize $\mathbb E\,\|\varepsilon - \varepsilon_\theta(x_t, t)\|^2$. This is a reweighted ELBO. The network is a denoiser with three equivalent readings: the noise $\varepsilon_\theta$; the clean estimate $\hat x = (x_t - \sqrt{1 - \bar\alpha_t}\,\varepsilon_\theta)/\sqrt{\bar\alpha_t}$; and the score $\nabla_{x_t} \log p_t(x_t) \approx -\varepsilon_\theta/\sqrt{1 - \bar\alpha_t}$ (Tweedie's formula).
• Sampling: ancestral sampling takes $T$ steps ($T = 1000$ in Ho et al., 2020). The deterministic DDIM sampler and higher-order ODE solvers need about 10–50 network calls; distillation and consistency models need 1–4.
• Likelihood: a bound (the ELBO), or exact through the probability-flow ODE at a high cost.
• Conditioning: classifier-free guidance drops the condition $c$ for about 10–20% of the training examples, then samples with $\tilde\varepsilon = \varepsilon_\theta(x_t) + w\,\big(\varepsilon_\theta(x_t, c) - \varepsilon_\theta(x_t)\big)$. $w = 1$ is plain conditional sampling; $w > 1$ trades diversity for fidelity to $c$.

Flow matching (rectified flow). The path is $x_t = (1-t)\,\varepsilon + t\,x$ for $t \in [0, 1]$.

• Training: sample $t \sim U[0, 1]$, $\varepsilon$ and $x$, form $x_t$, and minimize $\mathbb E\,\|v_\theta(x_t, t) - (x - \varepsilon)\|^2$.
• Sampling: start from $\varepsilon$ and integrate $dy/dt = v_\theta(y, t)$ from $t = 0$ to $t = 1$ (Euler or Heun, about 10–50 steps).

Theorem (what flow matching learns). Let $x$ and $\varepsilon$ be independent. The minimizer of the loss is the conditional mean velocity

If $v^*$ is regular enough for the ODE to have unique solutions (for example, Lipschitz in $y$), the ODE $dy/dt = v^*(y, t)$ started from $\mathcal N(0, I)$ has the same distribution as $x_t$ at every $t$; at $t = 1$ it gives the data distribution.

Why this is true, in one line. A regression with squared loss learns the average target at each input, and moving probability mass with its average velocity keeps the densities right.

Proof (sketch).

1. The loss is a regression of the target $Y = x - \varepsilon$ on the input $(x_t, t)$.
2. For fixed $(y, t)$, $\mathbb E\big[\|v - Y\|^2 \mid x_t = y\big]$ is minimized over vectors $v$ by the conditional mean $v = \mathbb E[Y \mid x_t = y]$. To see it, write $v - Y = (v - \bar Y) + (\bar Y - Y)$ with $\bar Y = \mathbb E[Y \mid x_t = y]$; the cross term has conditional mean 0, so the loss is $\|v - \bar Y\|^2$ plus a constant.
3. Along each training line, $dx_t/dt = x - \varepsilon = Y$. So $v^*$ is the average velocity of the probability mass that passes through $y$ at time $t$. A field that moves the mass with its average velocity keeps the same densities $p_t$ (the continuity equation). $\blacksquare$

Diffusion and flow matching are one family (Gao et al., 2024). Both use a Gaussian path $x_t = \alpha_t x + \sigma_t \varepsilon$: flow matching has $\alpha_t = t$, $\sigma_t = 1 - t$; DDPM has $\alpha = \sqrt{\bar\alpha}$, $\sigma = \sqrt{1 - \bar\alpha}$. The network outputs convert into each other. With $\hat x = \mathbb E[x \mid x_t]$ and $\hat\varepsilon = \mathbb E[\varepsilon \mid x_t]$, the identity $x_t = (1-t)\hat\varepsilon + t\hat x$ gives

DDIM sampling is the same as Euler sampling of the flow-matching ODE, once you rescale $x_t$ and match the times by their noise level $\sigma_t/\alpha_t$. What really differs: (1) the noise schedule, (2) the network output ($\varepsilon$, $\hat x$ or $v$), (3) the weight of each $t$ in the loss, and (4) a stochastic (SDE) or deterministic (ODE) sampler.

"Straight paths" is only half true. Each training line is straight. But lines from different pairs cross, and $v^*$ averages them, so the learned paths curve. One Euler step from $t = 0$ lands on $\hat x = \mathbb E[x]$, the mean of the data (predict 2; worked example 4B). Fixes for few-step sampling: reflow (retrain on the model's own noise–sample pairs; Liu et al., 2022), distillation and consistency models.

	Autoregressive	Diffusion	Flow matching	VAE	GAN
network learns	$p(x_i \mid x_{<i})$	the noise (a denoiser, equivalent to the score)	the velocity $\mathbb E[x - \varepsilon \mid x_t]$	encoder $q(z \mid x)$, decoder $p(x \mid z)$	generator $G(z)$ against a discriminator
training loss	cross-entropy	MSE on $\varepsilon$ (a weighted ELBO)	MSE on the velocity	ELBO: reconstruction + KL	minimax game
network calls per sample	$n$, sequential	10–1000 (1–4 distilled)	10–50 (1–4 after reflow or distillation)	1	1
likelihood	exact	bound, or exact through the ODE (costly)	exact through the ODE (costly)	bound	none
typical failure	slow sampling, compounding errors	slow sampling	slow sampling, curved paths	blurry samples, posterior collapse	mode collapse, unstable training
typical use	LLMs, code, other discrete tokens	images, video, audio, weather ensembles	images, video, science	the latent space of latent diffusion	fast one-step generation, adversarial loss terms

Which to choose. Discrete sequences (text, code): autoregressive. Continuous high-dimensional data (images, audio, physical fields): diffusion or flow matching, usually in the latent space of an autoencoder, which makes each call much cheaper. Sampling speed first (large ensembles, real time): a distilled or reflowed model, or a GAN, and check that the diversity survives. An exact density (anomaly scores, likelihood-based inference): autoregressive or a normalizing flow.

3.8 Debugging a training run #

Two checks before any long run:

1. The first loss. A fresh classifier over $K$ classes is close to uniform, so its cross-entropy loss is about $\ln K$ (2.30 for 10 classes, 6.91 for 1,000).
2. Overfit one small batch to a loss near 0. If this fails, there is a bug, not a tuning problem.

Symptom	Common causes	First checks and fixes
first loss far from $\ln K$	last layer too large at init, unnormalized inputs, wrong loss reduction, wrong labels	small init of the last layer; check the inputs, the reduction and the label alignment
cannot overfit one batch	inputs and targets misaligned, zero_grad missing, frozen or detached parameters, softmax applied before CrossEntropyLoss	print shapes and a few labels; check requires_grad
loss becomes NaN or inf	learning rate too high, fp16 overflow, $\log 0$ or a division by 0, a bad sample	find the first non-finite tensor; lower $\eta$; bf16 or loss scaling; add a small constant
loss flat from the start	learning rate far off, dead ReLUs, bad init, unnormalized inputs	learning-rate range test; activation statistics per layer
spike in mid-training	a bad batch, $\eta$ too high for the current curvature, an outdated Adam $v$, growing attention or output logits	gradient-norm history; rewind and skip the batch; clipping; $\beta_2 = 0.95$; qk-norm; z-loss
training loss falls, validation loss rises	overfitting	regularize, augment, stop early, get more data
validation loss below training loss	dropout and augmentation act only in training; an easier validation set	evaluate the training set in eval mode
good in training, bad in serving	model.eval() forgotten (BatchNorm, dropout), different preprocessing	run both pipelines on the same input

Interactive figure: loading…

Answers to section 1. (1) $\mathbb E[h_{12,i}^2] = 2^{-12} \approx 2.4 \cdot 10^{-4}$: much smaller. Each layer has $n\sigma_w^2 = 512 \cdot \tfrac{1}{512} = 1$, but each ReLU halves the second moment. The RMS of the signal falls by $2^{-6} = 1/64$, and the backward pass shrinks the gradients in the same way. He initialization, $\sigma_w^2 = 2/512$, keeps the second moment at its starting value. (2) Every sample lands on 0, the mean of the data, and none on $\pm 1$. At $t = 0$ the position $\varepsilon$ says nothing about which data point its line goes to, so $v^*(\varepsilon, 0) = \mathbb E[x] - \varepsilon = -\varepsilon$, and one step gives $\varepsilon + (-\varepsilon) = 0$.

4. Worked example #

A. Adam + L2 vs AdamW on two weights. Both weights start at $w = 1$. Their gradients are pure noise with mean 0 and RMS 10 ("noisy") or RMS 0.01 ("quiet"). Take $\eta = 10^{-3}$ and $\lambda = 0.1$.

1. Adam + L2 adds $\lambda w = 0.1$ to each gradient. The decay part of the step is about $\eta\lambda w/\sqrt{\hat v}$, with $\sqrt{\hat v} \approx \sqrt{\text{RMS}^2 + (\lambda w)^2}$.
2. Noisy weight: $\sqrt{\hat v} \approx 10$, so the decay is $10^{-3} \cdot 0.1/10 = 10^{-5}$ per step: about $0.99$ after 1,000 steps.
3. Quiet weight: $\sqrt{\hat v} \approx \sqrt{0.01^2 + 0.1^2} \approx 0.1$, so the decay is about $10^{-3} \cdot 0.1/0.1 = 10^{-3}$ per step, 100 times more. The simulation in section 5 gives $0.26$ after 1,000 steps.
4. AdamW: both shrink by $1 - \eta\lambda = 1 - 10^{-4}$ per step, so $(1 - 10^{-4})^{1000} = 0.905$, whatever their gradient scale (the simulation gives 0.904, because the noisy gradient steps add a small random walk).

B. Flow matching on two points. Data $\pm 1$ with probability $1/2$ each, noise $\varepsilon \sim \mathcal N(0, 1)$, path $x_t = (1-t)\varepsilon + t x$.

1. Given $x$, $x_t \sim \mathcal N(tx, (1-t)^2)$. Bayes' rule gives the posterior mean $\hat x(y, t) = \mathbb E[x \mid x_t = y] = \tanh\!\big(ty/(1-t)^2\big)$.
2. From $v^* = (\hat x - x_t)/(1 - t)$: $v^*(y, t) = \big(\hat x(y, t) - y\big)/(1-t)$.
3. One Euler step from $\varepsilon = 0.5$: $v^*(0.5, 0) = (\tanh 0 - 0.5)/1 = -0.5$, so $y = 0.5 - 0.5 = 0$.
4. Two Euler steps of size 0.5: first $y = 0.5 + 0.5 \cdot (-0.5) = 0.25$ at $t = 0.5$. There $\hat x = \tanh(0.5 \cdot 0.25/0.25) = \tanh(0.5) = 0.462$ and $v^* = (0.462 - 0.25)/0.5 = 0.424$, so $y = 0.25 + 0.5 \cdot 0.424 = 0.462$.
5. The exact ODE sends every $\varepsilon > 0$ to $+1$ and every $\varepsilon < 0$ to $-1$.

The last Euler step to $t = 1$ always lands on $\hat x$, the posterior mean at the current point. Few steps give averages (blurry samples); more steps commit to one mode.

C. One DDIM step by hand. A DDPM model at a step with $\bar\alpha_t = 0.36$ sees $x_t = 1.0$ and predicts $\varepsilon_\theta = 0.5$. So $x_t = 0.6\,x + 0.8\,\varepsilon$.

1. Clean estimate: $\hat x = (1.0 - 0.8 \cdot 0.5)/0.6 = 1.0$.
2. Score: $-\varepsilon_\theta/\sqrt{1 - \bar\alpha_t} = -0.5/0.8 = -0.625$.
3. Deterministic DDIM step to $\bar\alpha_{\text{prev}} = 0.64$: keep $\hat x$ and $\varepsilon_\theta$ and mix them again, $x_{\text{prev}} = \sqrt{0.64}\,\hat x + \sqrt{0.36}\,\varepsilon_\theta = 0.8 \cdot 1.0 + 0.6 \cdot 0.5 = 1.1$.

5. In code #

NumPy and SciPy check the claims of sections 3 and 4. The notebook has these cells and the exercises.

import numpy as np
from scipy.integrate import solve_ivp

# (A) Adam + L2 vs AdamW on two groups of 1,000 weights (gradients are pure noise)
def run(decoupled, steps=1000, lr=1e-3, wd=0.1, b1=0.9, b2=0.999, eps=1e-8):
    rng = np.random.default_rng(0)
    w = np.ones((2, 1000))                     # row 0: noisy weights, row 1: quiet weights
    scale = np.array([[10.0], [0.01]])         # RMS of the gradient noise per row
    m, v = np.zeros_like(w), np.zeros_like(w)
    for t in range(1, steps + 1):
        g = scale * rng.standard_normal(w.shape)
        if decoupled:
            w = w * (1 - lr * wd)              # AdamW: decay outside the adaptive step
        else:
            g = g + wd * w                     # Adam + L2: decay inside the gradient
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g * g
        mh, vh = m / (1 - b1**t), v / (1 - b2**t)
        w = w - lr * mh / (np.sqrt(vh) + eps)
    return w.mean(axis=1)

print("Adam + L2:", run(False).round(3))      # [0.99  0.262]
print("AdamW    :", run(True).round(3))       # [0.904 0.904]

# Predict 1: 12 ReLU layers of width 512, averaged over 30 random networks
x = np.random.default_rng(0).standard_normal((128, 512))
for name, var in [("1/n", 1 / 512), ("2/n (He)", 2 / 512)]:
    ms = []
    for seed in range(30):
        r = np.random.default_rng(seed)
        h = x
        for _ in range(12):
            h = np.maximum(h @ (np.sqrt(var) * r.standard_normal((512, 512))), 0)
        ms.append(np.mean(h ** 2))
    print(name, f"{np.mean(ms):.2e}")          # 1/n 2.41e-04 (2^-12 = 2.44e-04); 2/n (He) 9.88e-01

# (B) flow matching on data {-1, +1}: one Euler step, two steps, and the exact ODE
def velocity(y, t):
    """Exact field E[x - eps | x_t = y]; E[x | x_t = y] = tanh(t y / (1 - t)^2)."""
    return (np.tanh(t * y / (1 - t) ** 2) - y) / (1 - t)

eps = np.array([-1.3, -0.2, 0.05, 0.5, 2.0])            # five noise samples
print(eps + 1.0 * velocity(eps, 0.0))                   # one step: [0. 0. 0. 0. 0.]
y_half = 0.5 + 0.5 * velocity(0.5, 0.0)                 # two steps from eps = 0.5
print(round(y_half, 3), round(y_half + 0.5 * velocity(y_half, 0.5), 3))   # 0.25 0.462
sol = solve_ivp(lambda t, y: velocity(y, t), (0.0, 0.999), eps, rtol=1e-10, atol=1e-12)
print(sol.y[:, -1].round(3))                            # [-1.    -0.998  0.997  0.999  1.001]

The ODE stops at $t = 0.999$ because $v^*$ is singular at $t = 1$. There the state still holds a small share of noise ($x_t = 0.001\,\varepsilon + 0.999\,x$), so the values sit within about 0.003 of $\pm 1$. Every negative $\varepsilon$ goes to $-1$ and every positive $\varepsilon$ goes to $+1$.

6. Where it shows up #

• ML practice. A standard transformer recipe is sections 3.1–3.5 together: AdamW, warmup and a cosine or WSD decay, clipping at 1.0, bf16 with fp32 master weights, pre-LN with RMSNorm. When a run fails, the table of section 3.8 is the first pass. In ML-coding rounds (lessons 07–08) these are the follow-up questions.
• Generative AI. Text and code models are autoregressive. Image and video models (Stable Diffusion 3, FLUX, video diffusion models) use diffusion or flow matching in a latent space, with classifier-free guidance and few-step samplers.
• ML for science. Many weather and Earth-system models are transformers on patches of gridded fields; others, such as GraphCast, are GNNs on a mesh. Small batches per GPU rule out BatchNorm. Probabilistic forecasts need generative models: GenCast samples weather ensembles with a diffusion model. Ensemble size times calls per sample decides whether you need few-step sampling. Inductive bias decides the backbone: CNNs and neural operators for regular grids, GNNs for meshes and sensor networks, equivariant networks for molecules.

7. Common mistakes #

• "AdamW is Adam with L2." In Adam, the L2 gradient is divided by $\sqrt{\hat v}$; AdamW decouples the decay so that every weight shrinks by $1 - \eta\lambda$.
• "BatchNorm works by removing internal covariate shift." Say what it does: it normalizes each channel, smooths the optimization and allows larger learning rates.
• Forgetting model.eval() (BatchNorm, dropout), or scaling again at test time with inverted dropout, which already scales in training.
• Xavier or LeCun initialization for a deep ReLU network (each layer halves the second moment; use He), or zero initialization of all weights (no symmetry breaking).
• "Flow-matching paths are straight, so one step is enough." Only each training line is straight; the learned field curves, and one step gives the data mean.
• "Diffusion and flow matching are different model classes." With Gaussian paths they are one family. Also watch the time convention: $t = 0$ is data in DDPM and noise in rectified flow.

8. Check yourself #

1. Why do LLM recipes use $\beta_2 = 0.95$ instead of Adam's default 0.999?
   Answer

   $v$ averages $g^2$ over about $1/(1-\beta_2)$ steps: 20 steps for 0.95, 1,000 for 0.999. With 0.999, a parameter whose gradients were tiny for a long time has a tiny $v$. When its gradient suddenly grows, $\hat m$ follows within about 10 steps but $\sqrt{\hat v}$ lags, so the update $\eta\,\hat m/\sqrt{\hat v}$ is too large for many steps, often in many parameters at once: a loss spike. A lower $\beta_2$ lets $v$ follow the gradient scale faster.
2. Estimate the parameters of GPT-2 medium: 24 blocks, $d = 1024$, vocabulary 50,257, context 1,024.
   Answer

   Blocks: $12 \cdot 24 \cdot 1024^2 = 302$M. Token embeddings $50{,}257 \cdot 1024 = 51.5$M, position embeddings $1.0$M. Total about $355$M (biases and norms add less than 0.1%).
3. Why are attention scores divided by $\sqrt{d_k}$? What goes wrong without it for $d_k = 128$?
   Answer

   With unit-variance, independent entries, $q \cdot k$ has variance $d_k$. Without the scaling its standard deviation is $\sqrt{128} \approx 11.3$, so the softmax is close to one-hot at initialization, and its gradients are close to 0 for most entries. Dividing by $\sqrt{d_k}$ gives scores of variance 1.
4. A flow-matching network outputs $v_\theta(x_t, t)$. How do you get its estimate of the clean sample?
   Answer

   $\hat x = x_t + (1-t)\,v_\theta(x_t, t)$, from $v^* = (\hat x - x_t)/(1-t)$. A single Euler step from $t = 0$ to $1$ computes exactly this $\hat x$ at $t = 0$, where $x_t = \varepsilon$ says nothing about the data, so $\hat x = \mathbb E[x]$: the mean of the data, an average of all modes.
5. Inverted dropout with $p = 0.2$: an activation is 3.0. What values can it take in training, and what is its expected value?
   Answer

   0 with probability 0.2, and $3.0/(1 - 0.2) = 3.75$ with probability 0.8. The expected value is $0.8 \cdot 3.75 = 3.0$, the same as without dropout, so nothing changes at eval.
6. What is the output size of a convolution with input 64, kernel 5, stride 2, padding 2? And the receptive field of three stacked $3 \times 3$ convolutions with stride 1?
   Answer

   $\lfloor (64 + 4 - 5)/2 \rfloor + 1 = \lfloor 31.5 \rfloor + 1 = 32$. Receptive field $1 + 3 \cdot 2 = 7$, so $7 \times 7$.

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

1. What weight variance does He initialization use, and why the factor 2?

   Show answer

2. 12 ReLU layers with $\sigma_w^2 = 1/n$ and $\mathbb E[x^2] = 1$. What is the second moment of the last layer?

   Show answer

3. How does BatchNorm differ between training and evaluation?

   Show answer

4. Write pre-LN and post-LN for one transformer sub-layer.

   Show answer

5. In inverted dropout with drop probability $p$, what happens to the kept activations in training and at eval?

   Show answer

6. Why does AdamW decay weights outside the adaptive step?

   Show answer

7. How does the cost of a transformer layer grow with the sequence length $T$, and why?

   Show answer

8. What is the training loss of a DDPM diffusion model?

   Show answer

9. What does a flow-matching network learn, and what does one Euler step from noise give?

   Show answer

10. What should the first loss of a fresh $K$-class classifier be?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

Recognize the question from the wording.

The question says	Start with
"how do you initialize a deep network?"	the variance theorem, ReLU halves, He $2/n_{\text{in}}$, symmetry breaking
"why residual connections?", "why can ResNets be deep?"	the Jacobian $(I + J)$; the identity path for the gradient
"BatchNorm vs LayerNorm", "what changes at inference?"	the axis of the statistics; running averages; small batches; why transformers use LN/RMSNorm
"SGD vs Adam", "why AdamW?"	per-coordinate steps, bias correction, the decoupling proposition, $\lambda = 0.1$, $\beta_2 = 0.95$
"why warmup?", "loss spikes"	curvature at init, Adam's young $\hat v$, post-LN; clipping, $\beta_2$, qk-norm
"CNN or ViT?", "why $\sqrt{d_k}$?"	inductive bias vs data; variance of $q \cdot k$
"diffusion vs flow matching vs autoregressive"	what is learned, the loss, the calls, the likelihood, when to use; "one family"
"my training does not work"	first loss $\ln K$, overfit one batch, then the symptom table

Typical questions and short model answers.

1. "How do you initialize a deep ReLU network, and why?" — "I keep the variance of the activations constant from layer to layer. Each pre-activation is a sum of $n$ terms, so its variance is $n\sigma_w^2$ times the second moment of the input. ReLU zeroes half of the values, so it halves the second moment. So I use $\sigma_w^2 = 2/n_{\text{in}}$, He initialization. With $1/n$, a 12-layer network shrinks the signal to $2^{-12}$, and the gradients shrink the same way."
2. "BatchNorm vs LayerNorm?" — "BatchNorm normalizes each channel over the batch; in training it uses the batch statistics and in eval running averages. That makes it depend on the batch size and makes training and eval different functions. LayerNorm normalizes over the features of one sample, so it is the same in training and eval and works with batch size 1 and with autoregressive decoding. Transformers use LayerNorm or RMSNorm, which drops the mean. BatchNorm still works well in CNNs with large batches."
3. "Why AdamW and not Adam with L2?" — "For SGD they are the same. In Adam the L2 gradient goes into $m$ and $v$ and is divided by $\sqrt{\hat v}$, so weights with large gradients are barely decayed. AdamW applies $\theta \leftarrow (1 - \eta\lambda)\theta$ directly, so every weight shrinks at the same rate. Typical LLM values: $\lambda = 0.1$, $\beta_2 = 0.95$, warmup, then cosine decay."
4. "Explain diffusion vs flow matching vs autoregressive." — "Autoregressive models predict the next token with cross-entropy, give an exact likelihood, and need one call per token. Diffusion adds Gaussian noise and trains a network to predict it; sampling removes noise in 10–50 ODE steps, or 1–4 after distillation. Flow matching joins noise and data by straight lines and regresses the velocity; the network learns the average velocity, so the paths curve and one step gives the data mean. With Gaussian paths, diffusion and flow matching are the same family with different schedules, outputs and loss weights. Text: autoregressive; images and other continuous data: diffusion or flow matching in a latent space."
5. "Why residual connections?" — "The block computes $x + F(x)$, so its Jacobian is $I + J$. The product over many blocks contains the identity, so the gradient reaches early layers even when the $J$'s are small. Without it the product of $L$ Jacobians vanishes or explodes. In transformers, pre-LN keeps this identity path free of normalization, which is why it trains without warmup."
6. "The loss is NaN after 300 steps. What do you do?" — "Find the first non-finite tensor, then check the usual causes in order: a learning rate too high (try 3× smaller), fp16 overflow (bf16 or loss scaling), a $\log 0$ or a division by 0 in my code, a bad sample. I log the gradient norm and use clipping. Before that, I check that the first loss was $\ln K$ and that the model can overfit one batch."

Typical follow-ups. "What happens with batch size 1 in BatchNorm during training?" (in a conv layer the statistics come only from the spatial positions of one image, so they are very noisy; without spatial positions every value equals the batch mean, the normalized output is 0, and PyTorch raises an error). "Why does post-LN need warmup and pre-LN not?" (post-LN has large gradients near the output at initialization; pre-LN keeps a clean identity path). "What is classifier-free guidance, and what does a large $w$ do?" (the formula of section 3.7; more fidelity to the prompt, less diversity). "How many parameters per transformer block?" ($12d^2$).

Mock prompt (debug this training run). "A 1.3B-parameter pre-LN transformer trains in bf16 with AdamW ($\eta = 3 \cdot 10^{-4}$, $\beta_2 = 0.999$, 2,000 warmup steps, clipping at 1.0). At step 41,000 the loss jumps from 2.4 to 6.8 and the gradient norm from 0.3 to 30; it recovers to 2.9 after 2,000 steps. What do you do?" A strong answer contains: rewind to a checkpoint a few hundred steps before the spike; inspect or skip the batches around step 41,000, and replay the same data order to test whether the data caused it; per-layer gradient norms, the largest attention logit, the output logits and the Adam update size; mitigations ordered by cost (lower peak $\eta$ or longer warmup, $\beta_2 = 0.95$, qk-norm, z-loss, fp32 for softmax, norms and loss); a cheaper reproduction on a small model at a higher learning rate; an alert on the gradient norm.

10. Practice #

The notebook for this lesson runs in Colab, in NumPy. Each exercise has a check cell and a folded solution:

1. forward_variance: the second moment per layer of a deep ReLU MLP for several initialization scales; check the factor $2^{-L}$ and the flat line for He.
2. batchnorm_forward: BatchNorm in training and eval mode with running statistics; check against a direct computation and the batch-size-1 failure.
3. inverted_dropout: dropout with the $1/(1-p)$ scaling; check that the expected value is preserved and eval is the identity.
4. fm_velocity and euler_sample: the exact flow-matching velocity for two-point data and Euler sampling with $k$ steps; check that one step gives the mean and many steps give the two points.
5. first_loss: the cross-entropy of a randomly initialized softmax classifier; check that it is close to $\ln K$ with a small last layer and far from it with a large one.

Timed practice (25 minutes each, talking aloud): "derive He initialization"; "implement LayerNorm and BatchNorm forward, and say what changes at eval"; "explain diffusion vs flow matching on a whiteboard"; "the loss spikes at step 41,000: debug it"; "count the parameters of a transformer block"; the questions on initialization, normalization and optimization in Kashani and Ivry, Deep Learning Interviews (see below).

Say it aloud drill. Start a 2-minute timer for each question of section 9. Answer with the four parts (definition, mechanism, trade-off, number), record yourself, and compare with section 3.

11. Watch and read #

Building makemore Part 3: Activations & Gradients, BatchNormAndrej Karpathy · 1 h 56 min

Activation and gradient statistics per layer, the first-loss check, Kaiming initialization and BatchNorm in code.

Lecture 6 | Training Neural Networks IStanford University School of Engineering · CS231n (2017) · 1 h 20 min

Activation functions, weight initialization from variance, and batch normalization.

Adam Optimization Algorithm (C2W2L08)DeepLearningAI · 7 min

The Adam update with momentum, RMS scaling and bias correction, in a short lecture.

MIT 6.S184: Flow Matching and Diffusion Models - Lecture 01 - Flow and Diffusion Models (2026)Peter Holderrieth · 1 h 18 min

Flows and diffusion as ODEs and SDEs from noise to data, the setting of section 3.7.

• Simon J. D. Prince, Understanding Deep Learning (MIT Press, 2023, free PDF): §6.3–6.4 (momentum, Adam), §7.5 (initialization), §11.2–11.4 (residual blocks, batch normalization), chapter 12 (transformers), chapter 18 (diffusion models).
• Ruiqi Gao, Emiel Hoogeboom, Jonathan Heek, Valentin De Bortoli, Kevin Murphy and Tim Salimans (2024), Diffusion Meets Flow Matching: Two Sides of the Same Coin: DDIM as the Euler sampler of flow matching, and why the paths are not straight.
• Loshchilov and Hutter (2019), Decoupled Weight Decay Regularization: §2 (Adam + L2 vs AdamW).
• He, Zhang, Ren and Sun (2015), Delving Deep into Rectifiers: §2.2 (the derivation of He initialization).
• Xiong et al. (2020), On Layer Normalization in the Transformer Architecture: pre-LN vs post-LN and warmup.
• Santurkar, Tsipras, Ilyas and Madry (2018), How Does Batch Normalization Help Optimization?.
• Lipman et al. (2023), Flow Matching for Generative Modeling, and Liu, Gong and Liu (2022), Flow Straight and Fast: rectified flow.
• Andrej Karpathy, A Recipe for Training Neural Networks: the debugging habits of section 3.8.
• Shlomo Kashani and Amir Ivry, Deep Learning Interviews (free): chapter 8 (layers, training, optimization) with solutions.