Gradient Signal Propagation Through Depth at Initialization

The variance recursion that fixes an initializer is a statement about a single scalar per layer. It says nothing about the object that actually governs whether a deep network trains: the Jacobian of the whole network with respect to the input, and its counterpart with respect to the output, both of which are products of one matrix per layer. A product of matrices with well-behaved individual statistics can still have an ill-conditioned spectrum, and the analysis of why has produced the most useful statements available about depth.

The central object is a product $J = J_L J_{L-1}\cdots J_1$ where $J_\ell$ is the Jacobian of layer $\ell$‘s map. Backpropagated gradients at the input are $J^\top \cdot (\text{output gradient})$, so the loss’s sensitivity to early layers passes through $J^\top$, and forward activations pass through $J$.

Singular values of a product, not a variance

Each $J_\ell$ has a singular value distribution. The product’s singular values are not the products of the corresponding singular values when the matrices’ singular vectors are misaligned, and in a randomly initialised network they are indeed misaligned at every layer. The relevant statement, due to the multiplicative free convolution and its practical use in this context, is that the product’s squared singular values concentrate in a spectrum whose width depends on the ratio of the matrices’ dimensions and grows with depth. Two facts follow.

The first is that the typical singular value of the product can remain near one while the extreme ones diverge. Width matters here: a wide matrix’s singular values are tightly concentrated around their mean, so the product of wide matrices has a narrower spectrum than the product of narrow ones. This is the mathematical form of the observation that increasing width improves trainability at fixed depth, and decreasing width at fixed depth degrades it.

The second is that the log-singular values accumulate additively. If each layer’s log-spectrum has a spread, the depth-multiplied log-spectrum’s spread grows as the square root of depth while the mean grows linearly, so the mean’s deviation from zero dominates. A systematic gain slightly above one, even by a factor of $1 + \delta$ per layer, gives an alignment factor $e^{\delta L}$, which for $\delta = 0.01$ and $L = 100$ is $e$. This is why initialization constants are not approximate: the recursion’s assumption of exact variance preservation is the difference between exponential growth and a bounded spectrum.

The rectifier’s mean shift

There is a further effect that the variance recursion cannot see. A rectifier maps a symmetric distribution to one with positive mean, so the pre-activation distribution at layer $\ell+1$ is no longer symmetric: its mean is roughly $\sigma_\ell \sqrt{2/\pi}$ when the weights are zero-mean. The weights have zero mean, so this shift does not accumulate through the linear map’s mean — but the deviation from that mean passes through the rectifier with a different factor than the variance does, and the second moment is the sum of the squared mean and the variance. An analysis that tracks only the variance therefore overstates the correction needed as depth grows; the correct fixed point requires the weights’ variance to be chosen so that the second moment, not the variance, is preserved, and that fixed point depends on the weight distribution’s shape.

This is the substance of the mean-field treatment of rectified stacks: the propagation has a controlled fixed point only when the weights’ variance is near a critical value, and the approach to that fixed point can be ordered when it is. At the critical value, the correlations induced between two inputs decay with depth in a way that depends on their initial angle, so the network’s behaviour at initialization is a deterministic function of depth and of the initial correlation structure. Away from the critical value, either all inputs collapse to the same representation (ordered phase) or the representation of two nearby inputs decorrelates chaotically (chaotic phase).

Residual stacks change the product’s structure

A residual block’s Jacobian is $I + J_{F_\ell}$, so a stack’s Jacobian is a product of $\prod_\ell (I + J_{F_\ell})$. Expanding the product gives a sum of terms: the identity, all single-block contributions, all products of two, and so on. The identity term alone has singular values exactly one, so the product’s singular values are bounded below by those of the sum’s leading term, and the growth is governed by the sum of the branch Jacobians rather than their product.

Two consequences are immediate. The alignment problem disappears at leading order, because the identity is perfectly aligned with itself; the path through any subset of blocks contributes additively, and the number of paths of length $k$ out of $L$ is $\binom{L}{k}$, which is large for moderate $k$. The gradient’s norm through a residual stack with a unit-scale identity path therefore does not decay exponentially in depth at initialization, regardless of how small the branch contributions are, as long as they are not zero.

The variance recursion’s growth is the other side: the forward signal’s variance grows as $1 + Lq$ rather than $q^L$, which is why the pre-normalisation arrangement described in layer normalisation and the residual stream can be trained at depths where a plain stack cannot.

The critical point comes from the branch. If the branch’s Jacobian has spectral radius $\rho$, the block’s Jacobian has singular values near $1 \pm \rho$ in a basis aligned with the branch’s, so blocks with $\rho$ large amplify the stream at each step and accumulating variance grows faster than linearly. The branch’s own initialization must keep $\rho$ below one, which is exactly the content of the scaling of the branch’s output by a small constant: it makes the residual’s variance growth geometric rather than arithmetic.

import torch
from torch import nn

def spectral_radius(w):
    # largest singular value; torch.linalg.svdvals on a 2-D tensor
    return float(torch.linalg.svdvals(w)[0])

lin = nn.Linear(256, 256, bias=False)
nn.init.normal_(lin.weight, std=(2.0 / 256) ** 0.5)
print(round(spectral_radius(lin.weight), 3))

# a residual stack applies (I + W) rather than W; its radius is bounded by
# 1 + rho(W), so rho(W) < 1 is what keeps the stream from growing per layer

The printed value is the branch’s radius; the block’s is bounded by $1 + \rho$, and the bound is tight when the branch is linear.

Depth and width as separate budgets

The mean-field analysis produces a criterion that can be stated without reference to a particular architecture: the layer’s variance parameter must be set so that the propagation of the input’s norm is neutral, and the neutral point depends on the nonlinearity and on the distribution of the inputs. For a rectifier the neutral weight variance scales as $2/n$ for a forward pass, but a backward pass through the same rectifier sees a different effective factor because the derivative is zero on half the inputs, which halves the number of contributing paths. The asymmetry means a forward-neutral choice is not backward-neutral, and the observable symptom of choosing the forward criterion is that deep stacks train at first and then stop improving as depth grows, with gradients at early layers small but nonzero.

The width’s role is independent: it controls the concentration of the spectrum. An analysis of random matrices with independent entries of variance $\sigma^2$ and aspect ratio $c = n_{\text{in}}/n_{\text{out}}$ gives a support for the singular value distribution that narrows as the matrix’s dimensions grow, so a wide layer’s Jacobian has a spectrum concentrated near its mean and its product is better conditioned. Narrow layers have wider spectra, and a stack of narrow layers has a product whose extremes are far from one even when the mean is right. The practical form of this: at fixed parameter count, depth costs conditioning that width supplies.

Where the analysis stops predicting

The propagation theory is a statement about the network at initialization, and its predictions are about the first gradient. Training changes the weights, and no argument based on the initial spectrum survives to the weights after many steps. What the theory reliably predicts is whether training can start: whether the first steps have a well-conditioned direction and a gradient magnitude in a usable range. Networks whose initialization lies near the critical point in the sense above can be trained at depths that networks away from it cannot, which is the observed content of the analysis.

A second limit is that the theory assumes the network is a composition of layers with independent weights and that the nonlinearity acts elementwise. Attention mixes positions, so the Jacobian of an attention block is not a product of independent random matrices across positions and the spectra are structured rather than isotropic. Normalisation layers multiply the Jacobian by a projection that removes the components along the mean and along the current activation, which changes the spectrum’s extremes without changing its mean — a rank-one and a rank-two perturbation per layer. The theory still describes the scale correctly and the conditioning less so, and the correct response to an ill-conditioned stack is to change the architecture’s structure rather than to search for a better scalar constant.

A third limit is arithmetic. The analysis is a mean-field one: it tracks the expected propagation of norms and correlations and ignores fluctuations around that expectation. Fluctuations do not stay small in a deep stack — the singular value spectrum of a product of many independent matrices widens with depth even when each factor’s spectrum is narrow, so the expected behaviour and a single realised network part company as depth grows. This is why two runs from different random seeds at large depth differ in whether they train, and why the seed becomes a hyperparameter that no initializer constant can remove.