Training Neural Networks, machine learning lesson

A correct gradient is only one ingredient of training. The data have to be split and shuffled, the weights need starting values, an optimizer turns gradients into updates, and progress has to be judged on data the updates never saw. This lesson follows a small constructed experiment, a two-feature classifier trained with two learning rates, and uses it to define the training loop. It then explains how initialization and depth decide whether gradients stay usable, how momentum and Adam use past gradients, and why batch normalization behaves differently in training and in evaluation.

Key terms

Training loop

Definition A training loop fits the parameters $\boldsymbol{\theta}$ of a model to a training set of $n$ examples by repeating epochs, where an epoch is one pass through the training data. In each epoch it does the following.

Shuffle the order of the training examples and split it into mini-batches.
For each mini-batch of $m$ examples, run the forward pass, compute the mean loss over the batch, clear any stored gradients, run the backward pass, and let the optimizer update every parameter with the gradients from this same forward pass.
After the epoch, compute the loss on the full training set and on the validation set with the model in evaluation mode and without any update.
When the validation loss is the best so far, save a copy of the parameters as a checkpoint, with the epoch number and, if training may resume, the optimizer state.

Validation examples never enter a gradient, and the test set is not used inside the loop. For a binary label $y$ in $\{0, 1\}$ and a logit $z$ , the loss of one example computed directly from the logit is

\ell(z, y) = \ln\bigl(1 + e^{z}\bigr) - yz ,

which equals the binary cross-entropy $-y\ln\sigma(z) - (1 - y)\ln\bigl(1 - \sigma(z)\bigr)$ . The derivative of the batch mean with respect to the logit of example $i$ is $\bigl(\sigma(z_i) - y_i\bigr)/m$ . Evaluating $\ln(1 + e^{z})$ in a numerically stable way avoids taking the logarithm of a probability that has rounded to $0$ or $1$ .

In plain words The loop is the experiment written as code, and each line has one job. Many failed runs come from a line doing the wrong job, such as gradients left over from the previous batch, a loss summed where the learning rate assumes a mean, or validation run in training mode. A learning-rate schedule, if used, belongs to the loop as well and must say whether it advances per update or per epoch. A fixed random seed makes the shuffles and the initialization repeatable on the same software, but other hardware or library versions can still change the floating-point results.

Example In a constructed problem two features are drawn uniformly from $[-1, 1]$ , and the label is $1$ when $x_1 x_2 \ge 0$ , that is, in the first and third quadrants including the axes. No straight line separates the classes. Separate random draws give $96$ training and $48$ validation examples. The network has $2$ inputs, $8$ hidden $\tanh$ units and $1$ logit, both runs start from the same initialization, and plain mini-batch gradient descent with batches of $16$ makes $6$ updates per epoch for $250$ epochs, with a constant learning rate of $0.03$ or $3$ . With the examples of a batch stored as the rows of a $16 \times 2$ array, the hidden array is $16 \times 8$ , the logits form a vector of length $16$ , and every gradient array has the shape of its parameter. The validation results were

Run	Epoch	Validation loss	Correct of 48
Both rates, before training	0	0.7129	24
Rate 0.03	250	0.6549	34
Rate 3, best checkpoint	91	0.3641	43
Rate 3, last epoch	250	0.4205	42

Predicting $0.5$ for everyone gives a validation loss of $\ln 2 \approx 0.6931$ , and predicting class $1$ for everyone gets $26$ of $48$ right, so the slow run made only modest progress in $1500$ updates. The fast run did much better, but its validation loss fluctuated, for example $0.5210$ at epoch $100$ , and its last state is worse than the checkpoint from epoch $91$ , which is the model kept. This is one seed on one small constructed problem, so it does not show that large learning rates are better in general. The checkpoint was chosen for its validation loss, so $0.3641$ is an optimistic estimate of its loss on new data, and an untouched test set gives the honest final number.

Weight initialization

Definition Weight initialization is the rule that sets the parameters before the first update. It has two jobs. The first is to break symmetry. If two hidden units of a layer start with the same incoming weights, the same bias and the same outgoing weights, they compute the same value on every input and receive the same gradients, so gradient descent, momentum and Adam keep them identical forever. Random initial weights break this symmetry, and biases can then start at $0$ . The second job is to set the scale. Suppose a preactivation is $z = \sum_{j=1}^{n_{\text{in}}} w_j x_j$ with weights drawn independently of each other and of the inputs, each with mean $0$ and variance $s^2$ . Then, over the random draw,

\operatorname{Var}(z) = s^2 \sum_{j=1}^{n_{\text{in}}} E\bigl[x_j^2\bigr] .

The choice $s^2 = 1/n_{\text{in}}$ keeps $\operatorname{Var}(z)$ equal to the average of the $E[x_j^2]$ . Glorot initialization, also called Xavier initialization, uses $s^2 = 2/(n_{\text{in}} + n_{\text{out}})$ and was derived for activations such as $\tanh$ that are close to linear near $0$ . He initialization uses $s^2 = 2/n_{\text{in}}$ for ReLU layers. The factor $2$ makes up for ReLU setting the negative half of a preactivation that is symmetric around $0$ to zero, which halves its mean square.

In plain words Units that start identical stay copies of each other. Weights that start too large push $\tanh$ and sigmoid units into their flat regions, where the derivative is nearly $0$ , and weights that start too small shrink the signal layer after layer. Scale rules aim to keep typical preactivations and gradients at similar sizes in every layer. They are sensible defaults, not optimal values.

Example The experiment drew first-layer weights with standard deviation $1/\sqrt{2}$ , for two inputs, and output weights with standard deviation $1/\sqrt{8}$ , for eight hidden units, with all biases $0$ . Features uniform on $[-1, 1]$ have $E[x_j^2] = 1/3$ , so averaged over the random draw each first-layer preactivation starts with variance $2 \cdot \tfrac12 \cdot \tfrac13 = \tfrac13$ . If instead every weight and bias starts at $0$ , every hidden value is $\tanh(0) = 0$ . The output weights then get the gradient $H^{\top}\boldsymbol{\delta} = \mathbf{0}$ , where $H$ holds the hidden values and $\boldsymbol{\delta}$ the derivatives of the loss with respect to the logits, and the first-layer gradients are multiples of the output weights, which are $0$ . Only the output bias ever changes, so the network remains a constant predictor at every learning rate. Rerun this way at rate $3$ , the experiment ends with a validation loss of about $0.6927$ , essentially that of predicting $0.5$ for everyone.

Vanishing and exploding gradients

Definition Consider hidden layers $\mathbf{h}^{(k)} = \phi\bigl(W^{(k)}\mathbf{h}^{(k-1)} + \mathbf{b}^{(k)}\bigr)$ for $k = 1, \dots, M$ , and write $J_k = \operatorname{diag}\bigl(\phi'(\mathbf{z}^{(k)})\bigr)W^{(k)}$ for the Jacobian of layer $k$ with respect to its input. The chain rule gives, for $l < M$ ,

\bar{\mathbf{h}}^{(l)} = J_{l+1}^{\top} J_{l+2}^{\top} \cdots J_{M}^{\top}\, \bar{\mathbf{h}}^{(M)},

so the gradient reaching an early layer is a product of many factors, and in the spectral norm

\begin{aligned} \bigl\lVert \bar{\mathbf{h}}^{(l)} \bigr\rVert_2 &\le \Bigl( \prod_{k=l+1}^{M} \lVert J_k \rVert_2 \Bigr) \\[0.4em] &\quad \cdot \bigl\lVert \bar{\mathbf{h}}^{(M)} \bigr\rVert_2 . \end{aligned}

A vanishing gradient is one whose size shrinks roughly geometrically toward the early layers, so their parameters barely change. An exploding gradient grows roughly geometrically and produces huge or nonfinite updates. If $\lvert \phi' \rvert \le 1$ , as for $\tanh$ , and every $\lVert W^{(k)} \rVert_2 \le \rho < 1$ , the bound falls at least like $\rho^{M-l}$ , and for the sigmoid $\lvert \phi' \rvert \le 1/4$ . No norm bound forces an explosion, but factors larger than $1$ along the directions the gradient actually takes make one happen. Gradient clipping by norm with a threshold $\tau > 0$ replaces $\mathbf{g}$ by $\mathbf{g}\min\bigl(1, \tau/\lVert \mathbf{g} \rVert_2\bigr)$ , which limits the length of a step and keeps its direction.

In plain words Each layer multiplies the backward signal by its own factor. Many factors a little below $1$ starve the early layers of signal, and many factors above $1$ blow it up. Watching the gradient norm of each layer and checking for nonfinite values often shows the problem earlier than the loss does. Scale-aware initialization, ReLU-type activations, normalization layers and smaller learning rates all help. Clipping tames an explosion but cannot restore a gradient that has vanished, and it does not repair a wrong loss.

Example A chain of ten scalar factors equal to $0.5$ multiplies a gradient by $(1/2)^{10} = 1/1024 \approx 0.000977$ , and ten factors equal to $2$ multiply it by $1024$ . The derivative $1 - \tanh^2(z)$ is $1$ at $z = 0$ , about $0.42$ at $z = 1$ and about $0.00018$ at $z = 5$ , so a saturated $\tanh$ unit passes back less than $0.02\%$ of the gradient it receives. In a chain of ten sigmoid units joined by weights equal to $1$ , the gradient is multiplied by at most $(1/4)^{10} \approx 9.5 \times 10^{-7}$ . Clipping the gradient $(30, 40)$ , whose norm is $50$ , at $\tau = 5$ gives $(3, 4)$ .

Momentum

Definition Gradient descent with momentum keeps a running average $\mathbf{m}$ of past gradients. With a momentum constant $\beta$ in $[0, 1)$ , a learning rate $\eta > 0$ and $\mathbf{m}_0 = \mathbf{0}$ , step $t$ computes the mini-batch gradient $\mathbf{g}_t$ at $\boldsymbol{\theta}_{t-1}$ and sets

\begin{aligned} \mathbf{m}_t &= \beta\,\mathbf{m}_{t-1} + (1 - \beta)\,\mathbf{g}_t, \\[0.4em] \boldsymbol{\theta}_t &= \boldsymbol{\theta}_{t-1} - \eta\,\mathbf{m}_t . \end{aligned}

With $\beta = 0$ this is plain stochastic gradient descent. Unrolling gives $\mathbf{m}_t = (1 - \beta)\sum_{k=1}^{t}\beta^{t-k}\mathbf{g}_k$ , a weighted sum of past gradients whose weights decay geometrically. Many libraries use instead $\mathbf{u}_t = \beta\,\mathbf{u}_{t-1} + \mathbf{g}_t$ and $\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta'\mathbf{u}_t$ . Since $\mathbf{m}_t = (1 - \beta)\mathbf{u}_t$ , the two produce the same iterates when $\eta' = (1 - \beta)\eta$ , so a learning rate means something only together with its convention. The vector $\mathbf{m}$ is part of the state of the run.

In plain words Momentum keeps a running direction. Gradients that keep pointing the same way build up, and gradients that flip sign from step to step largely cancel. In a long narrow valley this damps the zigzag across the valley and keeps progress along it. The memory can also carry the parameters past a minimum, so $\beta$ and $\eta$ are tuned together on validation data.

Example With $\beta = 0.9$ and a constant gradient $1$ , the running average is $0.1$ , $0.19$ , $0.271$ and in general $1 - 0.9^t$ , which approaches $1$ . With gradients alternating $1, -1, 1, \dots$ it is $0.1$ , $-0.01$ , $0.091$ , $-0.0181$ and settles into swinging between about $\pm(1 - \beta)/(1 + \beta) \approx \pm 0.0526$ . The consistent signal passes through at full size and the oscillating one shrinks by a factor of $19$ . Under the library convention a constant gradient $1$ builds $\mathbf{u}_t$ up toward $1/(1 - \beta) = 10$ , which is why the two conventions need different learning rates.

Adam optimizer

Definition Adam has a learning rate $\eta > 0$ , decay rates $\beta_1$ and $\beta_2$ in $[0, 1)$ and a small constant $\varepsilon > 0$ , with common defaults $\eta = 0.001$ , $\beta_1 = 0.9$ , $\beta_2 = 0.999$ and $\varepsilon = 10^{-8}$ . Starting from $\mathbf{m}_0 = \mathbf{v}_0 = \mathbf{0}$ , step $t = 1, 2, \dots$ takes the gradient $\mathbf{g}_t$ at $\boldsymbol{\theta}_{t-1}$ and computes

\begin{aligned} \mathbf{m}_t &= \beta_1\mathbf{m}_{t-1} + (1 - \beta_1)\,\mathbf{g}_t, \\[0.4em] \mathbf{v}_t &= \beta_2\mathbf{v}_{t-1} + (1 - \beta_2)\,\mathbf{g}_t^2, \\[0.4em] \hat{\mathbf{m}}_t &= \mathbf{m}_t / (1 - \beta_1^t), \\[0.4em] \hat{\mathbf{v}}_t &= \mathbf{v}_t / (1 - \beta_2^t), \end{aligned}

and then

\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta\, \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \varepsilon} .

The square, the square root and the division act entry by entry, so every parameter gets its own effective step size. The second moment $\mathbf{v}_t$ averages squared gradients, which is not the square of the averaged gradient. The bias corrections undo the pull of the zero starting states, and for a constant gradient $\mathbf{g}$ they give exactly $\hat{\mathbf{m}}_t = \mathbf{g}$ and $\hat{\mathbf{v}}_t = \mathbf{g}^2$ . The vectors $\mathbf{m}_t$ and $\mathbf{v}_t$ and the counter $t$ are part of the state of the run. Adam guarantees neither convergence nor good accuracy on new data, and simple convex problems are known on which it does not converge to the minimizer.

In plain words Adam combines momentum with a separate scale for every parameter. Dividing the averaged gradient by the root of the averaged squared gradient makes typical steps about $\eta$ in size, whatever the scale of each parameter’s gradient, which helps when gradients differ greatly between layers. It is a common default for networks, but it still needs a tuned learning rate and validation checks.

Example At the first step $\hat{\mathbf{m}}_1 = \mathbf{g}_1$ and $\hat{\mathbf{v}}_1 = \mathbf{g}_1^2$ , so each entry moves by $-\eta g/(\lvert g \rvert + \varepsilon)$ , almost exactly $\eta$ against the sign of its gradient. With $\eta = 0.01$ , the default decay rates and a gradient of $2$ , the raw states are $m_1 = 0.2$ and $v_1 = 0.004$ , the corrected ones are $2$ and $4$ , and the parameter moves by $-0.01 \cdot 2/(2 + 10^{-8}) \approx -0.01$ . A gradient of $0.001$ moves another parameter by about $-0.0099990$ , nearly the same, while a gradient of $10^{-9}$ , smaller than $\varepsilon$ , moves it by only about $-0.00091$ . The state matters when a run resumes. If the second gradient of the first parameter is $-2$ , the saved state gives $\hat{m}_2 = -0.02/0.19 \approx -0.1053$ and $\hat{v}_2 = 4$ , a step of about $+0.00053$ . A run restarted with fresh zero states would instead step by $+0.01$ , about nineteen times as far.

Batch normalization

Definition Let a layer’s values for a mini-batch of $m \ge 2$ examples form an array with entry $h_{ij}$ for example $i$ and feature $j$ . In training mode batch normalization computes, for each feature $j$ over the batch axis,

\begin{aligned} \mu_j &= \frac{1}{m}\sum_{i=1}^{m} h_{ij}, \\[0.4em] \sigma_j^2 &= \frac{1}{m}\sum_{i=1}^{m} (h_{ij} - \mu_j)^2, \end{aligned}

and outputs

y_{ij} = \gamma_j\, \frac{h_{ij} - \mu_j}{\sqrt{\sigma_j^2 + \varepsilon}} + \beta_j ,

where $\varepsilon > 0$ is a small constant and the scale $\gamma_j$ and the shift $\beta_j$ are learned parameters, usually started at $1$ and $0$ . This shift has nothing to do with the decay rates of an optimizer. The backward pass differentiates through $\mu_j$ and $\sigma_j^2$ , so each example’s gradient depends on the other examples of its batch. Training also keeps running estimates of each feature’s mean and variance, typically exponential moving averages of the batch statistics. In evaluation mode the layer uses those stored estimates instead of the current batch, so it applies a fixed affine map to each feature and a prediction depends only on its own input. Libraries differ in details such as the averaging rate and whether the running variance divides by $m$ or $m - 1$ .

In plain words Batch normalization rescales each feature to mean about $0$ and variance about $1$ within the batch and then lets the network learn a new scale and shift. It often permits larger learning rates and makes training less sensitive to the initialization. Its two modes really compute different functions, so validation and prediction must run in evaluation mode, and the running estimates should come from training batches only.

Example One feature takes the values $1$ , $2$ , $3$ and $6$ in a batch of four. Then $\mu = 3$ and $\sigma^2 = (4 + 1 + 0 + 9)/4 = 3.5$ , and with $\gamma = 1$ , $\beta = 0$ and $\varepsilon = 10^{-5}$ the outputs are about $-1.069$ , $-0.535$ , $0$ and $1.604$ . In a batch with the values $3$ , $4$ , $5$ and $6$ the same input value $3$ becomes about $-1.342$ , so in training mode an example’s output changes with its batch mates. If the stored estimates were mean $3$ and variance $3.5$ , evaluation mode would map the value $3$ to $0$ in every batch. A training-mode batch of one example has $\sigma^2 = 0$ , so every output equals $\beta_j$ whatever the input, and the layer erases the information, which is why batch normalization needs batches of at least two and works poorly with very small ones.

Check yourself

Adam starts with $\eta = 0.01$ , $\beta_1 = 0.9$ , $\beta_2 = 0.999$ and $\varepsilon = 10^{-8}$ , and the first gradient of a parameter is $-3$ . Find the raw and the corrected moments and the change of the parameter. Compare with plain gradient descent at the same rate, and with a first gradient of $-0.03$ .

Show the answer

The raw moments are $m_1 = 0.1 \cdot (-3) = -0.3$ and $v_1 = 0.001 \cdot 9 = 0.009$ . The corrections divide by $1 - 0.9 = 0.1$ and $1 - 0.999 = 0.001$ , giving $\hat{m}_1 = -3$ and $\hat{v}_1 = 9$ . The update is $-0.01 \cdot (-3)/(3 + 10^{-8}) \approx +0.01$ , so the parameter increases by about $0.01$ . Moving against a negative derivative raises the parameter and lowers the loss to first order.

Plain gradient descent at rate $0.01$ would move the parameter by $+0.03$ . With a first gradient of $-0.03$ , Adam still moves it by about $+0.01$ , while plain gradient descent would move it by only $+0.0003$ . The first Adam step depends on the sign of the gradient and hardly on its size, as long as the gradient is much larger than $\varepsilon$ .

A model has one parameter and the training objective $\tfrac12 \theta^2$ , whose gradient is $\theta$ . Gradient descent starts at $\theta = 1$ with rate $0.1$ . The framework adds each new gradient into a stored buffer, and the loop never clears it. Compute the first three updates with and without the bug, and describe the trace of the loss.

Show the answer

The correct loop multiplies $\theta$ by $0.9$ at every step, giving $0.9$ , $0.81$ and $0.729$ , and after $50$ steps about $0.0052$ . With the bug the update uses the sum of all gradients so far. The first step uses $1$ and gives $0.9$ . The buffer then holds $1 + 0.9 = 1.9$ , so $\theta = 0.9 - 0.19 = 0.71$ . Next it holds $1.9 + 0.71 = 2.61$ , so $\theta = 0.71 - 0.261 = 0.449$ .

The values continue $0.1431$ , $-0.1771$ , $-0.4796$ and keep swinging between about $-1.01$ and $1.01$ without settling. Subtracting consecutive updates gives $\theta_{t+1} = 1.9\,\theta_t - \theta_{t-1}$ , whose characteristic roots are complex with modulus $1$ , so the oscillation never decays, and the loss keeps climbing back to about $0.5$ . A loss that falls at first and then oscillates or rises at a learning rate that should be safe is a typical sign of this bug. The fix is to clear the gradients before every backward pass, or, when several batches are accumulated on purpose, to scale the loss by the number of batches, update once and then clear.

A network with dropout and batch normalization is trained by a loop with the following evaluation step. After each epoch the code turns off gradient recording and computes the validation loss with the model still in training mode. It also computes the test loss, and whenever the test loss is the lowest so far it stores best = params, where params is the list of the live parameter arrays. Find the problems and fix them.

Show the answer

Validation in training mode is wrong. Dropout randomly removes units, so the validation loss is noisy and is not the loss of the model that will be used. Batch normalization uses the statistics of each validation batch instead of the stored ones, and in common libraries a training-mode pass also updates the running estimates, so validation data leak into the model. Switch to evaluation mode for validation and back to training mode afterward. Turning off gradient recording only saves memory and time and does not change how these layers behave, so both settings are needed.

Choosing the checkpoint by test loss makes the test set part of model selection, and its loss is then an optimistic estimate. Select the checkpoint by validation loss and evaluate the test set once, at the end.

Storing best = params keeps references to the arrays that later updates keep changing, so the saved model silently becomes the final model. Save copies of the parameters, together with the epoch and, if training may resume, the optimizer state.

Training Neural Networks

Key terms

Training loop

Weight initialization

Vanishing and exploding gradients

Momentum

Adam optimizer

Batch normalization

Check yourself

How this lesson connects

Builds on

Used later in