Training a Neural Network

The previous article introduced two classifiers for linear classification: the support vector machine (SVM) and softmax. An SVM computes a score for each class with an affine function $f$ and tries to separate the correct class's score from the others by at least a specified margin $\Delta$. Softmax uses an exponential normalization to turn those scores into normalized probabilities and minimizes their cross-entropy, or equivalently the relevant KL-divergence term. Ignoring regularization, the SVM stops contributing loss once its margin condition is satisfied, whereas softmax can continue improving the probability assigned to the correct class.

We will now look at how a loss function actually optimizes parameters. A course in optimization theory makes this material easier to understand. Deep-learning introductions tend to cover only gradient-based methods, which can make terms such as penalty term, projection, Jacobian, or Hessian startling when they appear in a paper. WGAN, one of the best-known generative-model papers, relies on ideas from linear programming and duality that ordinary deep-learning knowledge alone does not fully explain. For a rigorous understanding, I recommend also studying probability, random variables, and optimization theory.

Two concepts form the starting point for neural-network training:

1. A score function, such as $f$ in our example, is parameterized and maps raw image pixels to a score for each class.
2. A loss function measures how closely the scores produced by the current parameters meet the desired criterion. We previously saw the SVM hinge loss and softmax cross-entropy loss.

$$\begin{aligned} L_i =& \sum_{j \neq y_i} \max \left( 0, W_j^\top x_i - W_{y_i}^\top x_i + \Delta \right) \newline L_i =& -\log \left( \frac{e^{f_{y_i}}}{\sum_j e^{f_j}} \right) = -f_{y_i} + \log \sum_j e^{f_j} \end{aligned}$$

These are the SVM hinge loss and softmax cross-entropy loss, respectively. We define a score function and a loss for each classifier so that the current parameters $(W,b)$ tell us how far the result is from our target, after which we can search for a better score function $f$. Optimizing the weights $W$ and bias $b$ that constitute $f$ is an optimization problem. The following sections introduce the most fundamental optimization method used in neural networks and deep learning.

Optimization

If we want a loss function to attain its minimum, the easiest situation is one in which we know the shape of that function. Recall finding the minimum or maximum of a quadratic in high-school mathematics. Quadratics are representative convex or concave functions. Such problems have a unique optimum, and that optimum lies at the vertex where the gradient is zero.

Our loss would be easy to optimize if it had the same form. Unfortunately, with high-dimensional data such as images and text embeddings, we cannot assume a simple analytic shape. One common way to analyze a non-analytic function locally is to impose constraints on feasible directions.

Suppose an unknown function $L(\cdot)$ has value $L(W)$ at the current weight $W$. We do not know which direction leads to a global optimum, but we can define a direction vector $W_1$ with the same dimensionality as $W$ and, for step size $\alpha$, examine the loss along the line $L(W+\alpha W_1)$.

The plot on the left illustrates the value of the loss when moving in one particular direction from $L(W)$. A feasible direction does not reveal the entire function, but it does let us inspect the loss within a particular range. Adding another direction yields $L(W+\alpha W_1+\beta W_2)$. Two observable axes let us visualize the surface as a color map, like the center and right plots. Regions of equal color share the same function value and form contour lines. The gradient axis points in the direction in which those contours change, perpendicular to the contour line.

Consider the SVM hinge loss:

$$L_i = \sum_{j \neq y_i} \max \left( 0, W_j^\top x_i - W_{y_i}^\top x_i + 1 \right)$$

With $\Delta=1$, suppose the dataset consists of three one-dimensional samples. Their losses are

$$\begin{aligned} L_0 =& \sum_{j \neq y_0} \max \left( 0, W_j^\top x_0 - W_{y_0}^\top x_0 + 1 \right) \newline L_1 =& \sum_{j \neq y_1} \max \left( 0, W_j^\top x_1 - W_{y_1}^\top x_1 + 1 \right) \newline L_2 =& \sum_{j \neq y_2} \max \left( 0, W_j^\top x_2 - W_{y_2}^\top x_2 + 1 \right). \end{aligned}$$

Ignoring regularization for simplicity, the SVM uses the mean over all samples:

$$L = \frac{L_0 + L_1 + L_2}{3}$$

Because these are three samples made of one-dimensional points, $x_i$ and $w_j$ behave as scalars. Taking weight as the horizontal axis and each loss $(L_0,L_1,L_2)$ as the vertical axis gives the following graph.

For each of the three losses, every $W$ on the flat floor is a minimizer. This suggests that the hinge loss may remain convex in higher dimensions. In practice, however, the full loss landscape over high-dimensional data is too poorly known for us simply to assume and solve a convex problem. Methods that explicitly use a function's gradient and Hessian also require the loss to be differentiable with respect to every weight, which is itself a constraint.

Optimization Methods

Random Search

We need some other method to find weights that minimize the loss. The first idea is brute force: insert a random $W$, calculate $L(W)$, and keep the new weight whenever its loss is lower than the best one seen so far. The concept alone shows why optimization is effectively impossible this way. Real-valued weights have nearly infinitely many possible values, so the probability of guessing an optimum at random is essentially zero.

Instead, we begin with random weights—initialization—and ask in which direction to change them to improve performance. Finding the global optimum in one step is impossible, but finding a weight better than the current one is feasible. The usual analogy is a blindfolded hiker descending a mountain. The hiker does not know where the bottom of the mountain is, but by moving a little from the current position can determine which direction goes downhill.

Random Local Search

One method tries every feasible direction from the current point, moves in whichever one produces the smallest loss, and repeats. This is still inefficient because every step requires evaluating the function in every direction. The hiker is not merely blindfolded; this hiker cannot feel the slope either.

Following the Gradient

We do not need to explore every direction to learn which one changes the function most at our current position. We can calculate the gradient. Informally, a derivative is the rate of change or slope at a point. The gradient generalizes this idea and gives the direction in which the function value increases most rapidly at that point.

$$\frac{df(x)}{dx} = \lim_{h \rightarrow 0} \frac{f(x+h) - f(x)}{h}$$

The expression above defines a derivative. A gradient extends it to multiple dimensions:

$$\nabla f(x) = \begin{bmatrix} \frac{\partial f(x)}{\partial x_1} \newline \frac{\partial f(x)}{\partial x_2} \newline \vdots \newline \frac{\partial f(x)}{\partial x_n} \end{bmatrix}$$

We can estimate a numerical gradient from a small $\delta$, following the derivative definition, or derive an analytic gradient when the function is differentiable:

$$\begin{aligned} \frac{df(x)}{dx} \approx& \frac{f(x+\delta) - f(x)}{\delta} \newline \frac{df(x)}{dx} \approx& \frac{f(x+\delta/2) - f(x-\delta/2)}{\delta}. \end{aligned}$$

The second, symmetric estimate usually approximates the tangent slope more accurately. Because the gradient points toward the greatest increase, moving a negative amount along it decreases the function:

$$W_{i+1} = W_i - \alpha \nabla_W L(W_i)$$

The distance we move against the gradient is a hyperparameter, the step size $\alpha$. If it is too small, convergence is extremely slow. If it is too large, optimization can fail to converge and diverge instead. A step that is too large overshoots: it may oscillate across high points without descending or even move to a higher part of the surface.

Choosing an appropriate step size is therefore important. The hinge and cross-entropy losses defined earlier are differentiable with respect to $W$ almost everywhere, so their gradients can be written in analytic form.

For the hinge loss of sample $i$,

$$L_i = \sum_{j \neq y_i} \max \left( 0, W_j^\top x_i - W_{y_i}^\top x_i + 1 \right),$$

the partial derivative with respect to $W_{y_i}$, the row that scores the correct class, is

$$\nabla_{w_{y_i}} = -\sum_{j \neq y_i} \mathbb{1} \left(W_j^\top x_i - W_{y_i}^\top x_i + 1 > 0 \right) x_i.$$

Strictly speaking, the nondifferentiable corner in $\max$ means we use a subgradient. The loss compares how far every incorrect class's score is from the correct class's score. Whenever the violation is greater than zero,

$$\max \left(0,W_j^\top x_i-W_{y_i}^\top x_i+1\right) =W_j^\top x_i-W_{y_i}^\top x_i+1,$$

and the contribution to the correct-class row is

$$\nabla_{w_{y_i}}^j = \begin{cases} -x_i, & W_j^\top x_i - W_{y_i}^\top x_i + 1 > 0 \newline 0, & W_j^\top x_i - W_{y_i}^\top x_i + 1 \leq 0. \end{cases}$$

Summing these terms over every $j$ gives the analytic expression above. For an incorrect-class row $W_j$, the sign is reversed:

$$\nabla_{w_j} = \begin{cases} x_i, & W_j^\top x_i - W_{y_i}^\top x_i + 1 > 0 \newline 0, & W_j^\top x_i - W_{y_i}^\top x_i + 1 \leq 0. \end{cases}$$

Stochastic and Mini-Batch Gradient Descent

Calculating a gradient separately for every sample is expensive. With a heavy or complicated model, processing samples one by one is inefficient, so we often divide them into groups and use mini-batch gradient descent, which can be viewed as a practical variant of stochastic gradient descent (SGD).

“Stochastic” denotes a random process. If calculating the loss over all $N$ samples at once is difficult, we could draw one sample at a time without replacement during training and calculate only that sample's gradient. Treating the complete dataset as one batch gives a gradient that fully reflects the dataset distribution and, under convex assumptions, points in the ideal optimization direction. For a dataset containing a million images, however, calculating the entire loss at once is nearly impossible because of memory and speed constraints.

Pure one-sample SGD avoids that cost but produces noisy weight updates and can be inefficient relative to optimization over the full distribution. Mini-batch optimization is the compromise. Instead of drawing one sample, it draws $K$ samples, using that subset to estimate the shape of the whole dataset's population distribution.

This works because examples in a training set have some correlation. In an extreme thought experiment, suppose a set of 1.2 million images were made by copying the same 1,000 images. Computing all of them would repeat identical gradients 1,200 times. Real datasets are not built this way, but because samples remain related, a mini-batch gradient can approximate the broad structure of the full dataset without serious loss.

Conclusion

For softmax and SVM classifiers, weights $W$ combine with input $x$ to produce scores, and the difference between those scores and the true label $y$ produces a data loss. Adding a regularization loss on $W$ gives the full training objective. We optimize that objective by using an analytically calculated gradient to update the parameters in the direction that reduces it. Finally, rather than calculate the loss gradient over every sample at once—a costly procedure with its own practical drawbacks—we usually train with mini-batch gradient descent.