20  Fitting Neural Networks in Practice

As we’ve discussed, training a neural network is optimization, but of a certain kind. Unlike most mathematical optimization problems where we seek the true minimizer of a function, in deep learning the loss function is a surrogate: we care about the loss on the training data, but we care more about the loss when we generalize to new data. And the optimal training loss may not yield the optimal generalization error.

At the same time, neural networks are complex, nonlinear, nonconvex functions with many parameters. Optimizing them poses all kinds of challenges for our simple gradient-based methods.

Let’s discuss strategies for achieving better optimization for neural networks, keeping generalization error in mind as the final goal.

20.1 Common problems

To understand optimization, it may be useful to introduce some math.

Definition 20.1 (Hessian matrix) The Hessian matrix \(H\) of a function \(f(x) : \R^n \to \R\) is defined so that its \((i, j)\)th entry evaluated at a point \(x \in \R^n\) is \[ H(f)(x)_{i,j} = \frac{\partial^2}{\partial x_i \partial x_j} f(x). \] When the second partial derivatives are continuous, the order of differentiation can be swapped and the Hessian matrix is symmetric.

You may be used to thinking of the Hessian in one dimension: for a function \(f(x) : \R \to \R\), the Hessian is simply the second derivative. You use the second derivative to evaluate the function’s curvature at specific points. Similarly, the Hessian tells the curvature of a multidimensional function.

In univariate functions, you could use the Taylor expansion around a particular point \(x^*\) to approximate the function. A second-order Taylor expansion would be \[ f(x) \approx f(x^*) + (x - x^*) f'(x^*) + \frac{(x - x^*)^2}{2} f''(x^*). \] For a multivariate function, the Taylor expansion uses the gradient and the Hessian: \[ f(x) \approx f(x^*) + (x - x^*)\T \nabla f(x^*) + \frac{1}{2} (x - x^*)\T H(f)(x^*) (x - x^*). \]

This decomposition suggests that understanding the Hessian of a function will help us understand its shape. Because (for most functions) the Hessian is real and symmetric, it can be decomposed into real eigenvalues and eigenvectors. If all eigenvalues are positive, the function is convex at this point; if they are all negative, it is concave.

(TODO show why this is, or maybe have a homework exercise?)

We can relate the Hessian to the steps of gradient descent. Suppose that at a particular gradient descent step, we are at the point \(x^{(0)}\). We calculate the gradient and update the parameters to \(x^{(1)} = x^{(0)} - \gamma \nabla f(x^{(0)})\). In the Taylor approximation, the function’s value at the new point would be \[\begin{align*} f(x^{(1)}) &\approx f(x^{(0)}) + (x^{(1)} - x^{(0)})\T \nabla f(x^{(0)}) + \frac{1}{2} (x^{(1)} - x^{(0)})\T H(f)(x^{(0)}) (x^{(1)} - x^{(0)})\\ &\approx f(x^{(0)}) - \gamma \nabla f(x^{(0)})\T \nabla f(x^{(0)}) + \frac{\gamma^2}{2} \nabla f(x^{(0)})\T H(f)(x^{(0)}) \nabla f(x^{(0)}). \end{align*}\]

Let’s consider the implications of this for optimization. If \(f\) is the loss function for a particular network with a particular training dataset, and \(x^{(0)}\) is the vector of weights and biases, each gradient step will reduce the loss by an amount related to the gradient and the Hessian. If the Hessian term is smaller than the gradient term, the loss will decrease; but if the Hessian term is large, the gradient step may not decrease the loss at all.

20.1.1 Poor conditioning

That brings us to our first problem: poor conditioning of the Hessian matrix. By “poor conditioning” we mean it has a large ratio between its largest and smallest eigenvalues. If the Hessian is poorly conditioned at a point, it means that in one direction (the direction of the eigenvector corresponding to the largest eigenvalue), the gradient will change very quickly, while in another direction (corresponding to the smallest eigenvalue), the gradient will change very slowly. A smart optimizer would move in a direction where the gradient is consistent, so we consistently increase the loss; but gradient descent goes wherever looks best now, even if after a step or two the loss would be quite flat.

Complicated neural networks often have regions of poor conditioning. The visible symptom is that the loss increases between gradient descent steps, even if you set the learning rate \(\gamma\) to be small; to get the loss to decrease, you have to set it to very small values and wait ever-longer times for gradient descent to finish.

One way to detect this it to plot the norm of the gradient vector after each step: calculate \(\|\nabla f(x^{(i)})\|_2^2\) and track it along with the loss. If the gradient keeps increasing, but the loss does not decrease very fast, gradient descent is making large jumps that do not effectively reduce the loss. You may need to pick smaller learning rates.

20.1.2 Multiple local minima

Nontrivial neural networks almost always are not identifiable. An identifiable model \(f(x; \theta)\) is one where if \(f(x; \theta_1) = f(x; \theta_2)\) for all \(x\), then \(\theta_1 = \theta_2\): the shape of the function is uniquely determined by the parameters.

In a simple fully connected feed-forward network, the network is not identifiable because we can swap nodes within each hidden layer. By rearranging the weights and biases accordingly, the model will give identical output for any input.

Similarly, consider node \(i\) in a fully connected layer \(k\) using ReLU as its activation function. Its output vector is \[ (w_i\T h^{(k-1)} + b_i)_+ \] If we’re in the region where this is positive, we can multiply \(w_i\) by any constant \(c\) if we also multiple the weights in the previous layer by \(1/c\) (which would cause \(h^{(k-1)}\) to be a factor of \(c\) smaller). There is hence a hyperbola of weight values that give identical output.

If gradient descent reaches such a local minimum, the gradient will go to zero (or nearly zero, since we might land just outside the minimum) and stay there. Of course, in these nonidentifiable regions, every point has the same training loss; as long as the region is a minimum, it does not matter where in the region we land.

Because neural networks have many parameters, there may be many other local minima not arising from nonidentifiability. Normally this is a problem for optimization because you can end up in a local minimum with much larger loss than the true global minimum. But it is not clear how bad this is in large neural networks: it seems possible that most local minima have low values of the loss function, and so any one of them will do well (Goodfellow, Bengio, and Courville 2016, sec. 8.2.2).

20.1.3 Saddle points and plateaus

Functions can have points with no gradient that are not minima or maxima: saddle points. (In fact, for high-dimensional functions like neural networks, saddle points can be more common than minima and maxima.)

In principle, gradient descent could get stuck at a saddle point, and this could be detected by checking the Hessian matrix: if it has both positive and negative eigenvalues, this is a saddle point. In practice, it’s unlikely for a gradient descent step to land on a point where the gradient is exactly zero, so it is unclear how big of a problem this actually is.

More concerning are large regions where the loss function is flat. These “plateaus” may not be minima, and both the gradient and Hessian would be zero, preventing gradient descent from making progress. Again, tracking the norms of the gradient and Hessian during optimization may be useful to detect this.

20.1.4 Exploding and vanishing gradients

Loss functions may have very steep “cliffs”. (For example, in a network where the output involves the multiplication of many weights, the gradient with respect to those weights may be very steep, because the multiplication of many terms can change dramatically as those terms change.) If gradient descent lands in such a region, it can make a dramatic jump.

This problem can again be detected by tracking the norm of the gradient at each step: if it jumps to a huge value on one step, there may be a cliff.

One simple approach to avoid this is gradient clipping, which simply does not allow the gradients to be larger than a preset maximum value. In PyTorch, the torch.nn.utils.clip_grad_norm_ function goes through all parameters in a network, calculates the norm of the gradient, and rescales the gradient so the norm is no larger than a set value:

for data, labels in train_loader:
    # reset gradients
    optimizer.zero_grad()

    # get predictions and loss
    output = network(data)
    loss = F.nll_loss(output, labels)

    # calculate gradients
    loss.backward()

    # rescale all the gradients to have maximum norm 0.5
    torch.nn.utils.clip_grad_norm_(network.parameters(), 0.5)

    # perform the gradient descent step with the clipped gradients
    optimizer.step()

A different kind of explosion can occur when networks become very deep. (This will particularly be a problem in recurrent networks, which we will see later.)

Consider, for example, a deep feed-forward network using ReLU. At each layer, the activations from the previous layer are multiplied by a weight matrix, so after many layers the output has been multiplied by many times. The gradient of such a long product is itself a product of those weight matrices—and so if they have large eigenvalues, the gradient may get huge; if they have small eigenvalues, it may get very close to 0. This causes either exploding gradients or vanishing gradients.

Feed-forward networks use different weight matrices for every layer, so it’s unlikely they will all have large or small eigenvalues and cause an explosion. But recurrent networks use the same weight matrix many times, and so can be vulnerable.

20.1.5 Dead units

TODO ReLU dead units

20.2 Practical algorithms

Given all these potential problems, what optimization methods work well in practice?

20.2.1 Stochastic gradient descent

The most common optimizer is stochastic gradient descent, as introduced in Section 19.4.1, with a reasonable minibatch size. But usually we do not use a fixed learning rate \(\gamma\). Instead, we choose a sequence \(\gamma_k\) of learning rates such that \(\gamma_k \to 0\).

Ordinary gradient descent would not need this, because when we reach the minimum, the gradient would be zero and the optimizer would stop. But in SGD, we are estimating the gradient from a sample, so the gradient may not be exactly zero, and we’d continue going. According to Goodfellow, Bengio, and Courville (2016) (section 8.3.1), a sufficient condition for convergence is to choose the learning rates so that \[\begin{align*} \sum_{k=1}^\infty \gamma_k &= \infty \\ \sum_{k=1}^\infty \gamma_k^2 &< \infty. \end{align*}\] Goodfellow, Bengio, and Courville (2016) reports that

The learning rate may be chosen by trial and error, but it is usually best to choose it by monitoring learning curves that plot the objective function as a function of time. This is more of an art than a science, and most guidance on this subject should be regarded with some skepticism.

As we will see in Section 20.2.3, it’s also common to use adaptive strategies that update the learning rate automatically, to avoid needing so much hand-tuning.

20.2.2 Momentum

Another strategy to improve stochastic gradient descent is to add momentum. The analogy is to think of gradient descent like a small ball rolling downhill on the surface of the loss function; a very light ball will always roll directly downhill, but a heavy one can pick up speed and continue moving in a direction even as the loss changes slope, because its momentum keeps it moving.

Using momentum requires an additional tuning parameter \(\alpha \in [0, 1)\). We also need a velocity vector \(v\), which we can think of as the current velocity of the ball.

On each gradient descent step, we calculate the gradient estimate over the minibatch. Call the resulting gradient vector \(g\).

We now update the parameters \(\theta\) using the rule \[\begin{align*} v &\leftarrow \alpha v - \gamma g\\ \theta &\leftarrow \theta + v. \end{align*}\] At each step, we move by one unit in the direction of the velocity; the velocity depends on the gradient but also the previous velocity.

If you expand this out to write \(v\) in terms of the gradient of the past several steps, you will see it is essentially an exponentially weighted moving average of the previous gradients. When \(\alpha \approx 0\), the previous gradients get only slight weight; when \(\alpha \approx 1\), the previous gradients get high weight.

As another analogy, you can think of this as smoothing the gradient vectors over time. With the right amount of smoothing, we can reduce problems like sudden gradient cliffs; and if stochastic gradient descent would want to bounce back and forth around a minimum, using the average gradient would result in the bounces averaging out and SGD taking much smaller steps.

In PyTorch, we can tell SGD to use momentum by using the momentum parameter to set \(\alpha\):

optimizer = torch.optim.SGD(network.parameters(), lr=0.001, momentum=0.25)

There are other varieties of momentum that can be used, such as Nesterov momentum, but we won’t get into them here.

20.2.3 Adaptive learning rates

Because the learning rate is so important but can also be so difficult to tune, there are variety of methods for adjusting it automatically.

The basic idea is simple. If the gradient is small, don’t change the learning rate much; if the gradient is large, shrink it. But the details are complex. We’d like to apply this rule to each parameter separately—since some may have large gradients and some may have small—and we’d like to do so in an intelligent way that accounts for the randomness in stochastic gradient descent, so we aren’t making rapid wild changes to the learning rate.

RMSProp is one such algorithm (introduced by Geoffrey Hinton in a Coursera course, oddly enough). We start with a learning rate \(\gamma\), a decay rate \(\rho\), and a value \(\delta = 10^{-6}\). We initialize a vector \(r = 0\) with one entry per parameter.

On each gradient descent step, we:

  1. Calculate the gradient \(g\) with respect to the minibatch.
  2. Update \(r \leftarrow \rho r + (1 - \rho) g^2\), where \(g^2\) is calculated elementwise in the gradient vector. Hence \(r\) will be an exponentially weighted moving average of the squared gradients.
  3. Update the parameters as \(\theta \leftarrow \theta - \frac{\gamma}{\sqrt{\delta + r}} g\), where again the division and multiplication in the second term is done elementwise.

The effective learning rate for a particular parameter \(\theta_i\) is hence \[ \frac{\gamma}{\sqrt{\delta + r_i}}. \] If the parameter’s gradients were large (positive or negative) in recent steps, its learning rate will be smaller; if they were small, the learning rate will be larger. You can see that \(\delta\) is necessary to avoid dividing by zero in rare cases.

An improved version of RMSprop is Adam (Kingma and Ba 2015). Adam is like RMSProp with momentum, and so the learning rate adaptation is applied to the velocity vector. (For details, see also Goodfellow, Bengio, and Courville (2016), algorithm 8.7.) Momentum helps keep gradient descent from getting stuck or from being too affected by cliffs, while the adaptation helps it converge when it gets near a minimum. The result is that Adam is one of the most popular current methods for training neural networks. It works well without usually requiring much manual tuning of hyperparameters.

Adam is supported by PyTorch. When initializing the optimizer, instead of using SGD() you can use Adam:

optimizer = torch.optim.Adam(network.parameters())

It uses default values for the learning rate and momentum parameters, and these can be tweaked as needed.