Chapter 4: Neural Networks: Loss and Optimization

How a network learns

A neural network starts with random weights, so its initial predictions are useless. Learning means adjusting those weights until the predictions match the training data. To adjust anything, we need a number that tells us how wrong the network is: a loss function.

Think of the loss as the height of a landscape. The position on the map is the set of all weights, and the elevation is the loss. Training means walking downhill until we reach a low point. Gradient descent is the algorithm that tells us which direction is downhill.

Loss functions

The choice of loss depends on the task. For regression, where the model predicts a continuous value, mean squared error is common. It penalizes large errors more than small ones because the error is squared.

For classification, and especially for language modeling, cross-entropy loss is standard. It compares the model's predicted probability distribution to the true distribution. When the target is one-hot, only the log-probability of the correct class matters.

If the model is confident and correct, the loss is near zero. If it is confident and wrong, the loss grows large. This steep penalty for overconfidence is why cross-entropy is so effective for classification.

Convex and non-convex landscapes

A convex loss landscape has a single global minimum shaped like a bowl. Linear regression with MSE is convex, so gradient descent cannot get stuck in a bad place. Neural networks, however, have non-convex loss surfaces with many local minima, saddle points, and flat regions.

In practice, a deep network has so many parameters that we rarely find the global minimum. The good news is that many local minima are good enough. Modern networks are over-parameterized, meaning they have many more weights than training examples, which creates large, flat basins where the loss is low. Saddle points are often more problematic than local minima because the gradient can be very small in every direction.

Softmax and cross-entropy together

In a classifier, the last layer produces a vector of raw scores called logits. Softmax turns those logits into probabilities. Cross-entropy then compares those probabilities to the one-hot target. The two are usually written as one combined loss.

The log-softmax trick is common in code. Because softmax uses exponentials, large logits can overflow. Subtracting the maximum logit before exponentiating keeps the values stable without changing the result.

import numpy as np

def softmax(z):
    e = np.exp(z - np.max(z))
    return e / e.sum()

def cross_entropy(y_hat, y):
    # y is one-hot, y_hat is softmax probabilities
    return -np.sum(y * np.log(y_hat + 1e-9))

logits = np.array([2.0, 1.0, 0.1])
probs = softmax(logits)
print(probs)

y = np.array([1.0, 0.0, 0.0])
print("loss:", cross_entropy(probs, y))

Run the code and notice how softmax emphasizes the largest logit while still giving small probabilities to the other classes. The loss only cares about the probability assigned to the true class.

Temperature and softmax sharpness

Dividing the logits by a temperature T before softmax controls how sharp the distribution is. A high temperature makes the probabilities more uniform; a low temperature makes the highest logit more dominant. Temperature is central to text generation, where sampling from the softmax output lets the model produce diverse text.

Gradient descent

To minimize the loss, we compute the gradient of the loss with respect to each weight. The gradient points in the direction that increases the loss, so we move the weights a small step in the opposite direction. Repeating this on many examples is gradient descent.

The learning rate eta controls the step size. Too small and training crawls; too large and the optimizer overshoots the minimum or even diverges. Picking and scheduling the learning rate is one of the most important practical decisions in training.

import numpy as np

# Gradient descent on a simple quadratic f(x) = (x - 3)^2
x = 0.0
lr = 0.1
for step in range(20):
    grad = 2 * (x - 3)
    x = x - lr * grad
    print(f"step {step}: x = {x:.3f}")

Momentum and adaptive optimizers

Plain SGD is simple but can be slow in narrow valleys where the gradient oscillates. Momentum keeps a running average of past gradients, allowing the optimizer to build speed in consistent directions and dampen oscillations.

Adam extends this idea by maintaining both a moving average of gradients and a moving average of squared gradients. It uses the first to estimate the direction and the second to normalize the step size per parameter. Adam is the default optimizer for many language model training runs because it works well with little tuning.

Stochastic and mini-batch gradient descent

In practice, we rarely use the full dataset at once. Stochastic gradient descent uses one example at a time, and mini-batch SGD uses a small batch. Batches give a noisy but useful estimate of the true gradient and allow hardware parallelism.

Mini-batch SGD is the workhorse of modern deep learning. It balances the accurate gradient of full-batch descent with the speed of single-example updates. The batch size also affects generalization: very large batches can converge faster per step but sometimes find sharper minima that generalize worse.