Unit 3.3 · about 4 h
187 / 324Batch, layer, group and RMS normalization
Normalization decides which axes share statistics and whether evaluation depends on the batch. The wrong choice breaks single-example prediction or leaks batch mates into each output.
- derive
- implement
You can now
- Derive the normalized output, parameter count and normalized axes for batch, layer, group and RMS normalization.
- Implement layer normalization from tensor ops and match the reference on random input.
- Predict when batch statistics help, hurt or break, and choose the matching normalization.
Covers
- Batch, layer, group and RMS normalization axes
Study, in this order
5 steps
- Batch, layer and prediction behaviorD2L - optimization and regularization slice
Batch Normalization Layers (fully connected and convolutional), Layer Normalization, and Batch Normalization During Prediction
- Batch normalization (equations and behavior)Understanding Deep Learning - fitting, gradients and regularization
Chapter 11, Section 11.4 (pp. 192 to 194): batch equations, train versus test statistics, and costs with benefits
- Group normalization axesPyTorch GroupNorm reference
GroupNormpage: channel grouping, normalized axes and the affine parameters - RMS normalization formulaPyTorch RMSNorm reference
RMSNormpage: RMS formula,normalized_shapeaxes, and absent centering - HOML: batch norm and layer norm in PyTorchHands-On Machine Learning with Scikit-Learn and PyTorch
Chapter 11, Batch Normalization and Layer Normalization (printed pp. 375-382)
Depth · Read it for the PyTorch conventions: the BN
momentumis the weight of the new batch statistic (default 0.1), the opposite of the usual running-average meaning;BatchNorm1dexpects[batch, features, length];nn.LayerNormnormalizes the trailingnormalized_shapedimensions, with an explicitunbiased=Falseequivalent on p. 381.
Focus on
- Batch norm normalizes each channel over batch and spatial axes with learned per-channel scale and shift. One example's training output depends on the other batch members.
- Layer norm normalizes one example over its features with no running statistics, so train and test match. Group norm sits between: \(G\) channel groups per example, with \(G = C\) giving per-channel behavior.
- RMS norm drops mean centering and bias and scales by the root mean square only. Fewer statistics, same axis argument.
- PyTorch tracks batch statistics as exponential running averages during training. D2L's full-dataset recomputation after training describes the goal, not the implementation.
Reproduce closed-book
(batch, features) input from sums only. Then compute batch norm by hand for values [0, 2, 4] with unit scale, zero shift and zero \(\epsilon\): outputs -1.225, 0 and 1.225.Traps
- Normalizing with batch statistics at test time. Eval must use the tracked running statistics; with a batch of one, a fully connected batch-norm layer would map every feature to its shift parameter.
- Keeping a bias right before batch normalization. Batch norm subtracts the per-channel batch mean, which removes any per-channel constant, so that bias is redundant (a D2L exercise asks this). Before layer norm a per-feature bias is not fully removed, because layer norm subtracts one mean across features.
- Treating RMS norm as layer norm without bias. It also removes mean centering, so a large constant offset passes through scaled but unshifted.
Practice
TrenTorch: Layer Normalization, forward
Implement layer_norm_forward(x, gamma, beta, eps) over the last axis; one test compares against torch.nn.LayerNorm
Use the biased variance (divide by \(n\)) with \(\epsilon\) inside the square root; a test fails on the \(n - 1\) version.
D2L Batch Normalization from Scratch and Concise implementation
Implementation from Scratch and Concise Implementation: build batch norm, match the concise layer, and answer the bias-removal exercise
Self-test
-
Batch values
[0, 2, 4], \(\gamma\) 1, \(\beta\) 0, \(\epsilon\) 0. What are the normalized outputs?Mean 2, variance 8/3, std 1.633. Outputs (-2/1.633, 0, 2/1.633) = (-1.225, 0, 1.225). -
Conv input \((N, C, H, W)\). Which axes does batch norm use per channel, and why does batch size 1 still work?
Batch, height and width. Even one image supplies \(H\) times \(W\) values per channel, so mean and variance are defined. Fully connected batch norm has only the batch axis and degenerates at batch size 1. -
Group norm with \(G = 1\) versus \(G = C\) on one example. What does each compute?
\(G = 1\) normalizes all channels and spatial positions jointly with one mean and variance. \(G = C\) normalizes each channel separately, which is per-channel instance behavior. -
RMS norm input
[3, 4], \(\gamma\) 1. What is the output?\(\text{RMS} = \sqrt{(9+16)/2} = 3.536\). Outputs (0.849, 1.131). No centering is applied.
Completion criterion
Stop when
Your from-scratch layer norm matches the framework output on random input, and you state each method's axes and eval behavior without notes.
Mark passed, then open 3.4 Only mark this when you can do it without notes.Leads to
- 3.8 Diagnosing failed runs: divergence, stale state, slices and seeds Debug stale train and eval normalization state.
- 7.5 Transformer blocks: residual stream, normalization and feed-forward Place normalization inside residual blocks.
Comes after
- 3.1 Initialization and signal scale across depthSee how normalization re-centers and re-scales the signals derived here.
- 3.2 Dead, saturated and exploding units, and residual pathsControl the scale you just learned to measure.
Resources for this unit
- Understanding Deep Learning - fitting, gradients and regularization Use Chapters 6-9 and targeted parts of Chapter 20 for optimization, initialization, measurement, regularization and why training works. Re-derive the update rules you keep in the trainer.
- D2L - optimization and regularization slice Use MLP implementation, dropout, deep generalization, numerical stability, optimization introduction, SGD, minibatch SGD, momentum, AdaGrad, RMSProp, Adam and learning-rate scheduler. Implement momentum, AdamW, initialization, dropout, one normalization, clipping and one schedule.
- PyTorch GroupNorm reference Read how input channels are split into
num_groupsgroups and which axes each group normalizes; comparenum_groups1 againstnum_channels. - PyTorch RMSNorm reference Read the RMS formula and which axes
normalized_shapecovers; note the missing mean centering and bias. - Hands-On Machine Learning with Scikit-Learn and PyTorch Read the Chapter 11 sections named in the units for the PyTorch side of initialization, normalization, clipping and schedules.
Also in this section
- PyTorch - optimization and automatic mixed precision examples Check optimizer, scheduler, gradient clipping and AMP API semantics. Preserve the correct order for zeroing, scaling, backward, unscaling, clipping, stepping and updating.
- Deep-ML candidate - optimization and normalization problems After the platform admission check, use at most one of 146 Momentum, 169 AdamW, 197 Global-Norm Clipping, 109 LayerNorm or 115 BatchNorm as a changed-prompt drill.
- TrenTorch - training mechanics questions Solve only the questions the units name: Weight Initialization, Layer Normalization, Label Smoothing, Gradient Accumulation, and Warmup and Cosine Decay. Read the tests first; they run in the browser IDE.
- PyTorch torch.optim.AdamW reference Read the AdamW algorithm box for the decoupled weight-decay step, the bias-corrected moments and the defaults, and follow the Decoupled Weight Decay Regularization link for the argument.
- PyTorch CrossEntropyLoss reference Read Shape, Parameters (
weight,ignore_index,reduction,label_smoothing) and both Examples; extend the class-index example withignore_indexand a classweight. - Cyclical Learning Rates for Training Neural Networks (Smith, 2017) Read the abstract and the range test: increase the rate linearly for a few epochs and read usable bounds off the loss curve.
- Deep Learning Tuning Playbook FAQ: debugging optimization failures Read How can optimization failures be debugged and mitigated (identifying unstable workloads, warmup, clipping) and Batch normalization implementation details.