Depth-enablers

Last lecture assembled the machinery to train a network: automatic differentiation, stochastic gradient descent, and the schedules, momentum, and adaptive scaling to spend the gradient well. It also left a loose thread from the Neural Networks lecture. When we hand-computed a backward pass there, one hidden neuron was inactive, its ReLU derivative was zero, and the entire gradient path running through it went to zero with it. In a two-layer network that costs one neuron; in a fifty-layer network it compounds fifty times over, and for decades it decided whether deep networks could be trained at all.

Why does stacking more layers sometimes make a network harder to train, rather than simply slower?

Backpropagation through \(L\) layers multiplies \(L\) matrices, and a product of \(L\) factors is settled by whether each sits a little above or a little below \(1\). Today’s three tools (initialization, residual connections, and normalization) are three ways of pinning that per-layer factor near \(1\).

Vanishing and Exploding Gradients

Recall the matrix view of backpropagation from the Neural Networks lecture: the gradient crosses each layer by a matrix multiplication. Write \(\mathbf{h}^{(l)} \in \mathbb{R}^{d_l}\) for the activations of layer \(l\), with \(\mathbf{h}^{(0)} = \mathbf{x}\) the input and \(\mathbf{h}^{(L)}\) the output. Collect layer \(l\)’s derivatives into one matrix, its Jacobian: \[ \mathbf{J}^{(l)} = \frac{\partial \mathbf{h}^{(l)}}{\partial \mathbf{h}^{(l-1)}} \in \mathbb{R}^{d_l \times d_{l-1}}. \] For a linear layer followed by an elementwise activation, the Jacobian is the weight matrix \(\mathbf{W}^{(l)}\) with row \(i\) scaled by neuron \(i\)’s activation derivative, packaging that lecture’s mask-then-multiply backward pass into one matrix. One backward step across a layer multiplies the gradient by the transpose of that layer’s Jacobian: \[ \nabla_{\mathbf{h}^{(l-1)}}\mathcal{L} = \big(\mathbf{J}^{(l)}\big)^\top \nabla_{\mathbf{h}^{(l)}}\mathcal{L} . \] Applying that rule once per layer, from the output back to the input, puts \(L\) transposed Jacobians in a row: \[ \nabla_{\mathbf{h}^{(0)}}\mathcal{L} = \big(\mathbf{J}^{(1)}\big)^\top \big(\mathbf{J}^{(2)}\big)^\top \cdots \big(\mathbf{J}^{(L)}\big)^\top \, \nabla_{\mathbf{h}^{(L)}}\mathcal{L}. \] This is exact; the trouble is what a product of \(L\) matrices does to the length of the vector it multiplies, and one number per layer bounds it. That number is the spectral norm, the largest factor by which a matrix stretches any vector: \[ \|\mathbf{A}\|_2 = \max_{\mathbf{x}\neq\mathbf{0}}\frac{\|\mathbf{A}\mathbf{x}\|}{\|\mathbf{x}\|}, \] which the Decompositions lecture identified as \(\mathbf{A}\)’s largest singular value \(\sigma_{\max}\).

Claim: The gradient reaching the input is controlled by the product of the layers’ spectral norms: \[ \big\|\nabla_{\mathbf{h}^{(0)}}\mathcal{L}\big\| \leq \Big(\prod_{l=1}^{L} \big\|\mathbf{J}^{(l)}\big\|_2\Big) \big\|\nabla_{\mathbf{h}^{(L)}}\mathcal{L}\big\| . \]

Proof of Claim

The spectral norm gives \(\|\mathbf{A}\mathbf{x}\| \leq \|\mathbf{A}\|_2 \|\mathbf{x}\|\) for every \(\mathbf{x}\), since \(\|\mathbf{A}\|_2\) is by definition the largest such ratio. Apply that once per matrix, peeling from the outside in: \[ \begin{align*} \big\|(\mathbf{J}^{(1)})^\top (\mathbf{J}^{(2)})^\top \cdots (\mathbf{J}^{(L)})^\top \nabla_{\mathbf{h}^{(L)}}\mathcal{L}\big\| &\leq \big\|(\mathbf{J}^{(1)})^\top\big\|_2 \, \big\|(\mathbf{J}^{(2)})^\top \cdots (\mathbf{J}^{(L)})^\top \nabla_{\mathbf{h}^{(L)}}\mathcal{L}\big\| \\&\leq \big\|(\mathbf{J}^{(1)})^\top\big\|_2 \big\|(\mathbf{J}^{(2)})^\top\big\|_2 \, \big\|(\mathbf{J}^{(3)})^\top \cdots (\mathbf{J}^{(L)})^\top \nabla_{\mathbf{h}^{(L)}}\mathcal{L}\big\| \\&\;\;\vdots \\&\leq \Big(\prod_{l=1}^{L}\big\|\mathbf{J}^{(l)}\big\|_2\Big) \big\|\nabla_{\mathbf{h}^{(L)}}\mathcal{L}\big\|, \end{align*} \] where the last line used \(\|\mathbf{A}^\top\|_2 = \|\mathbf{A}\|_2\) at every step, since a matrix and its transpose have the same singular values.

The bound applies one multiplicative factor per layer, and it is one-sided: a small product proves the gradient at the input is small, while a large product proves nothing, the way any upper bound can be loose.

Suppose every layer’s Jacobian has spectral norm roughly \(\rho\), so the bound behaves like \(\rho^L\). Take \(\rho = 0.9\) and \(L = 50\): \(0.9^{50} \approx 0.005\), and whatever gradient left the output has all but vanished by the first layer; at \(\rho = 1.1\) the same arithmetic gives \(117\) instead. Neither is a strange number for a matrix, and nothing in the model asks the layers to sit at exactly \(\rho = 1\): a quantity raised to the \(L\)th power is decided by the third digit of its base, which is what makes depth different from width. So where does \(\rho\) come from? Before the first gradient step, only one place: the distribution we drew the weights from.

Keeping Variance Constant: Initialization

We can find the right weight scale by tracking a single number, the variance of a neuron’s pre-activation.

Consider one neuron in a layer with \(n_{\text{in}}\) inputs, whose pre-activation is the weighted sum of the values feeding it: \[ z = \sum_{i=1}^{n_{\text{in}}} w_i x_i, \] where the inputs \(x_i\) and weights \(w_i\) are independent of one another and mean zero, with common variances \(\mathrm{Var}(x)\) and \(\mathrm{Var}(w)\).

Claim: The pre-activation’s variance is the product of three factors, one of them the layer’s width: \[ \mathrm{Var}(z) = n_{\text{in}} \cdot \mathrm{Var}(w) \cdot \mathrm{Var}(x). \]

Proof of Claim

The \(n_{\text{in}}\) terms \(w_i x_i\) are independent of one another, so in the variance-of-a-sum claim from the Probability lecture every covariance term vanishes and the variances simply add: \[ \mathrm{Var}(z) = \sum_{i=1}^{n_{\text{in}}} \mathrm{Var}(w_i x_i). \] Now expand one of those terms: \[ \begin{align*} \mathrm{Var}(w_i x_i) &= \mathbb{E}\big[(w_i x_i)^2\big] - \big(\mathbb{E}[w_i x_i]\big)^2 \\&= \mathbb{E}[w_i^2]\,\mathbb{E}[x_i^2] - \big(\mathbb{E}[w_i]\,\mathbb{E}[x_i]\big)^2 \\&= \mathbb{E}[w_i^2]\,\mathbb{E}[x_i^2] \\&= \mathrm{Var}(w_i)\,\mathrm{Var}(x_i), \end{align*} \] where we used the definition of variance, then independence to split both expectations of a product, then \(\mathbb{E}[w_i] = 0\) to kill the second term, and finally \(\mathrm{Var}(w_i) = \mathbb{E}[w_i^2] - (\mathbb{E}[w_i])^2 = \mathbb{E}[w_i^2]\) for a mean-zero variable (and the same for \(x_i\)). Summing \(n_{\text{in}}\) identical terms gives the claim.

Widening a layer therefore inflates the signal unless the weights shrink to compensate, which pins the only weight scale that passes the signal through unchanged: \[ \mathrm{Var}(w) = \frac{1}{n_{\text{in}}} \quad\Longrightarrow\quad \mathrm{Var}(z) = \mathrm{Var}(x). \] Both initializations used in practice are this scale, adjusted for what sits on either side of the layer. Xavier initialization draws weights with \(\mathrm{Var}(w) = 2/(n_{\text{in}} + n_{\text{out}})\), splitting the difference between the forward pass, which wants \(1/n_{\text{in}}\), and the backward pass, which wants \(1/n_{\text{out}}\) since the gradient sums over the outputs instead. He initialization uses \(\mathrm{Var}(w) = 2/n_{\text{in}}\), where the extra \(2\) compensates for a ReLU zeroing about half of its inputs and so halving the variance that gets through.

The class demo runs three initialization scales through a \(30\)-layer \(\tanh\) network and reads the predicted slopes off a real gradient, vanishing at one end and exploding at the other.

Initialization, though, fixes the scale only at step zero; gradient descent then moves every weight, and nothing preserves \(\mathrm{Var}(w) = 1/n_{\text{in}}\) as it does. We need a fix that survives training.

Residual Connections

A residual (skip) connection gives up on controlling the per-layer factor and instead builds a path on which the factor is \(1\) by construction. Writing \(F^{(l)}\) for whatever the block computes, the block adds that output onto its own input rather than replacing it: \[ \mathbf{h}^{(l)} = \mathbf{h}^{(l-1)} + F^{(l)}\big(\mathbf{h}^{(l-1)}\big), \] which requires equal input and output widths, so a residual stack keeps one width \(d_l = d_{l-1}\) throughout. Differentiating, the block’s Jacobian picks up an identity term: \[ \mathbf{J}^{(l)} = \mathbf{I} + \mathbf{J}_F^{(l)}, \qquad \text{where } \mathbf{J}_F^{(l)} = \frac{\partial F^{(l)}}{\partial \mathbf{h}^{(l-1)}} \in \mathbb{R}^{d_l\times d_l} \text{ is the learned part's Jacobian}. \] If the learned part collapses, so \(\mathbf{J}_F^{(l)} \to \mathbf{0}\), the block’s Jacobian becomes the identity and passes the gradient through unchanged: the first section’s vanishing scenario is unreachable by any setting of the weights.

Suppose every learned Jacobian is small, \(\|\mathbf{J}_F^{(l)}\|_2 \leq \epsilon\) for some \(\epsilon < 1\), a fair description of a freshly initialized block; then the triangle inequality bounds the block’s effect on any vector \(\mathbf{x}\) from both sides: \[ (1-\epsilon)\|\mathbf{x}\| \leq \big\|(\mathbf{I} + \mathbf{J}_F^{(l)})\mathbf{x}\big\| \leq (1+\epsilon)\|\mathbf{x}\|, \] where the lower bound is the half a plain block cannot offer. Applying this once per block, a stack of \(L\) of them changes the gradient’s length by a factor somewhere between \((1-\epsilon)^L\) and \((1+\epsilon)^L\). Those are still exponentials in \(L\), but now we control the base: at \(\epsilon = 1/L\) the two ends are \((1 \mp 1/L)^L\), which sit between roughly \(1/e\) and \(e\) no matter how deep the network is. This is why residual networks hundreds of blocks deep train at all.

Problem 13 takes the same product apart a second way, expanding \(\prod_{l=1}^{L} (\mathbf{I} + \mathbf{J}_F^{(l)})\) into a sum with one term for each of the \(2^L\) choices of which blocks to route through and which to skip, and then asks which of those terms carry the mass.

What residual connections do not do is control the size of the forward signal. Each block adds its output onto the running stream, so the stream’s variance accumulates: in the thirty-block network of the class demo, the activation standard deviation climbs from about \(1.1\) at the first block to about \(3.7\) at the last. Nothing has broken, but a quantity that grows with depth is exactly what this lecture is trying to get rid of.

Normalization Layers

Normalization layers attack the scale directly, continuously during training rather than once at initialization. A normalization layer takes a vector \(\mathbf{z} \in \mathbb{R}^m\), centers and rescales it into \(\hat{\mathbf{z}}\), then applies a learned scale and shift to produce its output \(\tilde{\mathbf{z}}\): \[ \hat z_j = \frac{z_j - \mu}{\sqrt{\sigma^2 + \delta}}, \qquad \tilde z_j = \gamma_j \hat z_j + \beta_j, \] where \(\mu\) and \(\sigma^2\) are the mean and variance of the entries being normalized, \(\delta > 0\) is a small constant keeping the division safe, and \(\gamma_j, \beta_j \in \mathbb{R}\) are learned parameters, one pair per entry. That pair makes fixing the scale cost no expressiveness: the network can restore any scale and offset, but now it has to learn them. The two common versions differ only in which entries \(\mu\) and \(\sigma^2\) are computed over: for a batch of \(B\) examples with pre-activations \(z_j^{(b)}\) (example \(b\), feature \(j\)), there are two natural sets to average: \[ \mu^{\text{batch}}_j = \frac{1}{B}\sum_{b=1}^{B} z^{(b)}_j, \qquad \mu^{\text{layer},b} = \frac{1}{m}\sum_{j=1}^{m} z^{(b)}_j, \] with the corresponding variances defined over the same entries. Batch normalization uses the first, one mean per feature averaged down the batch; layer normalization uses the second, one mean per example across its own features. In practice, batch normalization’s output for one example depends on which other examples shared its batch, so it is noisy when batches are small and needs running averages at test time; layer normalization has neither problem and is the standard choice for the sequence models of the Architectures unit.

Either version removes the layer’s overall scale, and that is what makes it a depth tool rather than a preprocessing trick: multiply the layer’s weights by any constant \(c > 0\), and every pre-activation scales by \(c\), so the mean becomes \(c\mu\), the variance \(c^2\sigma^2\), and (taking \(\delta\) negligible) the constant cancels: \[ \frac{c z_j - c\mu}{\sqrt{c^2\sigma^2}} = \frac{c\,(z_j - \mu)}{c\sqrt{\sigma^2}} = \frac{z_j - \mu}{\sqrt{\sigma^2}} = \hat z_j . \] A layer whose output does not depend on the scale of its own weights cannot contribute a runaway factor to the product, however far training moves those weights.

The class demo runs both fixes on the same badly initialized network: residual connections let the forward signal grow, the accumulation the previous section warned about, while layer normalization pins it to exactly \(1\) at every layer, which is all that normalizing means. On the gradient side the two fixes land at very different absolute scales, and that is not the point: neither curve decays with depth, the only property the product bound cares about.

Both fixes control one number per layer, the largest factor by which the layer can stretch a vector; but a Jacobian has \(\min(d_l, d_{l-1})\) stretch factors, and the bound only ever looked at the largest.

Keeping the Jacobians Well Conditioned

The Optimization lecture gave the ratio of a matrix’s largest to smallest singular value a name, the condition number: \[ \kappa(\mathbf{A}) = \frac{\sigma_{\max}}{\sigma_{\min}}, \] and blamed it for the precision of the closed form and the speed of gradient descent; it comes back here in a third role. A layer with \(\sigma_{\max} = 1\) and \(\sigma_{\min} = 0.1\) has spectral norm exactly \(1\), so the product bound predicts no trouble, yet a gradient along the smallest singular direction shrinks by \(0.1\) per layer while one along the largest is untouched: after \(L\) layers the two differ by \(\kappa^L\), and the opening problem returns one singular direction at a time.

So the real target is stronger than “spectral norm near \(1\)”: we want every singular value near \(1\), each layer’s Jacobian close to a rotation, a map that turns a vector without changing its length. Our three tools approximate it to different degrees. Variance-preserving initialization only sets the average squared singular value to \(1\), and for a square layer \(\mathbf{W}\) of width \(n_{\text{in}}\) we can see why: \[ \sum_{i=1}^{n_{\text{in}}} \sigma_i^2 = \sum_{i,j}\big[\mathbf{W}\big]_{i,j}^2, \qquad \mathbb{E}\Big[\sum_{i,j}\big[\mathbf{W}\big]_{i,j}^2\Big] = n_{\text{in}}^2\cdot\frac{1}{n_{\text{in}}} = n_{\text{in}}, \] where the first identity writes the sum of squared singular values as the sum of squared entries and the second is linearity of expectation, leaving \(n_{\text{in}}\) of mass shared among \(n_{\text{in}}\) squared singular values, an average of \(1\). Nothing in an average constrains the spread, so the individual \(\sigma_i\) are free to range widely. A residual block does better: the two-sided bound above puts every stretch factor of \(\mathbf{I} + \mathbf{J}_F^{(l)}\) in \([1-\epsilon, 1+\epsilon]\), bounding the ratio of the largest to the smallest: \[ \kappa\big(\mathbf{I} + \mathbf{J}_F^{(l)}\big) \leq \frac{1+\epsilon}{1-\epsilon}, \] which at \(\epsilon = 0.1\) is at most \(1.22\), a nearly perfectly conditioned layer whatever the block has learned. Normalization sits in between: it removes the overall scale, fixing \(\sigma_{\max}\), while leaving the shape of the rest of the spectrum alone.

We will meet this target once more, from the other side. In the Low-Rank Structure unit, Muon is an optimizer whose entire job is to replace a gradient matrix’s singular values by \(1\) before stepping with it, so that the update is well conditioned by construction rather than by luck.

Looking Forward

Carry one sentence out of today: when a quantity will be raised to the \(L\)th power, do not tune its base; build the computation so the base is \(1\) before anything is learned.

Everything today was about whether a deep network can be trained, and a network that trains perfectly can still be useless: the Methodology lecture’s U-curve says that a model with enough capacity fits its training data exactly and generalizes worse for it, and the networks we can now train have far more parameters than data points. Next lecture asks why they work anyway, and it closes this unit.