Making deep networks train, from initialization and normalization to the tuning loop

Why this subject matters now

For the first decade of the deep learning era, getting a network deeper than a few layers to converge was a research contribution in itself. Unsupervised pretraining was invented in 2006 largely because direct training failed, and it was abandoned within a few years once the real culprits were identified, initializations that shrank or amplified signal exponentially with depth, saturating activations that silently zeroed gradients, and optimizers whose step sizes had no relationship to the geometry of the loss. The fixes, careful variance-preserving initialization, rectified and smooth activations, normalization layers, residual connections, and adaptive optimizers, are individually small ideas, but together they are the reason a thousand-layer network is trainable today and the reason "just stack more blocks" became a viable research strategy.

The subject has not become less important now that libraries ship good defaults. It has become the difference between practitioners who can act and practitioners who can only rerun. The defaults encode assumptions, ReLU-family activations, roughly unit-variance inputs, moderate depth, a particular optimizer, and the moment a project leaves that envelope (a new architecture, a new loss, a new data distribution, a precision change, a much larger batch) training failures reappear wearing the same three or four disguises they always wear. A loss stuck at exactly \( \ln(\text{num classes}) \), a gradient norm of \(10^{-18}\) in the first layer, an update that diverges on step 40 of warmup. Each of these has a short list of causes, and the practitioner who has derived the underlying mechanics can usually name the cause from the symptom in minutes. Interviews at strong labs test exactly this, not whether a candidate can call an optimizer but whether they know why Adam needs bias correction, why L2 regularization and weight decay are different things under adaptive methods, and what to check first when a network refuses to learn. Everything below is aimed at that level of fluency, and every claim that can be checked by running code has been checked on the H100 in this repository, with the measured numbers quoted.

Backpropagation, derived properly

The computational graph and the adjoint

A neural network's forward pass is a composition of primitive operations, and the natural data structure for reasoning about its derivatives is the computational graph, a directed acyclic graph whose nodes are intermediate values and whose edges record which values each operation consumed. The forward pass traverses the graph in topological order computing values. Differentiation traverses it once more, and the entire subject of automatic differentiation is the question of which direction to traverse.

Fix a scalar loss \( \L \) at the output. For every intermediate value \( v \) in the graph, define its adjoint \( \bar{v} = \partial \L / \partial v \), an object with the same shape as \( v \) holding the sensitivity of the loss to each of its entries. The chain rule, written in adjoint form, says that a node's adjoint is assembled from the adjoints of its consumers,

Reverse-mode automatic differentiation, of which backpropagation is the special case for neural networks, is nothing more than evaluating this recurrence in reverse topological order, starting from \( \bar{\L} = 1 \). The sum over consumers is why a value used in two places (a weight shared across time steps, an input feeding both a residual branch and a skip path) accumulates gradient from both. The graph structure does the bookkeeping that hand-derived gradients get wrong.

The vector-Jacobian product

The expression \( (\partial v / \partial u)^\top \bar{v} \) is a vector-Jacobian product, and the crucial engineering fact is that it is computed without ever forming the Jacobian. For a primitive \( f : \R^n \to \R^m \), the Jacobian \( J \in \R^{m \times n} \) can be enormous. For a linear layer \( Y = XW \) with batch 64 and width 4096, \( Y \) has \(64 \times 4096 \approx 2.6 \times 10^5\) entries and \( W \) has \(1.7 \times 10^7\), so the Jacobian of \( Y \) with respect to \( W \) has about \(4.4 \times 10^{12}\) entries. Nobody stores that. Instead each primitive ships a rule that maps an output adjoint directly to input adjoints. For the matrix multiply \( Y = XW \), write the forward in indices, \( y_{ij} = \sum_k x_{ik} w_{kj} \), and differentiate,

Both rules are themselves matrix multiplies of the same size as the forward one, which gives the standard FLOP accounting for training. The backward pass of a linear layer costs two matmuls against the forward's one, so a full training step costs roughly three times the forward pass. Every framework is organized around this idea. PyTorch's rules live in a single declarative table, tools/autograd/derivatives.yaml, one VJP formula per operation. JAX exposes the machinery directly as jax.vjp, and its grad is a thin wrapper that calls the VJP with cotangent 1.0.

Reverse versus forward mode, and the memory bill

Forward-mode differentiation propagates a directional derivative (a Jacobian-vector product) alongside the forward pass. Seed a direction \( \dot{\theta} \) in parameter space, and one augmented forward pass yields \( J \dot{\theta} \), the derivative of every output along that one direction. It needs no stored activations, but recovering the full gradient of a scalar loss with respect to \( P \) parameters requires \( P \) passes, one per basis direction. Reverse mode inverts the tradeoff. One backward pass recovers the gradient with respect to every parameter simultaneously, because the loss is a single scalar and the recurrence fans sensitivities out from it. For a network with a billion parameters that is one pass instead of a billion, which is the entire reason deep learning trains in reverse mode.

The price is memory. The VJP for \( Y = XW \) needs \( X \), so the input of every layer must be kept alive from the moment the forward pass computes it until the backward pass consumes it. Activation memory therefore scales with depth times the size of each layer's activations, and for large models it, not the parameters, is what fills the GPU. The standard escape, activation checkpointing, drops most activations and recomputes them during the backward pass, trading roughly one extra forward pass of compute for an order-of-magnitude reduction in activation memory. The arithmetic is worked in a problem later on this page. The rule of thumb worth memorizing is that reverse mode costs a small constant times the forward FLOPs regardless of parameter count, and its memory cost is the thing all the systems tricks exist to manage.

forward mode (JVP)                    reverse mode (VJP)
─────────────────                     ──────────────────
one pass per input direction          one pass for the whole gradient
P passes for P parameters             ~2x forward FLOPs, once
no activation storage                 stores every layer input
cheap for few inputs, many outputs    cheap for many inputs, one output
(used: Jacobian columns, HVPs)        (used: training every deep net)

Signal propagation and initialization

The variance of a linear layer

Whether a deep network is trainable at step zero is a question about how the scale of activations and gradients evolves with depth, and the whole analysis rests on one computation. Take a single output of a linear layer, \( y_i = \sum_{j=1}^{n_{\text{in}}} w_{ij} x_j \), with the weights drawn independently with mean zero and variance \( \sigma_w^2 \), independent of the inputs, and the inputs zero-mean with variance \( \sigma_x^2 \). Independence kills the cross terms in the expansion of \( \E[y_i^2] \),

One layer multiplies the variance of the signal by \( n_{\text{in}} \sigma_w^2 \). Stack \( L \) such layers with an (approximately) identity activation and the output variance is the input variance times \( (n \sigma_w^2)^L \), an exponential in depth whose base is set entirely by the initialization. The base must be almost exactly 1 or the network is numerically dead or exploding before the first gradient step. Glorot and Bengio (2010) made this argument for the forward pass, then repeated it for the backward pass. The gradient entering a layer is \( \bar{x} = W^\top \bar{y} \), the same computation with \( n_{\text{out}} \) playing the role of the fan, so backward variance is preserved when \( n_{\text{out}} \sigma_w^2 = 1 \). The two conditions conflict whenever \( n_{\text{in}} \neq n_{\text{out}} \), so they proposed the harmonic compromise now called Xavier or Glorot initialization,

The uniform bound follows from \( \Var(U[-a,a]) = a^2/3 \). Setting \( a^2/3 = 2/(n_{\text{in}}+n_{\text{out}}) \) gives \( a = \sqrt{6/(n_{\text{in}}+n_{\text{out}})} \). The analysis assumed the activation is roughly linear around zero, which is true for tanh at small inputs and false for ReLU everywhere.

The He condition and where the factor of 2 comes from

A ReLU zeroes half of a symmetric input distribution, and that factor of one half must appear in the variance bookkeeping. Let \( z \) be a pre-activation, symmetric about zero (true at initialization because the weights are symmetric), and let \( x = \max(0, z) \) be the post-activation that feeds the next layer. The next layer's variance formula needs the second moment of \( x \), and symmetry computes it directly,

Half the second moment survives rectification. (Note the derivation tracks second moments rather than variances, since \( x \) is not zero-mean but the next layer's weights are, so the cross terms still vanish and second moments are the right currency.) Propagating through the next linear layer, \( \Var(z^{(l+1)}) = n \sigma_w^2 \cdot \tfrac{1}{2} \E[(z^{(l)})^2] \), so variance is preserved through a linear-plus-ReLU pair exactly when

which is He initialization (He, Zhang, Ren, and Sun, 2015). The factor of 2 is not a fudge. It is the reciprocal of the fraction of the distribution a ReLU keeps. The same paper reports the consequence of using the Glorot value instead. Their 30-layer ReLU network simply failed to start converging, because the per-layer factor was \( \tfrac{1}{2} \) rather than 1.

Getting it wrong, measured at depth 50

The exponential is easy to state and more convincing to measure. The compounding factor for a ReLU stack is \( c = \tfrac{1}{2} n \sigma_w^2 \) per layer, so after \( L \) layers the standard deviation scales by \( c^{L/2} \). At width \( n = 512 \), Xavier-style \( \sigma_w^2 = 1/n \) gives \( c = 1/2 \) and a predicted factor of \( 2^{-25} \approx 3.0 \times 10^{-8} \) at depth 50. Doubling the He standard deviation gives \( c = 4 \) and a predicted factor of \( 2^{50} \approx 1.1 \times 10^{15} \). A naive \( \sigma_w = 0.01 \) gives \( c = 0.0256 \), which underflows float32 within about twenty layers. Running exactly this experiment on the H100 in this repository (a 50-layer, width-512 ReLU stack on unit-variance input, batch 1024) produced the table below. The agreement with the predictions is essentially exact, which is the point. This failure mode is arithmetic, not folklore.

The He row wanders (finite width makes each layer's factor a random variable with mean 1) but stays order-one for fifty layers. The Xavier row loses a factor of \( \sqrt{2} \) per layer like clockwork. And because the backward pass is the same linear algebra transposed, gradients suffer the same exponential. The training experiment later on this page measures a first-layer gradient norm of \(1.5 \times 10^{-18}\) for a badly initialized 20-layer network, against \(2.7\) for the He-initialized twin.

Orthogonal initialization and dynamical isometry

Variance arguments control the average scale of the signal, but averages can hide a spread. A layer can preserve variance overall while stretching some directions and crushing others. The sharper criterion is the spectrum of singular values of the input-to-output Jacobian. Saxe, McClelland, and Ganguli (2014) analyzed deep linear networks exactly and showed that random Gaussian layers, even variance-correct ones, compound into a Jacobian whose singular values spread out with depth, while orthogonal weight matrices (\( W^\top W = I \), every singular value exactly 1) compose into an orthogonal product, keeping the entire spectrum at 1 at any depth. They named the resulting property dynamical isometry, and showed deep linear networks with orthogonal initialization train in a number of steps independent of depth. Pennington, Schoenholz, and Ganguli (2017) extended the analysis to nonlinear networks with random matrix theory. Tanh networks can achieve approximate dynamical isometry with orthogonal weights, while ReLU networks cannot regardless of initialization (the rectifier's random diagonal mask spreads the spectrum), a rare clean theoretical separation between the two activation families. The practical descendant is that nn.init.orthogonal_ remains the default recommendation for RNN recurrence matrices, where the same matrix is applied hundreds of times and spectral discipline matters most, and delta-orthogonal initialization let Xiao et al. (2018) train 10,000-layer plain CNNs without residuals or normalization, as a proof that initialization alone can carry extreme depth.

Initialization for residual networks

Residual connections change the recursion from multiplicative to additive. With \( x_{l+1} = x_l + F_l(x_l) \), at initialization, where the branch is approximately independent of the stream, variances add, \( \Var(x_{l+1}) \approx \Var(x_l) + \Var(F_l(x_l)) \). If every branch is initialized to emit unit variance, the stream's variance grows linearly with the number of residual additions, and the activations entering late layers grow as \( \sqrt{\text{depth}} \). The modern fix scales the last linear layer of each residual branch so the branch's contribution is small. With \( 2L \) residual additions (a transformer with \( L \) blocks has two per block, attention and MLP), initializing each branch's output projection with weights scaled by \( 1/\sqrt{2L} \) makes each branch contribute variance \( \propto 1/(2L) \), and the total added variance across the whole stack sums to order one instead of order \( L \). This is exactly the GPT-2 initialization rule, and its relatives replace the scale with zero. Initializing the branch's final layer to zero (or the normalization gain \( \gamma \) to zero, the "zero-gamma" trick for batch-norm ResNets) makes every block the identity at step zero, which is the most trainable possible starting point, and schemes like Fixup and ReZero show that with such scaling residual networks train stably even with normalization removed. The unifying principle is that an untrained deep network should be as close to an identity (or at least an isometry) as its architecture allows, and initialization is the tool that puts it there.

Activation functions

Saturation, with the numbers

The sigmoid \( \sigma(z) = 1/(1+e^{-z}) \) has derivative \( \sigma'(z) = \sigma(z)(1 - \sigma(z)) \), maximized at \( z = 0 \) where it equals exactly \( 1/4 \). Two consequences, both fatal for depth. First, even at the best operating point, each sigmoid layer multiplies the backward signal by at most 0.25 (times the weight matrix), so ten stacked sigmoid layers attenuate gradients by up to \( 4^{-10} \approx 10^{-6} \) before the weights are even considered. Second, away from zero the derivative collapses. At \( z = 5 \), \( \sigma(5) = 0.9933 \) and \( \sigma'(5) = 0.9933 \times 0.0067 = 0.0066 \), and at \( z = 10 \) the derivative is \( 4.5 \times 10^{-5} \). A unit pushed into saturation by a bad batch or a large weight stops learning almost permanently, because the gradient that would pull it back is itself multiplied by the tiny derivative. The sigmoid has a further defect. Its outputs are all positive, so the gradient with respect to a layer's weights, which is an outer product with the all-positive input vector, has all entries sharing the sign of the upstream adjoint, forcing zig-zag updates. Tanh fixes the centering (it is \( 2\sigma(2z) - 1 \), zero-centered, maximum slope 1) but not the saturation, since \( \tanh'(z) = 1 - \tanh^2(z) \), and at \( z = 3 \) that is \( 1 - 0.99505^2 = 0.0099 \). These gradients are why pre-2011 networks deeper than a few layers were nearly untrainable.

ReLU and its failure mode

\( \mathrm{ReLU}(z) = \max(0, z) \) has derivative exactly 1 on the active half, so an active path through any number of ReLU layers passes gradient with no attenuation from the activation at all. This, more than any other single change, is what made depth 10-100 practical, and it costs one comparison. The price is the dead half, a derivative of exactly 0 for \( z < 0 \). A unit whose pre-activation is negative for every input in the data distribution contributes nothing forward and receives no gradient backward, permanently. Units die in practice when a large gradient step (often from a too-high learning rate) knocks a bias or weight row far negative. The diagnostic is to run a large batch through the network and count units that never activate, and the training-loop implementation later on this page does exactly that. Leaky ReLU (\( \max(\alpha z, z) \), typically \( \alpha = 0.01 \)) keeps a trickle of gradient on the negative side so dead units can recover, and PReLU (He et al., 2015) learns \( \alpha \) per channel. Both remove the death mode at negligible cost, and neither reliably improves final accuracy on healthy networks, which says the death mode, not the exact negative-side slope, was the issue.

Smooth activations, GELU, SiLU, and SwiGLU

The modern default in transformers is the GELU (Hendrycks and Gimpel, 2016), \( \mathrm{GELU}(z) = z\, \Phi(z) \) where \( \Phi \) is the standard normal CDF. This is the input times the probability a standard Gaussian falls below it, a smooth interpolation between killing the input (far negative) and passing it (far positive), with a shallow dip below zero (minimum about \(-0.17\) near \( z = -0.75 \)). The exact form uses \( \mathrm{erf} \). The widely used tanh approximation is \( 0.5\,z\,(1 + \tanh[\sqrt{2/\pi}\,(z + 0.044715 z^3)]) \). SiLU, also called Swish, is \( z\, \sigma(z) \), nearly indistinguishable from GELU in shape. The argument for smoothness is about optimization rather than expressiveness. The gradient of a smooth activation varies continuously with the input, so the loss surface has no creases from the activation, small negative inputs still receive gradient (no hard death), and the stochastic-regularizer derivation of GELU frames it as the expectation of randomly zeroing inputs weighted by their magnitude. SwiGLU (Shazeer, 2020) upgrades the transformer MLP rather than the pointwise function. Instead of \( \mathrm{act}(xW_1)W_2 \) it computes a gated product \( (\mathrm{SiLU}(xW_1) \odot xW_3)\, W_2 \), three matrices instead of two, with the hidden width scaled by \( 2/3 \) (typically to \( 8d/3 \)) to hold parameters constant. It is the FFN in Llama-family and PaLM-family models on the strength of consistent small perplexity gains.

The honest empirical picture deserves stating plainly. The jump from saturating activations to the ReLU family was a cliff, worth orders of magnitude in trainable depth. Within the modern family, ReLU versus GELU versus SiLU versus SwiGLU, differences are real but small, fractions of a percent of accuracy or a percent or two of perplexity, frequently within the noise of a seed change, and Shazeer's own SwiGLU paper declines to offer an explanation beyond the benchmarks ("divine benevolence," in its words). Activation choice is a second-order decision, while initialization, normalization, and learning rate are first-order. The one first-order rule is to match the initialization to the activation, since the He factor of 2 assumed an exact half-rectifier, and GELU/SiLU networks conventionally keep the same factor, the approximation error being absorbed by normalization.

Normalization

Batch normalization forward

Batch normalization (Ioffe and Szegedy, 2015) standardizes each feature over the current mini-batch and then gives the network back the freedom it just removed, through a learned scale and shift. For one feature across a batch of \( N \) values \( x_1, \dots, x_N \),

Every feature (channel) gets its own \( \mu, \sigma^2, \gamma, \beta \). In a convolutional network the statistics are taken over batch and spatial positions per channel. The \( \epsilon \) (typically \(10^{-5}\)) guards the division. Because \( \gamma \) and \( \beta \) can represent any affine map, the layer does not restrict what the network can compute. It restricts how the computation is parameterized, which turns out to be the entire point.

The batch normalization backward pass in full

The backward pass is the classic exercise because \( \mu \) and \( \sigma^2 \) are functions of every element of the batch, so each input \( x_i \) influences every output \( y_j \), and the naive chain rule sprouts terms in three directions. Given the upstream adjoints \( \bar{y}_i = \partial \L / \partial y_i \), the parameter gradients are immediate from \( y_i = \gamma \hat{x}_i + \beta \),

For \( \bar{x}_i \), treat \( \hat{x} \) as a function of three quantities, \( x_i \) directly, \( \mu \), and \( \sigma^2 \), and sum the three paths. Write \( s = (\sigma^2 + \epsilon)^{-1/2} \) for the inverse standard deviation. The path through \( \sigma^2 \) is

The path through \( \mu \) has two contributions, the explicit \( -\mu \) in the numerator of \( \hat{x} \) and the appearance of \( \mu \) inside \( \sigma^2 \). The second vanishes because \( \partial \sigma^2 / \partial \mu = -\tfrac{2}{N}\sum_i (x_i - \mu) = 0 \) (deviations from the mean sum to zero), leaving

Now assemble \( \bar{x}_i \) from its three paths, directly through the numerator (coefficient \( s \)), through \( \sigma^2 \) (which depends on \( x_i \) with derivative \( \tfrac{2}{N}(x_i - \mu) \)), and through \( \mu \) (with derivative \( \tfrac{1}{N} \)),

Substituting the two partials and using \( \hat{x}_i = s (x_i - \mu) \) to eliminate raw deviations in favor of \( \hat{x} \) gives

This closed form is the one to implement (the middle expression shows where each of the three terms came from). It has a geometric reading that explains a lot of batch norm's behavior. Up to the scale \( s \), the input gradient is the upstream gradient with its batch-mean removed (the \( \sum_j \bar{\hat{x}}_j \) term) and its component along \( \hat{x} \) removed (the last term). The gradient leaving a batch-norm layer is orthogonal to the all-ones direction and to the normalized activation itself. Shifts and rescalings of the batch, precisely the directions the forward pass is invariant to, are projected out of the gradient. The implementation section verifies this formula against autograd on the H100. The maximum absolute error in float64 is \( 1.8 \times 10^{-15} \) on the input gradient and exactly zero on \( \bar{\gamma}, \bar{\beta} \).

Train versus eval, and the batch-size dependence

Batch statistics are only available when there is a batch. At inference the layer must be a deterministic per-example function, so during training it maintains exponential moving averages, \( \mu_{\text{run}} \leftarrow (1 - m)\, \mu_{\text{run}} + m\, \mu_{\text{batch}} \) and likewise for the variance (PyTorch's momentum is this \( m \), default 0.1, and the running variance uses the unbiased \( N/(N{-}1) \) correction), and at evaluation it normalizes with the running statistics instead. This split is the source of the single most common bug in the vicinity, forgetting model.eval(), so evaluation batches are normalized by their own statistics (wrong, and batch-size dependent), or forgetting model.train() afterward, so training silently stops updating the running averages. It also means train-mode and eval-mode losses differ slightly even on the same data, which surprises people watching both curves.

The deeper limitation is statistical, since \( \mu \) and \( \sigma^2 \) are estimates from \( N \) samples, with variance \( O(1/N) \). At batch 256 the noise is negligible. At batch 2 it is enormous, and the network trains against normalization statistics that jitter wildly. Wu and He (2018) measured the effect directly. ResNet-50's ImageNet error with batch norm degrades from 23.6 percent at batch 32 to 27.3 percent at batch 4 and 34.7 percent at batch 2 per GPU, while their group norm holds essentially flat. Small per-device batches are common (detection and segmentation with high-resolution inputs, memory-constrained fine-tuning), which is the practical reason batch-independent normalizers exist. When batch norm must be used with small per-device batches, synchronized batch norm computes statistics across devices at the cost of a communication round per layer.

Layer norm, RMSNorm, group norm

Layer normalization (Ba, Kiros, and Hinton, 2016) rotates the axis of normalization. Statistics are computed per example across the features, \( \mu_i = \tfrac{1}{d}\sum_k x_{ik} \), rather than per feature across the batch. Nothing depends on the batch, so train and eval are identical, batch size 1 works, and there are no running statistics to desynchronize. This is why it, not batch norm, is the normalizer in transformers, where variable-length sequences and autoregressive inference make batch statistics awkward. RMSNorm (Zhang and Sennrich, 2019) deletes the re-centering. It computes \( y = \gamma \odot x / \mathrm{RMS}(x) \) with \( \mathrm{RMS}(x) = \sqrt{\tfrac{1}{d}\sum_k x_k^2 + \epsilon} \), no mean subtraction and no \( \beta \). Their experiments showed the re-centering contributes little to nothing to quality across machine translation and language modeling while the norm costs 7 to 64 percent of the layer's runtime in their measurements, and subsequent adoption has been the strongest kind of evidence, since Llama, T5, and most current LLMs use RMSNorm. Group normalization (Wu and He, 2018) is the convolutional compromise. It splits channels into groups (32 by default) and normalizes per example per group over channels and spatial positions, recovering layer norm at one group and instance norm at \( C \) groups. It is the standard choice in detection and segmentation backbones where per-device batches are small.

Why normalization actually helps

The original paper explained batch norm as reducing internal covariate shift. Each layer's input distribution moves as the layers below it update, and normalization pins the first two moments so each layer trains against a stationary target. Santurkar, Tsipras, Ilyas, and Madry (2018) tested the story directly and broke it. They trained a batch-norm network with time-varying, non-zero-mean, non-unit-variance noise injected after every normalization layer, reintroducing severe distributional shift downstream, and it trained essentially as well as clean batch norm and far better than no batch norm. The benefit survives the removal of its official explanation. Their alternative account, supported by both measurement and analysis, is optimization geometry. Batch norm smooths the loss landscape, in the sense that the loss and gradient obey much better effective Lipschitz bounds along the training trajectory, so gradients change less abruptly, larger learning rates are stable, and training tolerates worse initialization. This reparameterization view also explains observations the covariate story cannot. Normalization makes the forward pass invariant to the scale of the incoming weights (scaling \( W \) by \( c \) changes nothing after normalization, while the gradient scales by \( 1/c \), an automatic per-layer step-size equalizer), and the projection structure of the backward pass derived above removes exactly the gradient components that would fight the normalization. The honest summary is that normalization demonstrably enables higher learning rates, weaker sensitivity to initialization, and faster convergence. The smoothing account is the best-supported explanation, the covariate-shift account is historically important and empirically insufficient, and a complete first-principles theory is still open.

Residual connections

The identity path, derived

A residual block computes \( x_{l+1} = x_l + F_l(x_l) \). The layer learns a correction to its input rather than a replacement for it (He, Zhang, Ren, and Sun, 2016). The motivating observation was the degradation problem. Plain 56-layer networks had higher training error than 20-layer ones, which is not overfitting but an optimization failure, since the deeper network contains the shallower one (set the extra layers to identity) and gradient descent still could not find it. Making identity the zero point of the parameterization fixes the asymmetry. A residual block outputs identity when its weights are near zero, which is where initialization puts them.

The gradient argument is short enough to write completely. Unroll the recursion from layer \( l \) to the top layer \( L \),

and differentiate the loss with respect to \( x_l \),

The identity term means the top-layer gradient reaches layer \( l \) additively and untouched, no matter how deep the stack. For the total gradient to vanish, the sum of branch Jacobians would have to conspire to equal \( -I \), which does not happen for random or trained weights. Contrast the plain stack, where \( \partial x_L / \partial x_l = \prod_{i=l}^{L-1} J_i \) is a product of \( L - l \) Jacobians and vanishes or explodes geometrically unless every factor is spectrally near 1, which is exactly the condition initialization struggled to maintain. Residuals convert a product into (an identity plus) a sum, and sums are numerically kind.

Pre-activation ordering

The derivation above assumed the skip path is a pure identity. The original ResNet block was \( x_{l+1} = \mathrm{ReLU}(x_l + F(x_l)) \), which puts a nonlinearity on the trunk. The unrolled telescoping breaks, because every block's output passes through another ReLU before reaching the top. He et al.'s follow-up (2016) tested every ordering and found that moving batch norm and ReLU inside the branch, \( x_{l+1} = x_l + F(\mathrm{BN}, \mathrm{ReLU};\, x_l) \), the "pre-activation" block, leaves the trunk as a true identity from input to logits, and with it they trained a 1001-layer CIFAR ResNet that outperformed its 110-layer counterpart, where the post-activation version of the same depth diverged. Pre-norm transformer blocks are the same discovery in different clothes. Apply the norm on the branch input, keep the residual stream clean, and depth stops being the enemy. The one cost, documented in both literatures, is that a perfectly clean trunk lets late blocks be partially bypassed, which slightly weakens the final-loss ceiling. Post-norm can reach marginally better loss when it can be stabilized at all, which is why hybrid orderings occasionally resurface.

Ensembles of shallow paths

Veit, Wilber, and Belongie (2016) proposed reading a residual network not as one deep computation but as an implicit ensemble. Expand the product form of the whole network,

a sum over \( 2^L \) paths, where each path passes through the branches in \( S \) and skips the rest. Path lengths follow a binomial distribution centered at \( L/2 \). Their experiments back the reading. Deleting a single block from a trained 110-layer ResNet barely moves test error (deleting a layer from a plain VGG-style network destroys it), shuffling blocks degrades gracefully, and, most tellingly, measuring per-path gradient magnitude shows almost all gradient flows through paths of length 5 to 17 even in a 110-block network. The effective depth during training is far smaller than the nominal depth, and the long paths contribute negligible gradient. This resolves the apparent paradox of trainable thousand-layer networks. Nothing requires gradient to traverse a thousand nonlinear transformations. The architecture provides an exponential family of shallow, trainable paths, and deep paths come online, if at all, later. Stochastic depth, covered under regularization, turns this observation into a training method by literally sampling the ensemble.

Optimizers

SGD and momentum

Stochastic gradient descent updates \( \theta \leftarrow \theta - \alpha\, \hat{g} \), where \( \hat{g} \) is the gradient of the loss on a mini-batch, an unbiased but noisy estimate of the full gradient. Momentum (the heavy-ball method, Polyak 1964) replaces the raw gradient with an exponential moving average of past gradients. In the convention most frameworks use,

Two derivations give its character. First, the effective step size. Suppose the gradient is constant at \( g \) for many steps, as it approximately is along a slowly varying descent direction. The velocity converges to the fixed point of \( v = \beta v + g \), namely \( v_\infty = g / (1 - \beta) \), so the steady-state update is \( \alpha g / (1 - \beta) \). With \( \beta = 0.9 \) momentum amplifies persistent gradient directions tenfold, and with \( \beta = 0.99 \) a hundredfold. Raising \( \beta \) without cutting \( \alpha \) is therefore a tenfold learning-rate increase in disguise, a classic source of "momentum made it diverge." Second, the filtering view. Components of the gradient that alternate sign step to step (the across-the-ravine direction, and mini-batch noise) largely cancel in the moving average, while components that persist (the along-the-valley direction) accumulate. On an ill-conditioned quadratic this is worth more than a constant. Optimally tuned momentum improves the convergence factor from \( (\kappa - 1)/(\kappa + 1) \) to \( (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) \), a square-root reduction in the effective condition number. Nesterov's variant evaluates the gradient at the look-ahead point \( \theta + \beta v \) rather than at \( \theta \), which tightens the classical convergence guarantee. In deep learning practice the difference is usually minor, and SGD with plain momentum 0.9 remains the reference optimizer for convnets.

AdaGrad and RMSProp, per-coordinate scaling

A single global learning rate is wrong whenever coordinates have very different gradient scales, which is the normal situation. Embedding rows for rare tokens receive rare, large gradients, and biases and gains live on different scales than weight matrices. AdaGrad (Duchi, Hazan, and Singer, 2011) divides each coordinate's step by the root of its accumulated squared gradients, \( \theta_j \leftarrow \theta_j - \alpha\, g_j \big/ \sqrt{\textstyle\sum_{\tau \le t} g_{\tau,j}^2 + \epsilon} \), equalizing progress across coordinates and giving rare features proportionally larger steps. Its flaw for deep learning is that the accumulator only grows, so the effective step decays toward zero on a schedule set by history rather than by progress. RMSProp (Tieleman and Hinton, 2012, circulated unpublished) replaces the sum with an exponential moving average, \( v_t = \rho v_{t-1} + (1 - \rho) g_t^2 \), so the denominator tracks the recent gradient scale instead of the total.

Adam, with the bias correction derived

Adam (Kingma and Ba, 2015) combines both moving averages, a first moment for direction and a second for scale,

with defaults \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \), \( \epsilon = 10^{-8} \). The hats are the bias correction, and they are needed because the averages start from zero. Unroll the first-moment recursion,

for stationary gradients, by the geometric series. The average underestimates the true moment by exactly the factor \( 1 - \beta^t \), which is severe early (at \( t = 1 \), \( m_1 = 0.1\, g_1 \) is ten times too small) and vanishes as \( t \) grows. Dividing by \( 1 - \beta^t \) removes it exactly. The reason this matters is a ratio effect. The two moments warm up at different speeds. Without correction, the first update is

more than three times the intended step, taken at the moment the optimizer knows the least, and the mis-scaling persists for roughly \( 1/(1-\beta_2) = 1000 \) steps as the second moment slowly fills. With correction, the first step is exactly \( \alpha\, \mathrm{sign}(g_1) \). This early-variance problem is also the deepest reason warmup helps Adam, discussed below. Note what the update has become, a signal-to-noise ratio. Coordinates whose gradients are consistent (\( \hat{m} \approx \sqrt{\hat{v}} \)) move at nearly \( \alpha \), while coordinates whose gradients thrash move less. Adam is approximately scale-invariant to the loss (multiply \( \L \) by 10 and nothing changes), which SGD is not. Reddi, Kale, and Kumar (2018) showed the EMA second moment breaks Adam's convergence proof and constructed convex counterexamples where Adam diverges because rare, informative large gradients are forgotten too quickly. Their AMSGrad fix (a running max on \( \hat v \)) repairs the theory but rarely changes deep learning practice, where the working fixes are warmup, a larger \( \epsilon \), and gradient clipping.

AdamW, and why L2 and weight decay are different things

Under plain SGD, penalizing the loss with \( \tfrac{\lambda}{2} \lVert \theta \rVert^2 \) and decaying the weights are identical. The gradient of the penalty is \( \lambda \theta \), so the update is \( \theta \leftarrow \theta - \alpha g - \alpha \lambda \theta = (1 - \alpha\lambda)\, \theta - \alpha g \), a multiplicative shrink plus the usual step. Under any adaptive method the equivalence breaks, and the proof is one line. With L2, the penalty gradient enters the moment estimates, and the update direction becomes

so the shrink applied to each coordinate is divided by that coordinate's gradient scale. Weights with large, active gradients are barely regularized, and weights with quiet gradients are regularized hardest. No choice of \( \lambda \) makes this equal a uniform decay unless \( \hat{v} \) is the same in every coordinate, which defeats the point of an adaptive method. Decoupled weight decay (Loshchilov and Hutter, 2019) instead applies the shrink outside the adaptive machinery,

restoring uniform, optimizer-independent regularization. Their experiments showed it decouples the best \( \lambda \) from the best \( \alpha \) (with L2 the two must be tuned jointly) and improves generalization, and AdamW is now the default optimizer for transformers essentially everywhere. The practical corollary is that weight_decay in torch.optim.Adam is L2 (coupled), while in torch.optim.AdamW it is decoupled. They are different regularizers and their good values differ.

Gradient clipping and the newer optimizers

Clipping by global norm, \( g \leftarrow g \cdot \min\big(1,\, c / \lVert g \rVert \big) \), leaves typical steps untouched and caps the rare pathological ones, a bad batch, an attention logit blowup, or a spike after a data distribution seam. It was introduced for exploding gradients in RNNs and survives as cheap insurance in nearly every large training run (a common choice is \( c = 1.0 \) with the threshold sanity-checked against the observed gradient-norm distribution so that clipping is rare rather than constant, since a run that clips every step is a run whose learning rate is wrong). Monitoring the clip rate is a diagnostic in itself.

The post-Adam optimizer field is active and deserves an honest reading. Lion (Chen et al., 2023), found by evolutionary program search at Google, updates with the sign of an interpolated momentum, halving optimizer memory. It matched or beat AdamW across vision and some language tasks in the original paper, and has been mixed in independent LLM-scale replications. Sophia (Liu, Li, Hall, Liang, and Ma, 2023) preconditions with a cheap diagonal Hessian estimate refreshed every few steps and reported roughly 2× fewer steps than AdamW on GPT-2-scale pretraining. Independent re-evaluations under matched, tuned baselines (Kaddour et al., 2023) found much of the advantage evaporates, which is a recurring pattern in optimizer papers, an under-tuned AdamW baseline (see also Schmidt et al., 2021, who benchmarked fifteen optimizers and found no consistent winner over well-tuned classics). Muon (Jordan et al., 2024) is the current strongest challenger. It applies momentum, then orthogonalizes each hidden weight matrix's update via a few Newton-Schulz iterations (approximately replacing the update with the nearest semi-orthogonal matrix), with AdamW retained for embeddings and scalars. It set community speed records on small GPT training, and Moonshot AI (Liu et al., 2025) reported roughly 2× computational efficiency versus AdamW at LLM scale and trained the trillion-parameter Kimi K2 with a clipped variant, the first frontier-scale model trained on a non-Adam optimizer in years. The prudent summary is that AdamW remains the default, Muon has the most credible at-scale evidence among challengers, and any claimed optimizer win should be discounted until the baseline's learning rate, schedule, and weight decay were retuned with equal care.

The learning rate

Why it dominates every other hyperparameter

The learning rate is the only hyperparameter that multiplies every update of every parameter on every step, and its usable range is bounded on both sides by failure. Too low, and training is slow and can settle into worse minima reachable only with large early steps. Too high, and the iterates diverge or bounce at a high loss floor. On a quadratic with curvature \( \lambda \) along an eigendirection, gradient descent converges only if \( \alpha < 2/\lambda \), so the largest curvature bounds the stable learning rate, and (a measured surprise of the last few years, Cohen et al., 2021) deep networks trained at constant \( \alpha \) drive their sharpness up to sit at that \( 2/\alpha \) boundary, the "edge of stability," rather than below it. Practically, the loss-versus-learning-rate curve on a log axis is a U with a cliff on the right, the best value sits uncomfortably close to the cliff, and no other hyperparameter's curve looks like this. Tuning budgets should be spent on the learning rate first, on a log grid, before anything else is touched.

Warmup, and its interaction with Adam

Warmup ramps the learning rate linearly (or geometrically) from near zero over the first few hundred to few thousand steps. Three independent reasons converge on it. For adaptive optimizers, the second-moment estimate \( \hat{v} \) is computed from a handful of samples early on, so the per-coordinate step sizes are high-variance precisely when the parameters are least trained, and a small \( \alpha \) during that window caps the damage (this is the analysis behind RAdam, which derives a correction term and shows warmup approximates it). For large-batch SGD, the linear-scaling prescription below calls for learning rates that are unstable from a cold start but fine after the network leaves its initialization. Goyal et al. (2017) introduced the 5-epoch gradual warmup for exactly this. And in post-norm transformers, early gradient magnitudes vary sharply across depth, and warmup was historically what kept them alive. Warmup costs little, fixes real failure modes, and there is no reason to skip it in a modern run.

Schedules, cosine and WSD

After warmup, the standard schedule for years has been cosine decay (Loshchilov and Hutter's SGDR, 2017, minus the restarts),

a smooth anneal spending most of the budget near the peak and slowing gently into the end, typically with \( \eta_{\min} \) set to a tenth or a hundredth of the peak. Its inconvenient property is that \( T \) appears in the formula. The entire trajectory depends on knowing the total step budget in advance, and a checkpoint from the middle of a cosine run is not a good model (its learning rate never annealed) nor a good starting point for extension. The warmup-stable-decay (WSD) schedule answers this. Hold the learning rate constant after warmup for the bulk of training, then decay sharply (linearly or exponentially) over the final 10 to 20 percent. Hu et al. (MiniCPM, 2024) popularized it for LLMs, and Hägele et al. (2024) showed systematically that WSD matches cosine at equal budget while allowing the decay to be branched off any checkpoint of the stable phase, so one long run yields models at many budgets and scaling-law experiments get dramatically cheaper. Most of the benefit of any schedule is concentrated in the final decay. The loss drops substantially during that phase as the iterates stop bouncing and settle into the local basin.

Batch size, the linear-scaling rule, and gradient noise

Batch size and learning rate are not independent knobs. The classical argument (Goyal et al., 2017) runs as follows. \( k \) consecutive small-batch steps with learning rate \( \alpha \) move the parameters by roughly the sum of \( k \) gradients evaluated at nearby points, and one large-batch step on the concatenated batch with learning rate \( k\alpha \) moves by the same amount if the gradients change slowly over the window. Hence the linear scaling rule. Multiply batch size by \( k \), multiply the learning rate by \( k \), and expect a nearly unchanged training trajectory per epoch. With warmup added it carried ResNet-50 to batch 8192 with no accuracy loss in their one-hour ImageNet run, and it breaks down at some model-dependent scale where the "gradients change slowly" premise fails.

The principled version of "some scale" is the gradient noise scale (McCandlish, Kaplan, Amodei, and Team, 2018). Model the mini-batch gradient as the true gradient plus zero-mean noise with covariance \( \Sigma / B \). The batch size at which the noise stops dominating, the critical batch size, is approximately

the ratio of total gradient variance (per example) to squared true-gradient norm. Below \( B_{\text{crit}} \), gradient noise dominates, and doubling the batch nearly halves the number of steps needed (perfect data parallelism). Far above it, the gradient estimate is already accurate and extra batch is almost pure waste. The noise scale is measurable during training (compare gradient norms at two batch sizes), it varies over orders of magnitude across tasks, and it grows as the loss falls, which is why large runs ramp batch size over training. This one quantity organizes the whole batch-size playbook. Measure it, train near it, scale the learning rate with batch below it, and do not expect miracles above it.

The learning-rate range test

Smith (2017) proposed the cheapest useful experiment in the field, one short run in which the learning rate ramps geometrically from far too small to far too large while loss is recorded per step. The smoothed loss traces a characteristic shape, flat where \( \alpha \) is too small to move, descending through the useful range, then blowing up past the stability cliff. The working prescription is to pick a value several-fold below the cliff, near where descent is steepest. Running it on the 20-layer MLP used throughout this page (SGD momentum 0.9, 300 steps, \(10^{-6}\) to 10) gave a smoothed loss flat near 6.5 below \( 10^{-5} \), a steep descent through \(10^{-4}\) to \(10^{-3}\), a minimum of 0.002 at \( \alpha = 8.8 \times 10^{-3} \), and divergence past \( 1.6 \times 10^{-2} \), with the smoothed loss at \( 4.3 \times 10^5 \) by \( \alpha = 5.5 \times 10^{-2} \). The cliff sits within a factor of two of the best value, which is typical and is the reason to run the test rather than guess. The follow-on observation (super-convergence, Smith and Topin, 2018) is that a cyclical or one-cycle schedule peaking near the cliff can train some networks in far fewer steps than a conservative constant rate.

Regularization

L2 and its geometry

The L2 penalty \( \tfrac{\lambda}{2} \lVert \theta \rVert^2 \) has a precise geometric action that a quadratic approximation makes visible. Near a minimum \( \theta^* \) of the unpenalized loss with Hessian \( H = Q \Lambda Q^\top \), the penalized minimum is \( \tilde{\theta} = (H + \lambda I)^{-1} H\, \theta^* \). In the Hessian eigenbasis, the component of \( \theta^* \) along the eigenvector with curvature \( \lambda_i \) is scaled by \( \lambda_i / (\lambda_i + \lambda) \). Directions in which the loss is sharply curved (\( \lambda_i \gg \lambda \)), where the data strongly pins the parameter, are almost untouched. Flat directions (\( \lambda_i \ll \lambda \)), where the data barely cares, are shrunk toward zero. L2 is a filter that discards parameter components the loss cannot defend. This is also why its interaction with normalization is odd. For weights feeding a normalization layer the loss is exactly scale-invariant, the "shrink" direction is a flat direction of the loss, and weight decay's real effect there is to control the effective learning rate rather than the function.

Dropout, the ensemble view and inference-time scaling

Dropout (Srivastava, Hinton, Krizhevsky, Sutskever, and Salakhutdinov, 2014) multiplies each activation by an independent Bernoulli mask during training, keeping each with probability \( q = 1 - p \) and zeroing it otherwise. The inference-time question is what deterministic network best represents the distribution of trained subnetworks, and expectation matching answers it. A unit that was present with probability \( q \) during training contributes \( \E[m\, h] = q\, h \) on average, so at inference its outgoing weights should be scaled by \( q \) (or, the modern "inverted" convention, activations are divided by \( q \) during training so inference needs no change at all, \( \E[\tfrac{m}{q} h] = h \)). The ensemble reading is that a network with \( n \) droppable units defines \( 2^n \) weight-sharing subnetworks, each mini-batch trains one of them, and the scaled inference network approximates the ensemble's geometric-mean prediction, exactly so for a single softmax layer (shown in the original paper). Gal and Ghahramani (2016) gave the Bayesian reading. Dropout training approximates variational inference over the weights, and keeping dropout active at test time (Monte Carlo dropout) yields cheap, imperfect uncertainty estimates. As matters stand, it is essential in the small-data regimes it was invented for, largely displaced in convnets by batch norm and augmentation, and used in transformers at modest rates (0 to 0.1) with pretraining on web-scale data often using none at all, because the strongest regularizer is data the model has not seen twice.

Stochastic depth and label smoothing

Stochastic depth (Huang, Sun, Liu, Sedra, and Weinberger, 2016) applies the dropout idea to whole residual branches, \( x_{l+1} = x_l + b_l\, F_l(x_l) \) with \( b_l \sim \mathrm{Bernoulli}(p_l) \), survival probability ramping down linearly with depth (their recommendation ramps from 1 at the input to 0.5 at the top). Training samples shallow members of Veit's path ensemble, shortening expected depth and gradient paths. At inference each branch is scaled by \( p_l \), the same expectation-matching argument as dropout. It let the authors train a 1202-layer CIFAR ResNet to record accuracy, and it survives today as "drop path" in vision transformer recipes, where it is among the few regularizers that reliably matter.

Label smoothing (introduced in Szegedy et al., 2016) replaces the one-hot target with \( q = (1 - \varepsilon)\, \mathbf{1}_{y} + \varepsilon / K \), and its effect on the optimum is derivable in closed form. Cross-entropy against \( q \) is minimized when the softmax equals \( q \), so consider symmetric logits with the correct class at \( z_c \) and the rest at \( z_o \). The optimum requires

Against hard targets the same ratio is infinite. The loss keeps paying for ever-larger logit gaps, driving weight growth and overconfidence forever. Smoothing caps the demanded gap at a finite number, which is its real mechanism. Müller, Kornblith, and Hinton (2019) mapped the consequences, better calibration and often slightly better accuracy, but visibly collapsed intra-class structure in the penultimate layer (correct-class representations get pulled into tight, equidistant clusters), which measurably hurts knowledge distillation from a smoothed teacher, since the "dark knowledge" in the relative wrong-class logits is exactly what smoothing erased.

Mixup, CutMix, and augmentation

Data augmentation is regularization by encoding invariances, such as flips, crops, and color jitter in vision. The model sees the label survive transformations and must learn features that survive them too. Mixup (Zhang, Cissé, Dauphin, and Lopez-Paz, 2018) abandons realism. Training on convex combinations \( \tilde{x} = \lambda x_i + (1-\lambda) x_j \), \( \tilde{y} = \lambda y_i + (1-\lambda) y_j \) with \( \lambda \sim \mathrm{Beta}(\alpha, \alpha) \) teaches the model to behave linearly between training points (vicinal risk minimization), smoothing decision boundaries and improving calibration and robustness to label noise. It is worth about a point of ImageNet top-1 on ResNet-50 in the original paper. CutMix (Yun et al., 2019) makes the combination spatial, pasting a rectangular patch of one image onto another with the label mixed by patch area, keeping locally realistic statistics. Modern vision recipes stack these (RandAugment, mixup, CutMix, random erasing) and the stack is collectively worth several points. In NLP pretraining, by contrast, augmentation plays a minor role and scale of unique data does the regularizing.

Early stopping

Early stopping, halting when validation loss stops improving, is implicit regularization with a clean linear-model account. For a quadratic loss trained by gradient descent from the origin, the parameter after \( \tau \) steps at rate \( \alpha \) has each Hessian eigencomponent filled in proportion to \( 1 - (1 - \alpha \lambda_i)^\tau \), so high-curvature (well-determined) directions converge first and flat (noise-dominated) directions are still near zero when training halts. Stopping at time \( \tau \) acts like L2 with \( \lambda \sim 1/(\alpha \tau) \) (the correspondence is made exact in Goodfellow, Bengio, and Courville, chapter 7). It costs one held-out set and a patience counter, composes with everything else, and its byproduct, the gap between training and validation curves over time, is the primary diagnostic signal of the next section.

Diagnosing a training run

Check the initial loss against theory

A classifier over \( K \) classes that knows nothing should emit the uniform distribution, and its cross-entropy is exactly \( -\ln(1/K) = \ln K \), which is 2.3026 for ten classes, 4.6052 for a hundred, and about 11.0 for a 60k-token vocabulary. Comparing the measured step-zero loss against this number is a free test that catches a large fraction of real bugs, and it ran as promised on this page's 20-layer MLP. With the final layer initialized small, the measured initial loss was 2.3026 to five figures, while He initialization all the way to the logits gave 6.83, because variance-preserving initialization produces logits of standard deviation well above 1 and the softmax starts out confidently wrong. (That is why zero- or small-initializing the final classifier layer is standard. It starts training from the uniform prediction.) Initial loss far above \( \ln K \) means logits are too large (initialization, missing normalization). Far below means the task is leaking labels, and \( \ln K \) with no subsequent movement means gradient is not reaching the weights.

Overfit a single batch

A correct architecture, loss, and optimizer must be able to drive the loss to approximately zero on one small batch. Memorizing 64 examples requires no generalization, only a functioning gradient path. Zhang et al. (2017) demonstrated the strong form of this, standard convnets fitting completely random ImageNet labels to zero training error, so failure to overfit a small batch is never "the model isn't big enough". It is a bug, such as labels misaligned with inputs (a shuffle applied to one but not the other is the classic), a loss reading the wrong axis, a softmax applied twice, gradients detached somewhere, a learning rate orders off, or data preprocessing that destroys the signal. The measured version on this page's network reached mini-batch loss 0.001 within 50 steps with sane initialization. The complementary check in the other direction, feeding the model shuffled labels and confirming it does not reach low validation accuracy, catches leakage.

Reading loss curves and monitoring norms

Beneath the loss, three families of time series carry most of the diagnostic information. Per-layer gradient norms should be broadly comparable across depth (the 20-layer experiment on this page measured \( 2.7 \) at the first layer under He initialization versus \( 1.5 \times 10^{-18} \) under \( \sigma_w = 0.01 \), and five minutes of logging finds in one glance what a week of retuning would not). The update-to-parameter ratio \( \lVert \Delta\theta \rVert / \lVert \theta \rVert \) should sit near \(10^{-3}\) per step as a rule of thumb. \(10^{-1}\) is thrashing and \(10^{-6}\) is not training. Activation statistics catch saturation (fraction of tanh units above 0.99, fraction of ReLU units dead across a large batch). The implementation below counts 0.7 percent dead at initialization, harmless, versus the 40-plus percent that a divergence event leaves behind. Layer-wise histograms of all three over time, which is what wandb.watch or TensorBoard histograms give, make most training pathologies visible before the loss curve shows them.

A checklist for a network that will not train

0. fix seed; make the run repeatable before touching anything
1. initial loss == ln K (or task equivalent)?     no -> init/loss bug
2. can it overfit 64 examples to ~0?              no -> wiring bug:
     labels aligned? loss axis? softmax twice? detach? requires_grad?
3. per-layer grad norms sane, no 1e-15s, no 1e+6s? no -> init/norm/depth
4. LR range test; is current LR in the descending band?
5. inputs actually normalized? (print mean/std of a real batch)
6. train/eval modes correct? (BN/dropout state is a silent killer)
7. only now: tune capacity, schedule, regularization

The ordering is the point. Each step assumes the previous ones passed, and the expensive activities (architecture search, regularization tuning) come last because they cannot fix what the earlier steps catch. The discipline that makes the list work is changing one thing at a time against a fixed seed, and keeping a log. Debugging by simultaneously changing the learning rate, the initialization, and the data pipeline produces a network that works and an engineer who does not know why.

Error analysis as a discipline

Splits and reference points

Once a network trains, the question becomes what to do next, and that question has a procedure, not a vibe. The data is split so that each number answers one question, training data (fit), a development set (all tuning decisions), and a test set (touched rarely, never tuned against, because every peek leaks information, which is also why leaderboard chasing overfits communities to test sets). Dev and test must be drawn from the distribution the model will face in deployment, and they must be large enough to resolve the differences being chased. Detecting a 0.1 percent improvement reliably needs tens of thousands of dev examples, and a 500-example dev set resolves only multi-point differences. When training data comes from a different distribution than deployment (scraped web images versus user phone photos), carve a second small set from the training distribution, a train-dev set, whose role appears below.

The second reference is an estimate of the Bayes error, in practice proxied by human-level performance measured on the same data with the same interface. Its use is arithmetic. Avoidable bias is training error minus the human reference, variance is dev error minus training error, and the larger of the two names the next project. A model at 8 percent training error looks underfit until human error on the task is measured at 7.5 percent, at which point almost all remaining headroom is variance and capacity increases are wasted effort.

The decision tree

train err >> human/Bayes ref?          [avoidable bias]
  -> bigger model, train longer, better optimizer/LR,
     architecture closer to the task, reduce regularization
train-dev err >> train err?            [variance]
  -> more data, augmentation, stronger regularization,
     smaller/simpler model, early stopping
dev err >> train-dev err?              [distribution mismatch]
  -> collect/synthesize data matching deployment,
     reweight training data, domain adaptation
test err >> dev err?                   [overfit the dev set]
  -> larger or fresh dev set; stop tuning against it
deployed metric bad, test fine?        [shift after deployment]
  -> monitor input stats, retrain cadence, data pipeline audit

The train-dev set is the hinge. Because it shares the training distribution, a gap between train and train-dev isolates variance, while a gap between train-dev and dev isolates distribution mismatch. Without it the two are confounded and teams routinely fight variance with more of the wrong data. The other pillar is manual error analysis, unglamorous and consistently the highest-leverage hour available. Take one to two hundred dev errors, read them, and tally categories (blurry inputs, a mislabeled class pair, long sequences, rare entities). The categories bound the value of every candidate fix, since eliminating a category entirely recovers only its share of the errors, and a category of mislabeled examples caps everything. Deciding between "improve the blur-robustness" (30 percent of errors) and "handle rare entities" (4 percent) stops being a debate.

Ablation discipline

The experimental habit that separates trustworthy claims from noise is to change one variable per run, rerun the baseline in the same code state, and repeat key comparisons across at least three seeds, because seed variance on mid-size benchmarks is frequently as large as the effect being claimed. Report the spread, not the best seed. When a change helps, ablate it back out of the final system to confirm it still matters in combination. Improvements are frequently non-additive, and a technique that helped a weak baseline often does nothing for a strong one (the recurring lesson of the optimizer-benchmark literature). A tuned baseline is the most valuable artifact a team owns. Most published "wins" over it are noise, and the fastest way to find the real ones is to make every comparison differ by exactly one thing.

Hyperparameter search

Random beats grid, provably in the case that matters

Bergstra and Bengio (2012) made the case that killed grid search for continuous hyperparameters. Loss surfaces over hyperparameters have low effective dimensionality. A few parameters (the learning rate above all) move the metric, and the rest barely matter, but which few is not known in advance. A grid of \( k \) values per dimension in \( d \) dimensions spends \( k^d \) runs but tests only \( k \) distinct values of each individual parameter. If one parameter dominates, the grid has explored it at resolution \( k \). The same \( n = k^d \) budget of random draws tests \( n \) distinct values of every parameter simultaneously, so whichever parameter turns out to matter has been explored at resolution \( n \). Their experiments confirmed random search matching or beating grid at a fraction of the budget across image and text tasks. Draw scale parameters log-uniformly (learning rate over \( [10^{-5}, 10^{-1}] \), weight decay over \( [10^{-6}, 10^{-1}] \)). A uniform draw over a log-scaled quantity wastes most of the budget on one decade.

Model-based and multi-fidelity search

Bayesian optimization treats the metric as an expensive black box. Fit a surrogate (Gaussian process, or the tree-structured Parzen estimator that Optuna implements) to the runs so far, and choose the next configuration by an acquisition rule (expected improvement) that trades exploring uncertain regions against exploiting good ones. It shines when each run is expensive, the space is under ~20 dimensions, and sequential decisions are acceptable. Hyperband (Li et al., 2018) attacks the budget from the other side. Most bad configurations reveal themselves early, so run many configurations at a small budget (epochs, data fraction), keep the top fraction \( 1/\eta \), multiply their budget by \( \eta \), and repeat (successive halving), with multiple brackets hedging against metrics whose early ranking is misleading. ASHA, its asynchronous version, promotes configurations as results arrive and is the default in Ray Tune at cluster scale. Population-based training (Jaderberg et al., 2017) folds search into a single training run. A population trains in parallel, and periodically poor performers copy the weights of good ones (exploit) and perturb their hyperparameters (explore), which discovers schedules, hyperparameters that change over training, that static search cannot express, at the price of a lineage-dependent, hard-to-reproduce result.

To budget concretely, spend the first and largest slice on the learning rate alone (a range test plus a log grid of ~5 values), add weight decay and warmup length as the second tier, treat batch size as set by hardware and the noise scale rather than searched, and for everything else prefer random draws under ASHA with a small final-budget rerun of the top three. On a fixed GPU budget, many cheap runs with early stopping buy more information than few full runs. The exception is when the early-training metric ranks configurations differently than the final metric, which is exactly the case Hyperband's brackets and PBT exist for.

More worked problems

Implementation

Everything in this section was run on the H100 80GB in this repository (PyTorch 2.7 and CUDA 12.8, with JAX 0.6 and optax), and the numbers in the comments are the numbers the runs printed. First, batch normalization forward and backward from scratch, with the analytic backward formula derived above checked against autograd. This is the canonical way to trust a hand-derived gradient. Compute both in float64 and look at the maximum absolute difference, which should sit at accumulation roundoff (\( \sim 10^{-15} \)), not at approximation error (\( 10^{-3} \)-ish would mean a wrong formula).

import torch

def bn_forward(x, gamma, beta, eps=1e-5):
    # x: (N, D); statistics per feature over the batch
    mu = x.mean(dim=0)                        # (D,)
    var = x.var(dim=0, unbiased=False)        # (D,)  biased, as in training
    xhat = (x - mu) / torch.sqrt(var + eps)   # (N, D)
    out = gamma * xhat + beta                 # (N, D)
    return out, (xhat, var, eps)

def bn_backward(dout, cache, gamma):
    # implements dx = (s/N) * (N*dxhat - sum(dxhat) - xhat*sum(dxhat*xhat))
    xhat, var, eps = cache
    N = dout.shape[0]
    dgamma = (dout * xhat).sum(dim=0)         # (D,)
    dbeta = dout.sum(dim=0)                   # (D,)
    dxhat = dout * gamma                      # (N, D)
    s = 1.0 / torch.sqrt(var + eps)           # inverse std, (D,)
    dx = (s / N) * (N * dxhat - dxhat.sum(dim=0)
                    - xhat * (dxhat * xhat).sum(dim=0))
    return dx, dgamma, dbeta

torch.manual_seed(0)
dev = "cuda"
x = torch.randn(64, 128, device=dev, dtype=torch.float64, requires_grad=True)
gamma = torch.randn(128, device=dev, dtype=torch.float64, requires_grad=True)
beta = torch.randn(128, device=dev, dtype=torch.float64, requires_grad=True)

out, cache = bn_forward(x, gamma, beta)
dout = torch.randn_like(out)
out.backward(dout)                            # autograd's answer
dx, dgamma, dbeta = bn_backward(dout, cache, gamma)

print((x.grad - dx).abs().max())      # measured: 1.78e-15  (float64)
print((gamma.grad - dgamma).abs().max())  # measured: 0.0
print((beta.grad - dbeta).abs().max())    # measured: 0.0
# same check in float32 gives 9.54e-07 on dx: same formula,
# float roundoff. ~1e-3 here would mean the derivation is wrong.

import jax, jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

def bn_forward(x, gamma, beta, eps=1e-5):
    mu = x.mean(axis=0)                       # (D,)
    var = x.var(axis=0)                       # (D,)
    xhat = (x - mu) / jnp.sqrt(var + eps)     # (N, D)
    return gamma * xhat + beta, (xhat, var, eps)

def bn_backward(dout, cache, gamma):
    xhat, var, eps = cache
    N = dout.shape[0]
    dgamma = (dout * xhat).sum(axis=0)
    dbeta = dout.sum(axis=0)
    dxhat = dout * gamma
    s = 1.0 / jnp.sqrt(var + eps)
    dx = (s / N) * (N * dxhat - dxhat.sum(axis=0)
                    - xhat * (dxhat * xhat).sum(axis=0))
    return dx, dgamma, dbeta

k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(0), 4)
x = jax.random.normal(k1, (64, 128), dtype=jnp.float64)
gamma = jax.random.normal(k2, (128,), dtype=jnp.float64)
beta = jax.random.normal(k3, (128,), dtype=jnp.float64)
dout = jax.random.normal(k4, (64, 128), dtype=jnp.float64)

# scalar contraction with dout makes jax.grad compute the same VJP
def probe(x, gamma, beta):
    out, _ = bn_forward(x, gamma, beta)
    return (out * dout).sum()

dx_ref, dg_ref, db_ref = jax.grad(probe, argnums=(0, 1, 2))(x, gamma, beta)
_, cache = bn_forward(x, gamma, beta)
dx, dg, db = bn_backward(dout, cache, gamma)

print(jnp.abs(dx - dx_ref).max())   # measured: 8.88e-16  (float64)
print(jnp.abs(dg - dg_ref).max())   # measured: 0.0
print(jnp.abs(db - db_ref).max())   # measured: 0.0

Next, Adam from scratch, bias correction and all, verified against the framework implementation not on one step but over 200 training steps of a real network, where any discrepancy in the moment updates, correction factors, or epsilon placement would compound into a visible divergence. After 200 steps in float64 the maximum absolute difference over all parameters was \( 8.3 \times 10^{-16} \) against torch.optim.Adam and \( 3.7 \times 10^{-16} \) against optax.adam. The scratch implementations are the algorithm, not an approximation of it.

import torch, torch.nn as nn, torch.nn.functional as F

class ScratchAdam:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
        self.params = list(params)
        self.lr, self.b1, self.b2, self.eps = lr, betas[0], betas[1], eps
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.t = 0

    @torch.no_grad()
    def step(self):
        self.t += 1
        for p, m, v in zip(self.params, self.m, self.v):
            if p.grad is None:
                continue
            g = p.grad
            m.mul_(self.b1).add_(g, alpha=1 - self.b1)        # first moment
            v.mul_(self.b2).addcmul_(g, g, value=1 - self.b2) # second moment
            mhat = m / (1 - self.b1 ** self.t)                # bias correction
            vhat = v / (1 - self.b2 ** self.t)
            p.addcdiv_(mhat, vhat.sqrt() + self.eps, value=-self.lr)

    def zero_grad(self):
        for p in self.params:
            p.grad = None

def make_net(seed):
    torch.manual_seed(seed)
    return nn.Sequential(nn.Linear(32, 64), nn.ReLU(),
                         nn.Linear(64, 10)).to("cuda", torch.float64)

net_a, net_b = make_net(1), make_net(1)          # identical weights
opt_a = ScratchAdam(net_a.parameters(), lr=1e-3)
opt_b = torch.optim.Adam(net_b.parameters(), lr=1e-3)

torch.manual_seed(2)
xs = torch.randn(256, 32, device="cuda", dtype=torch.float64)
ys = torch.randint(0, 10, (256,), device="cuda")
for step in range(200):
    for net, opt in ((net_a, opt_a), (net_b, opt_b)):
        opt.zero_grad()
        F.cross_entropy(net(xs), ys).backward()
        opt.step()

print(max((pa - pb).abs().max().item()
          for pa, pb in zip(net_a.parameters(), net_b.parameters())))
# measured: 8.33e-16 after 200 steps, float64

import jax, jax.numpy as jnp, optax
jax.config.update("jax_enable_x64", True)

def adam_init(params):
    return {"m": jax.tree.map(jnp.zeros_like, params),
            "v": jax.tree.map(jnp.zeros_like, params), "t": 0}

def adam_update(grads, state, params,
                lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    t = state["t"] + 1
    m = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g, state["m"], grads)
    v = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g * g, state["v"], grads)
    def upd(p, m, v):
        mhat = m / (1 - b1 ** t)              # bias correction
        vhat = v / (1 - b2 ** t)
        return p - lr * mhat / (jnp.sqrt(vhat) + eps)
    return jax.tree.map(upd, params, m, v), {"m": m, "v": v, "t": t}

def init_net(key, dims=(32, 64, 10)):
    ks = jax.random.split(key, len(dims) - 1)
    return [{"W": jax.random.normal(k, (a, b)) * jnp.sqrt(2.0 / a),
             "b": jnp.zeros((b,))}
            for k, a, b in zip(ks, dims[:-1], dims[1:])]

def apply(params, x):
    for i, layer in enumerate(params):
        x = x @ layer["W"] + layer["b"]
        if i < len(params) - 1:
            x = jax.nn.relu(x)
    return x

def loss_fn(params, x, y):
    return optax.softmax_cross_entropy_with_integer_labels(
        apply(params, x), y).mean()

kd, kl = jax.random.split(jax.random.PRNGKey(1))
xs = jax.random.normal(kd, (256, 32))
ys = jax.random.randint(kl, (256,), 0, 10)

p_a = init_net(jax.random.PRNGKey(2))
p_b = jax.tree.map(lambda x: x, p_a)             # identical copy
st_a = adam_init(p_a)
tx = optax.adam(1e-3)
st_b = tx.init(p_b)
grad_fn = jax.jit(jax.grad(loss_fn))
for _ in range(200):
    p_a, st_a = adam_update(grad_fn(p_a, xs, ys), st_a, p_a)
    upd, st_b = tx.update(grad_fn(p_b, xs, ys), st_b, p_b)
    p_b = optax.apply_updates(p_b, upd)

diffs = jax.tree.map(lambda a, b: float(jnp.abs(a - b).max()), p_a, p_b)
print(max(jax.tree.leaves(diffs)))
# measured: 3.75e-16 after 200 steps, float64

Third comes the experiment that makes initialization theory concrete. The same 20-layer, width-256 ReLU MLP is trained on a synthetic 10-class problem three times, changing nothing but the weight initialization, with the diagnostic suite from the checklist attached, the initial-loss check, per-layer gradient norms, the dead-unit count, and the overfit test. The measured results are in the comments and are worth reading as a story. The small initialization passes the initial-loss check perfectly (2.3026, exactly \( \ln 10 \)) and is nonetheless completely untrainable, with a first-layer gradient norm of \( 10^{-18} \), which is precisely why the checklist monitors gradients and not just the loss.

import math, torch, torch.nn as nn, torch.nn.functional as F

class DeepMLP(nn.Module):
    def __init__(self, depth=20, width=256, d_in=64, classes=10, init="he"):
        super().__init__()
        dims = [d_in] + [width] * (depth - 1) + [classes]
        self.layers = nn.ModuleList(nn.Linear(a, b)
                                    for a, b in zip(dims[:-1], dims[1:]))
        for lin in self.layers:
            if init == "he":
                nn.init.kaiming_normal_(lin.weight, nonlinearity="relu")
            elif init == "small":
                nn.init.normal_(lin.weight, std=0.01)
            elif init == "big":     # 2x the He std
                nn.init.normal_(lin.weight,
                                std=2 * math.sqrt(2 / lin.weight.shape[1]))
            nn.init.zeros_(lin.bias)

    def forward(self, x):
        for lin in self.layers[:-1]:
            x = F.relu(lin(x))
        return self.layers[-1](x)

@torch.no_grad()
def dead_fraction(net, x):
    h, dead, total = x, 0, 0
    for lin in net.layers[:-1]:
        h = F.relu(lin(h))
        dead += (h.amax(dim=0) == 0).sum().item()  # never fires in batch
        total += h.shape[1]
    return dead / total

def diagnose(init, xtr, ytr, steps=300):
    torch.manual_seed(0)
    net = DeepMLP(init=init).to("cuda")
    loss0 = F.cross_entropy(net(xtr[:1024]), ytr[:1024])
    print(f"init loss {loss0.item():.4f}  (theory ln10 = {math.log(10):.4f})")
    loss0.backward()
    g = [lin.weight.grad.norm().item() for lin in net.layers]
    print(f"grad norm first {g[0]:.3e}  last {g[-1]:.3e}")
    print(f"dead units at init: {100 * dead_fraction(net, xtr[:1024]):.1f}%")
    net.zero_grad()
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for step in range(steps):
        i = torch.randint(0, xtr.shape[0], (256,), device="cuda")
        opt.zero_grad()
        loss = F.cross_entropy(net(xtr[i]), ytr[i])
        loss.backward()
        opt.step()
        if step in (49, 99, steps - 1):
            print(f"step {step + 1}: loss {loss.item():.4f}")

# measured on the H100, synthetic 10-class task, depth 20:
#  init="he":    init loss 6.832 (logits too hot: He all the way to the
#                classifier head; small-init the last layer to fix),
#                grads 2.7e0 / 4.0e1, trains to 0.001 by step 50
#  init="small": init loss 2.3026 == ln 10 exactly -- and grad norm
#                1.5e-18 at layer 1: passes the loss check, dead anyway,
#                still at 2.30 after 300 steps
#  init="big":   init loss 7.1e6, grads ~1e6; Adam's normalization
#                eventually rescues it (~step 100), SGD would not

import math, jax, jax.numpy as jnp, optax

def init_mlp(key, depth=20, width=256, d_in=64, classes=10, mode="he"):
    dims = [d_in] + [width] * (depth - 1) + [classes]
    ks = jax.random.split(key, len(dims) - 1)
    std = {"he":   lambda n: math.sqrt(2 / n),
           "small": lambda n: 0.01,
           "big":  lambda n: 2 * math.sqrt(2 / n)}[mode]
    return [{"W": jax.random.normal(k, (a, b)) * std(a),
             "b": jnp.zeros((b,))}
            for k, a, b in zip(ks, dims[:-1], dims[1:])]

def apply(params, x):
    for layer in params[:-1]:
        x = jax.nn.relu(x @ layer["W"] + layer["b"])
    return x @ params[-1]["W"] + params[-1]["b"]

def loss_fn(params, x, y):
    return optax.softmax_cross_entropy_with_integer_labels(
        apply(params, x), y).mean()

def diagnose(mode, xtr, ytr, steps=300):
    params = init_mlp(jax.random.PRNGKey(0), mode=mode)
    l0, grads = jax.value_and_grad(loss_fn)(params, xtr[:1024], ytr[:1024])
    gnorms = [float(jnp.linalg.norm(g["W"])) for g in grads]
    print(f"init loss {l0:.4f} (ln10 = {math.log(10):.4f}); "
          f"grad first {gnorms[0]:.3e} last {gnorms[-1]:.3e}")
    tx = optax.adam(1e-3)
    st = tx.init(params)

    @jax.jit
    def step(params, st, x, y):
        loss, g = jax.value_and_grad(loss_fn)(params, x, y)
        upd, st = tx.update(g, st, params)
        return optax.apply_updates(params, upd), st, loss

    key = jax.random.PRNGKey(1)
    for i in range(steps):
        key, sub = jax.random.split(key)
        idx = jax.random.randint(sub, (256,), 0, xtr.shape[0])
        params, st, loss = step(params, st, xtr[idx], ytr[idx])
        if i in (49, 99, steps - 1):
            print(f"step {i + 1}: loss {loss:.4f}")

# same measured story as the PyTorch tab: "small" passes the
# initial-loss check (2.3026) with a 1e-18 first-layer gradient;
# monitoring the loss alone would misdiagnose it for hours.

Last, the learning-rate range test, which turns learning-rate selection from a superstition into a 30-second measurement. The loop below produced the measured trace quoted earlier, descent from \(10^{-5}\), a minimum smoothed loss of 0.002 at \( 8.8 \times 10^{-3} \), and catastrophic divergence just past \( 1.6 \times 10^{-2} \). The chosen rate should sit a factor of 2 to 5 below the cliff.

import math, torch, torch.nn.functional as F

def lr_range_test(make_net, xtr, ytr,
                  lr_min=1e-6, lr_max=10.0, steps=300, beta=0.9):
    net = make_net()
    opt = torch.optim.SGD(net.parameters(), lr=lr_min, momentum=0.9)
    gamma = (lr_max / lr_min) ** (1.0 / steps)   # geometric ramp
    trace, smooth = [], None
    for step in range(steps):
        lr = lr_min * gamma ** step
        for g in opt.param_groups:
            g["lr"] = lr
        i = torch.randint(0, xtr.shape[0], (256,), device=xtr.device)
        opt.zero_grad()
        loss = F.cross_entropy(net(xtr[i]), ytr[i])
        loss.backward()
        opt.step()
        l = loss.item()
        smooth = l if smooth is None else beta * smooth + (1 - beta) * l
        trace.append((lr, smooth))
        if not math.isfinite(l) or smooth > 50:   # stop after the cliff
            break
    return trace

# measured on the 20-layer MLP (SGD momentum 0.9):
#   lr 1.0e-06 smoothed 6.56     flat: too small to move
#   lr 1.3e-04 smoothed 1.32     descending
#   lr 3.2e-03 smoothed 0.009    near best
#   lr 8.8e-03 smoothed 0.002    minimum
#   lr 1.6e-02 smoothed 3.5      past the cliff
#   lr 5.5e-02 smoothed 4.3e5    gone
# pick ~3e-3 here: several-fold below the cliff, in the steep zone.

import math, jax, jax.numpy as jnp, optax

def lr_range_test(params, loss_fn, xtr, ytr,
                  lr_min=1e-6, lr_max=10.0, steps=300, beta=0.9):
    gamma = (lr_max / lr_min) ** (1.0 / steps)
    # inject_hyperparams makes the schedule's lr writable per step
    tx = optax.inject_hyperparams(optax.sgd)(
        learning_rate=lr_min, momentum=0.9)
    st = tx.init(params)

    @jax.jit
    def step(params, st, x, y, lr):
        st.hyperparams["learning_rate"] = lr
        loss, g = jax.value_and_grad(loss_fn)(params, x, y)
        upd, st = tx.update(g, st, params)
        return optax.apply_updates(params, upd), st, loss

    trace, smooth, key = [], None, jax.random.PRNGKey(0)
    for i in range(steps):
        lr = lr_min * gamma ** i
        key, sub = jax.random.split(key)
        idx = jax.random.randint(sub, (256,), 0, xtr.shape[0])
        params, st, loss = step(params, st, xtr[idx], ytr[idx],
                                jnp.float32(lr))
        l = float(loss)
        smooth = l if smooth is None else beta * smooth + (1 - beta) * l
        trace.append((lr, smooth))
        if not math.isfinite(l) or smooth > 50:
            break
    return trace

# read the trace the same way: choose a rate a factor of 2-5 below
# the divergence point, near the steepest descent.

How it is done in practice

Mixed precision, bf16 versus fp16

The economics are not subtle. On the H100 in this repository, a 4096-square matmul measures 51.3 TFLOP/s in fp32, 384.0 in TF32, and 744.6 in bf16. The tensor cores pay roughly 14× over fp32 for using 16-bit inputs, and no serious training run declines that. The standard recipe (Micikevicius et al., 2018) keeps a master copy of weights in fp32 inside the optimizer, runs forward and backward in 16-bit, and accumulates matmul partial products and gradient reductions in fp32. The choice within 16-bit is an exponent-versus-mantissa trade. Fp16 has 5 exponent bits and 10 mantissa bits, so it is more precise but overflows above 65,504 and, the real killer, flushes gradients below about \( 2^{-24} \) to zero. Bf16 has fp32's 8 exponent bits with only 7 mantissa bits, so it almost never over- or underflows but is coarse. Fp16 therefore requires loss scaling, multiplying the loss by a factor \( S \) (dynamically adjusted, doubling every couple of thousand clean steps and halving on overflow) so small gradients rise above the flush threshold, then unscaling before the optimizer step. Bf16 needs none of this machinery, which is why on hardware that supports it (Ampere onward, TPUs) bf16 is the default and fp16 survives mainly on older GPUs and in inference. The rough memory arithmetic for Adam-family training is 2 bytes per parameter for the working weights, plus fp32 master weights and two fp32 moments, about 16 bytes per parameter of state before activations, which is why optimizer-state precision and sharding are where memory work goes next.

Gradient accumulation and activation checkpointing

Gradient accumulation decouples the optimization batch from the memory batch. Run \( k \) forward/backward passes summing gradients, step once, and divide the loss (or gradients) by \( k \) so the result matches one large batch exactly, up to batch-statistics layers (batch norm sees the micro-batch, one of several reasons transformers' batch-independent norms simplify life). Activation checkpointing implements the memory/compute trade derived in Problem 5: PyTorch's torch.utils.checkpoint at block granularity, or jax.checkpoint (remat) in JAX, buying an order-of-magnitude activation-memory reduction for about a 33 percent step-time increase. The two compose. Accumulation shrinks the live batch, checkpointing shrinks the per-batch footprint, and together they let a single device train models far past its naive limit before any distributed machinery is invoked (that machinery, ZeRO-style state sharding and model parallelism, belongs to the LLM page).

Kernel fusion and the data pipeline

A training step's elementwise arithmetic, activations, residual adds, norm gains, optimizer math, is memory-bound. At the H100's measured ~3000 GB/s, a chain of \( k \) separate elementwise kernels reads and writes every tensor \( k \) times, and the arithmetic is nearly free compared to the traffic. Fusing the chain into one kernel does all the math per element in registers in a single read-write pass. On this machine, a bias-GELU-scale-residual chain that takes 0.911 ms as separate PyTorch ops runs in 0.214 ms after torch.compile fuses it, a 4.26× speedup with zero code change beyond the decorator, and the same fusion logic is why compiled training steps commonly gain tens of percent end to end. JAX gets the same effect from XLA by default. The other silent tax is the input pipeline. A GPU that finishes a step faster than the loader can produce the next batch idles, and the signature is unmistakable in a profile, gaps between kernels on the GPU timeline aligned with DataLoader activity on the CPU side. The standard fixes are worker processes (num_workers sized to the CPU), pinned memory with async host-to-device copies, prefetching a batch or two ahead, and moving decode-heavy augmentation to the GPU (DALI) or ahead of time. The rule for all of it is to profile before optimizing, with torch.profiler or Nsight Systems, and remember CUDA is asynchronous, so timing with time.time() without a synchronize measures launch overhead, not execution. Use CUDA events or the profiler.

Reproducibility on a GPU

"Same seed, same result" has fine print on a GPU. Seeding (torch.manual_seed, plus per-worker seeding in data loaders, plus the Python and NumPy generators if augmentation uses them) fixes the sampled randomness, but two further sources remain. First, algorithm selection. cuDNN autotuning (benchmark=True) picks kernels by timing, and timing varies, so different runs can pick different algorithms with different rounding. Second, genuine kernel nondeterminism. Floating-point addition is not associative, and any kernel that reduces with atomics (scatter-add, embedding-bag backward, some index operations) sums in hardware-scheduling order, so bitwise results differ run to run even with every seed fixed. torch.use_deterministic_algorithms(True) (with CUBLAS_WORKSPACE_CONFIG=:4096:8) swaps in deterministic implementations or errors where none exists, typically at a modest speed cost. Even then, "reproducible" means bitwise-identical on the same hardware, library versions, and parallel layout. Changing the GPU model, CUDA version, tensor-parallel degree, or even the reduction tree in a collective changes summation order and therefore the bits. The professional posture is to make single-machine runs deterministic during debugging (a heisenbug you cannot replay is not fixable), to accept statistical rather than bitwise reproducibility across hardware, and to treat "the metric moved when nothing changed" as a claim requiring three seeds before belief.

The current research frontier

The liveliest thread is optimizers, covered above. Muon (community origin, scaled by Moonshot AI to the trillion-parameter Kimi K2), Lion (Google), Sophia, and schedule-free methods (Defazio et al. at Meta, 2024, which drops the decay schedule entirely by averaging iterates) compete against a backdrop of benchmark skepticism, with MLCommons' AlgoPerf benchmark (Dahl et al., 2023) built specifically to force tuned-baseline comparisons. Its first round crowned distributed Shampoo, a second-order method from the Google line of work, on the self-tuning track, which says the field's oldest idea, preconditioning, still has headroom. Hyperparameter transfer is the second thread. The maximal update parameterization (muP, from Yang and Hu with OpenAI collaborators, 2021-2022) reparameterizes width so that the optimal learning rate is stable across scale, letting small proxies tune large models, and it has spread quietly into production LLM stacks (and its failure modes at depth are an active topic). Third, the norm wars continue, with normalization-free residual networks via signal-propagation-correct scaling (Brock et al., DeepMind, 2021, NFNets), QK-norm and other stability patches inside attention, and periodic claims that careful initialization plus residual scaling can retire normalization, which so far win benchmarks only with compensating machinery. Fourth, training-dynamics theory is catching up to practice, with edge-of-stability analyses (Cohen et al., 2021) showing optimizers hover at the stability boundary rather than below it, loss-spike forensics at LLM scale, and the empirics of warmup and schedules being reverse-engineered into theory (the WSD results above). The through-line of all four is that the craft on this page is being converted, result by result, from folklore into measurement, and the reliable way to evaluate any new entrant remains the same, a tuned baseline, matched compute, and more than one seed.

Open source to read

pytorch/pytorch is the reference implementation of everything on this page. Open tools/autograd/derivatives.yaml first, the entire backward pass of the framework as a declarative table of VJP formulas, one line per op, then torch/optim/adamw.py for decoupled decay exactly as derived above.

google/jax presents autodiff as a program transformation. Open docs/autodidax.py first, a self-contained reimplementation of JAX's core (tracing, JVP, VJP) in one readable file. It is the best backprop tutorial in any repo.

google-deepmind/optax treats optimizers as composable gradient transformations. Open optax/_src/transform.py first. scale_by_adam is Adam's moments and bias correction in thirty lines, and chaining it with weight decay reproduces AdamW.

karpathy/nanoGPT is the most readable serious training loop in public. Open train.py first and read it top to bottom, warmup plus cosine decay, AdamW with parameter groups that exclude norms and biases from decay, gradient accumulation, clipping, bf16 autocast, and torch.compile, every practice from this page in ~300 lines.

Lightning-AI/pytorch-lightning shows what a production-hardened training loop accretes (precision plugins, callbacks, sanity checks). Open src/lightning/pytorch/loops/training_epoch_loop.py first to see the anatomy under the abstraction.

huggingface/accelerate is the thinnest useful layer over device placement, mixed precision, and accumulation. Open src/accelerate/accelerator.py first. The prepare() method is the whole idea.

wandb/wandb covers experiment tracking. Open wandb/sdk/wandb_watch.py first, the hook-based gradient and parameter histogram logging that implements the monitoring this page's diagnostics section prescribes.

facebookresearch/hydra handles configuration for sweeps. Composable configs and multirun are how ablation discipline survives contact with a cluster. Open hydra/main.py first, the decorator that is the entire user contract.

optuna/optuna implements hyperparameter search with pruning. Open optuna/samplers/_tpe/sampler.py first, the tree-structured Parzen estimator behind suggest_float, then the median pruner for the Hyperband connection.

NVIDIA/apex is historically the origin of mixed-precision training in PyTorch. Its amp module is superseded by native torch.amp, but the fused kernels remain instructive. Open apex/normalization/fused_layer_norm.py first and trace it down to the CUDA kernel to see what "fused" means concretely.

Common misconceptions

"Adam adapts the learning rate, so tuning it barely matters." Adam normalizes per-coordinate scale, but the global \( \alpha \) still multiplies every update and still has a divergence cliff, and Adam's loss-versus-LR curve is the same U shape as SGD's, typically centered one to two orders of magnitude lower. The learning rate remains the first hyperparameter to tune under any optimizer.

"Batch norm works by reducing internal covariate shift." The experiment that would confirm this refutes it. Injecting non-stationary distributional shift after every batch-norm layer leaves its benefits intact (Santurkar et al., 2018). The supported account is optimization-geometric, smoother loss and gradients, tolerance of larger learning rates, plus the scale-invariance the normalization imposes on incoming weights.

"L2 regularization and weight decay are the same thing." True under vanilla SGD, false under any adaptive method. Routed through Adam's preconditioner, the L2 gradient decays well-trained weights least and quiet weights most. Decoupling the decay from the gradient path (AdamW) restores uniform shrinkage and changes results enough that the two must be treated as different regularizers with different good values.

"The network won't fit a small batch, so it needs more capacity." Standard networks can fit random labels (Zhang et al., 2017), and memorizing 64 examples requires almost no capacity. Failure to overfit a single batch is evidence of a wiring bug, misaligned labels, a broken loss, a detached graph, a wild learning rate, and adding capacity to fix it treats the wrong disease.

"A loss starting at the theoretical \( \ln K \) means initialization is fine." The measured counterexample on this page is \( \sigma_w = 0.01 \), which produced an initial loss of exactly 2.3026 on ten classes, with a first-layer gradient norm of \( 1.5 \times 10^{-18} \). The loss check catches wrong logit scales. Only gradient monitoring catches dead signal paths. Run both.

"Grid search is more systematic than random search." With a budget of \( n \) runs over \( d \) hyperparameters of which few matter, grid tests only \( n^{1/d} \) distinct values of the one that does, while random tests \( n \). Grid's apparent rigor is resolution wasted on dimensions that turn out not to matter (Bergstra and Bengio, 2012). Grids remain fine for small discrete choices.

"Warmup is a superstition." It has at least three derivable justifications. Adam's second-moment estimate is high-variance for roughly \( 1/(1-\beta_2) \) steps and mis-scales early updates (the Problem 3 arithmetic), linear-scaled large-batch learning rates are unstable specifically from a cold start, and post-norm transformer gradients are depth-skewed at initialization. Skipping warmup is free risk with no offsetting benefit.

"Fixing the seed makes GPU training reproducible." Seeds fix sampled randomness, not summation order. Autotuned kernel selection and atomic-based reductions make many kernels nondeterministic at fixed seed, and cross-hardware bitwise reproducibility is effectively unobtainable. Deterministic modes exist for debugging on fixed hardware. Across machines, reproducibility means distributions over seeds, not bits.

Self-check

References

1. Goodfellow, I., Bengio, Y., Courville, A. Deep Learning. MIT Press, 2016. deeplearningbook.org (chapters 6-8, 11).
2. Bishop, C. M., Bishop, H. Deep Learning: Foundations and Concepts. Springer, 2024. bishopbook.com.
3. Zhang, A., Lipton, Z., Li, M., Smola, A. Dive into Deep Learning. Cambridge University Press, 2023. d2l.ai.
4. Glorot, X., Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. AISTATS 2010. pmlr v9.
5. He, K., Zhang, X., Ren, S., Sun, J. Delving deep into rectifiers: surpassing human-level performance on ImageNet classification. ICCV 2015. arXiv:1502.01852.
6. Saxe, A., McClelland, J., Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. ICLR 2014. arXiv:1312.6120.
7. Pennington, J., Schoenholz, S., Ganguli, S. Resurrecting the sigmoid in deep learning through dynamical isometry. NeurIPS 2017. arXiv:1711.04735.
8. He, K., Zhang, X., Ren, S., Sun, J. Deep residual learning for image recognition. CVPR 2016. arXiv:1512.03385.
9. He, K., Zhang, X., Ren, S., Sun, J. Identity mappings in deep residual networks. ECCV 2016. arXiv:1603.05027.
10. Veit, A., Wilber, M., Belongie, S. Residual networks behave like ensembles of relatively shallow networks. NeurIPS 2016. arXiv:1605.06431.
11. Ioffe, S., Szegedy, C. Batch normalization: accelerating deep network training by reducing internal covariate shift. ICML 2015. arXiv:1502.03167.
12. Ba, J., Kiros, J., Hinton, G. Layer normalization. 2016. arXiv:1607.06450.
13. Wu, Y., He, K. Group normalization. ECCV 2018. arXiv:1803.08494.
14. Zhang, B., Sennrich, R. Root mean square layer normalization. NeurIPS 2019. arXiv:1910.07467.
15. Santurkar, S., Tsipras, D., Ilyas, A., Madry, A. How does batch normalization help optimization? NeurIPS 2018. arXiv:1805.11604.
16. Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., Salakhutdinov, R. Dropout: a simple way to prevent neural networks from overfitting. JMLR 15, 2014. jmlr v15.
17. Huang, G., Sun, Y., Liu, Z., Sedra, D., Weinberger, K. Deep networks with stochastic depth. ECCV 2016. arXiv:1603.09382.
18. Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., Wojna, Z. Rethinking the Inception architecture for computer vision. CVPR 2016. arXiv:1512.00567.
19. Müller, R., Kornblith, S., Hinton, G. When does label smoothing help? NeurIPS 2019. arXiv:1906.02629.
20. Zhang, H., Cissé, M., Dauphin, Y., Lopez-Paz, D. mixup: beyond empirical risk minimization. ICLR 2018. arXiv:1710.09412.
21. Kingma, D., Ba, J. Adam: a method for stochastic optimization. ICLR 2015. arXiv:1412.6980.
22. Reddi, S., Kale, S., Kumar, S. On the convergence of Adam and beyond. ICLR 2018. arXiv:1904.09237.
23. Loshchilov, I., Hutter, F. Decoupled weight decay regularization. ICLR 2019. arXiv:1711.05101. Also SGDR: stochastic gradient descent with warm restarts. ICLR 2017. arXiv:1608.03983.
24. Smith, L. Cyclical learning rates for training neural networks. WACV 2017. arXiv:1506.01186. Smith, L., Topin, N. Super-convergence. 2018. arXiv:1708.07120.
25. Goyal, P., Dollár, P., Girshick, R., et al. Accurate, large minibatch SGD: training ImageNet in 1 hour. 2017. arXiv:1706.02677.
26. McCandlish, S., Kaplan, J., Amodei, D., et al. An empirical model of large-batch training. 2018. arXiv:1812.06162.
27. Bergstra, J., Bengio, Y. Random search for hyper-parameter optimization. JMLR 13, 2012. jmlr v13.
28. Li, L., Jamieson, K., DeSalvo, G., Rostamizadeh, A., Talwalkar, A. Hyperband: a novel bandit-based approach to hyperparameter optimization. JMLR 18, 2018. arXiv:1603.06560.
29. Jaderberg, M., Dalibard, V., Osindero, S., et al. Population based training of neural networks. 2017. arXiv:1711.09846.
30. Micikevicius, P., Narang, S., Alben, J., et al. Mixed precision training. ICLR 2018. arXiv:1710.03740.
31. Chen, T., Xu, B., Zhang, C., Guestrin, C. Training deep nets with sublinear memory cost. 2016. arXiv:1604.06174.
32. Zhang, C., Bengio, S., Hardt, M., Recht, B., Vinyals, O. Understanding deep learning requires rethinking generalization. ICLR 2017. arXiv:1611.03530.
33. Chen, X., Liang, C., Huang, D., et al. Symbolic discovery of optimization algorithms (Lion). NeurIPS 2023. arXiv:2302.06675.
34. Liu, H., Li, Z., Hall, D., Liang, P., Ma, T. Sophia: a scalable stochastic second-order optimizer for language model pre-training. 2023. arXiv:2305.14342.
35. Liu, J., Su, J., Yao, X., et al. Muon is scalable for LLM training. 2025. arXiv:2502.16982.
36. Hägele, A., Bakouch, E., Kosson, A., et al. Scaling laws and compute-optimal training beyond fixed training durations. NeurIPS 2024. arXiv:2405.18392.