Batch Normalization · Coordinate Descent · Polyak Averaging

§8.7.1 Batch normalization

Batch normalization (Ioffe & Szegedy, 2015) is one of the most important recent innovations — and it is not an optimization algorithm. It is an adaptive reparametrization that tames the core difficulty of very deep nets: the gradient tells you how to update each layer assuming the others stay fixed, but you update them all at once. Take a deep linear chain $\hat{y} = x\,w_1 w_2 \cdots w_l$ :

Term by term

\hat{y}_\text{new} = x\,(w_1 - \epsilon g_1)(w_2 - \epsilon g_2)\cdots(w_l - \epsilon g_l)

$\text{first order}$	the gradient predicts ŷ drops by ε·gᵀg — so you might pick ε from that	the plan
$\text{higher order}$	but the real update multiplies l changed factors, adding 2nd…lth-order terms like ε²g₁g₂∏wᵢ that can be negligible OR exponentially large	the surprise
$\Rightarrow$	the effect of updating one layer depends on ALL the others, so the learning rate is nearly impossible to choose. Batch norm sidesteps this entirely	the problem

Batch norm’s fix: standardize each unit’s activations across the minibatch, then give the network back a clean way to choose the scale and shift. Drag $\gamma$ and $\beta$ to see the three stages:

Batch norm on one unit: raw H → normalized (H−µ)/σ (always zero-mean, unit-variance) → γĤ+β (learned scale/shift, mean = β, std = γ)

γ = 1.00β = 0.00

The middle row is always zero-mean, unit-variance — that is what removes the cross-layer coordination problem. Then γ and β restore expressive power, but now the output's mean is set solely by β and its std solely by γ, instead of by a tangled interaction of every lower-layer parameter — so gradient descent can adjust them cleanly.

Term by term

\boldsymbol{H}' = \frac{\boldsymbol{H} - \boldsymbol{\mu}}{\boldsymbol{\sigma}} \quad\longrightarrow\quad \gamma\boldsymbol{H}' + \beta

$(\boldsymbol{H}-\boldsymbol{\mu})/\boldsymbol{\sigma}$	standardize each unit to zero mean, unit variance using the MINIBATCH mean/std — and crucially back-propagate THROUGH µ and σ, so the gradient can never just shift the mean or std	normalize
$\gamma\boldsymbol{H}' + \beta$	learned scale γ and shift β restore full expressive power — but now the output mean is set solely by β and the std solely by γ, easy for gradient descent to adjust	restore
$\text{test time}$	replace µ,σ with running averages collected during training, so the model can be evaluated on a single example	inference

Why back-propagate through the normalization?

That is the key innovation. Earlier attempts either added a penalty encouraging normalized statistics (imperfect) or renormalized after each gradient step (wasteful — the optimizer kept proposing to change the mean and variance, and the renormalization kept undoing it). By making the normalization part of the model and differentiating through it, the gradient’s request to change a unit’s mean or variance is zeroed out automatically. In a purely linear net this would render the lower layers useless (their only influence — first and second-order statistics — is normalized away); in a real nonlinear net they still perform useful nonlinear transformations. Apply batch norm to $\boldsymbol{XW}$ (dropping the redundant bias $b$ , absorbed by $\beta$ ), and in CNNs use the same µ, σ at every spatial location.

§8.7.2 Coordinate descent

Sometimes a problem splits cleanly. Coordinate descent minimizes with respect to one variable at a time (or a block — block coordinate descent ), cycling until convergence. It shines when the variables separate into fairly isolated groups. The classic case is sparse coding, $J(\boldsymbol{H}, \boldsymbol{W}) = \sum_{i,j}|H_{i,j}| + \sum_{i,j}(\boldsymbol{X} - \boldsymbol{W}\boldsymbol{H})^2_{i,j}$ : jointly non-convex, but convex in the dictionary $\boldsymbol{W}$ with the codes $\boldsymbol{H}$ fixed, and vice versa — so alternate between the two. It is a poor strategy, though, when one variable strongly influences another’s optimum (e.g. $f = (x_1 - x_2)^2 + \alpha(x_1^2 + x_2^2)$ with small $\alpha$ ), where changing one coordinate far from the other is penalized and progress crawls.

§8.7.3 Polyak averaging

Polyak averaging averages the points visited along the optimization trajectory rather than returning the last one:

Term by term

\hat{\boldsymbol{\theta}}^{(t)} = \frac{1}{t}\sum_{i=1}^{t} \boldsymbol{\theta}^{(i)} \qquad \text{(non-convex: } \hat{\boldsymbol{\theta}}^{(t)} = \alpha\hat{\boldsymbol{\theta}}^{(t-1)} + (1-\alpha)\boldsymbol{\theta}^{(t)}\text{)}

$\frac{1}{t}\sum \boldsymbol{\theta}^{(i)}$	the plain average of the whole trajectory — if the optimizer leaps back and forth across a valley without hitting the bottom, the average of both sides lands near it	convex form
$\alpha\hat{\boldsymbol{\theta}}^{(t-1)} + (1-\alpha)\boldsymbol{\theta}^{(t)}$	for non-convex problems, an EXPONENTIALLY-DECAYING average — distant past points may be behind cost barriers and are down-weighted	non-convex form
$\text{guarantee}$	strong convergence guarantees for convex gradient descent; for neural nets it is heuristic but works well in practice	the caveat

Next: pretraining, designing models to aid optimization, and continuation / curriculum learning — pretraining, model design & curriculum.

Check yourself — batch norm, coordinate descent, Polyak averaging

1.What problem does batch normalization address, using the deep linear chain ŷ = x·w₁w₂···wₗ as an example?

2.In BatchNormLab, the middle row (H−µ)/σ is always zero-mean and unit-variance regardless of γ and β. Why does that matter?

3.Why is it crucial to BACK-PROPAGATE through the mean µ and std σ in batch normalization?

4.When does coordinate (or block coordinate) descent work well, and when does it struggle?

5.What is Polyak averaging, and how is it adapted for non-convex problems?

5 questions