Training a Model: Gradients, Optimizers, and Learning Rates

Training turns examples and an objective into parameter updates. A forward pass computes predictions, the loss measures their disagreement with targets, backpropagation computes derivatives, and an optimizer decides how to use those derivatives. These are separate jobs: changing Adam to SGD does not change what the labels mean or repair a mislabeled dataset.

The Deep Learning optimization chapter distinguishes minimizing a training objective from learning a predictor that generalizes. Keep that distinction visible throughout training: the training loss directs updates; validation tells you whether those updates remain useful on unseen examples.

One update you can calculate

Use the model $\hat y=wx$ without an intercept and the invented observations $(1,2),(2,4)$. Define half mean squared error:

$$L(w)=\frac{1}{2n}\sum_i(wx_i-y_i)^2, \qquad g(w)=\frac{1}{n}\sum_i(wx_i-y_i)x_i.$$

At $w=0$, the loss is $5$ and the gradient is $-5$. With learning rate $\eta=0.1$, the update $w\leftarrow w-\eta g$ gives $w=0.5$; the loss becomes $2.8125$. The negative gradient says increasing the weight reduces local error. It does not mean every arbitrarily large increase is helpful.

For these particular data, $L(w)=1.25(w-2)^2$. The distance from the optimum follows

$$w_{t+1}-2=(1-2.5\eta)(w_t-2).$$

Convergence therefore requires $0<\eta<0.8$. At $\eta=0.8$ the distance keeps its magnitude and changes sign; larger rates diverge. This exact threshold belongs to this quadratic example, not to neural networks in general.

x, y = [1.0, 2.0], [2.0, 4.0]
w, learning_rate = 0.0, 0.1
for step in range(50):
    gradient = sum((w * a - b) * a for a, b in zip(x, y)) / len(x)
    w -= learning_rate * gradient
loss = sum((w * a - b) ** 2 for a, b in zip(x, y)) / (2 * len(x))
assert abs(w - 2.0) < 2e-6
print(round(w, 6), round(loss, 10))

This standard-library example has no held-out set because it demonstrates arithmetic, not predictive performance. For a model with an intercept and a visual experiment, see linear regression.

Why updates use batches

A full gradient reads every training example before one step. A minibatch estimates it from a subset. With suitable random sampling and an averaged per-example loss, this estimate targets the full-data gradient; successive updates still see different parameters and noisy directions. Dive into Deep Learning develops this stochastic-gradient viewpoint.

A step is an optimizer update; an epoch is one pass over the training set. With 1,000 examples and batches of 100, one epoch normally has 10 steps. Accumulating four minibatches before updating changes that count and the effective batch size. Average gradients consistently: summing four already averaged losses without dividing by four multiplies the gradient scale. Equivalence to a single large batch also depends on batch-sensitive layers and randomness.

Larger batches may improve device utilization but require more activation memory and make fewer updates per example pass. Comparing only epoch counts can therefore hide substantial differences. Record examples or tokens processed, update count, and wall-clock time alongside quality.

What an optimizer remembers

SGD uses the current gradient. Momentum also carries a moving direction from previous steps, which can reduce alternating movement. Adam maintains moving first and second gradient moments and rescales parameter updates. These states are additional arrays; a model that fits for inference may exceed memory once gradients, optimizer states, and saved activations are included.

Open full-size image

The blue curves join points with equal loss; the orange dots are successive parameter values. This separate two-dimensional example uses f(x₁,x₂) = 0.1x₁² + 2x₂², learning rate 0.6, and momentum coefficient 0.5. The vertical oscillations shrink as the trajectory approaches (0, 0). Ordinary gradient descent diverges at this learning rate on the same objective. These are deterministic gradients of a quadratic, not minibatch noise or the scalar w example above.

Zhang et al., Dive into Deep Learning, §12.6.1.3. · CC BY-SA 4.0

PyTorch’s AdamW definition separates weight decay from the adaptive gradient update. Consequently, adding an L2 penalty to the loss and setting AdamW’s weight decay are not generally interchangeable. State which convention you use rather than transferring a regularization coefficient between them unchanged.

An optimizer cannot choose a universally correct learning rate. Feature scale, initialization, batch composition, and the loss normalization affect its useful range. A warmup gradually raises the rate at the start; a decay schedule reduces it later. Each is a control over update size, not evidence that a run has learned the intended task.

Reading failures rather than just curves

If both training and validation losses remain high, first try fitting a tiny clean subset. Failure even there suggests a broken target, disconnected gradient, inappropriate objective, or unsuitable step size. If the tiny subset fits but the full data does not, capacity and optimization budget become plausible explanations.

If training improves while validation worsens, more updates are fitting distinctions that do not transfer. Check duplicate groups and split boundaries before changing regularization. Select stopping time and hyperparameters on validation data; opening the test set after every run turns it into another validation set. The data-splitting note explains that boundary.

Sudden non-finite losses call for checking inputs, divisions, exponentials, gradient norms, and precision. Gradient clipping can cap an update’s gradient norm but cannot repair invalid labels or missing values. A flat loss can also hide a bookkeeping bug: gradients may have been cleared too late or accumulated across steps unintentionally.

Resuming the same training process

Saving weights is sufficient to reuse a predictor. Continuing the same optimization trajectory also needs optimizer and scheduler states, update count, and relevant random/data-loader state. Resuming with fresh momentum can produce different updates even when the weights are identical. Record preprocessing and label definitions with the checkpoint so that the recovered model receives the same representation.

Once this loop is understood, pretraining and post-training change the source of targets, while LoRA and distillation change which parameters or which model learn from them.

My reading · Opened notes are remembered in this browser.