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 without an intercept and the invented observations . Define half mean squared error:
At , the loss is and the gradient is . With learning rate , the update gives ; the loss becomes . The negative gradient says increasing the weight reduces local error. It does not mean every arbitrarily large increase is helpful.
For these particular data, . The distance from the optimum follows
Convergence therefore requires . At 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 by four. 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 imageThe 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.
Each optimizer starts at the dashed ring and takes 100 steps with the same learning rate η. The gray lines join points of equal loss; ✕ marks the minimum.
f(x, y) = 0.1x² + 2y²Click or tap the plot to start from another point.
| Optimizer | (x, y) | loss | last step length |
|---|---|---|---|
| SGD | (−0.001, −1.3 × 10⁻²²) | 1.4 × 10⁻⁷ | 1.0 × 10⁻⁴ |
| Momentum | (−8.0 × 10⁻¹¹, −2.6 × 10⁻¹⁶) | 6.5 × 10⁻²² | 2.3 × 10⁻¹¹ |
| Adam | (0.021, 0.009) | 2.1 × 10⁻⁴ | 0.002 |
step 100: SGD loss 1.4 × 10⁻⁷; Momentum loss 6.5 × 10⁻²²; Adam loss 2.1 × 10⁻⁴
Along y the curvature is λ = 4, so each SGD step multiplies the y error by 1 − ηλ, of size 0.6 here. Below η = 2/λ = 0.5 that size stays under 1 and SGD settles, zigzagging when the factor is negative. Momentum with β = 0.5 stays stable while ηλ < 2(1 + β), here η < 0.75. Momentum adds 0.5 of its previous step to each new one, so it gathers speed along the valley floor. Adam divides each coordinate by a running size of its gradient, so its steps stay close to η whatever the slope: it does not blow up, but with a small η it is slow.
Illustrative: exact gradients of a two-variable function with PyTorch's update rules (momentum 0.5, Adam β₁ = 0.9, β₂ = 0.999). Real training adds minibatch noise and has far more dimensions.
PyTorch’s AdamW definition separates weight decay from the adaptive gradient update. In current PyTorch, AdamW is simply Adam with decoupled_weight_decay=True; the two spellings are equivalent. 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.
How much memory fine-tuning a 7B model takes
Training holds four things in GPU memory: the weights, their gradients, the optimizer state, and the activations the backward pass needs. Work through a 7-billion-parameter model. Weights in bf16 take 2 bytes per parameter, 14 GB in all; the gradients take another 14 GB; Adam keeps a first and a second moment for every parameter, 8 bytes in fp32, 56 GB in all. Full fine-tuning needs 84 GB before any activations, and many mixed-precision setups also keep an fp32 master copy of the weights, another 28 GB.
LoRA freezes the original weights and trains only a small set of adapter parameters. Suppose the adapters have 40 million parameters, about 0.6%. With the adapters and their gradients in bf16 and Adam's two moments in fp32, gradients plus optimizer state come to under 0.5 GB. The frozen weights still have to sit in memory, so the 14 GB term stays. QLoRA then stores the frozen weights in 4 bits, about half a byte per parameter, cutting 14 GB to roughly 3.5 GB plus a few quantization constants; gradients flow back through this frozen 4-bit model and update only the adapters. The paper used it to fine-tune a 65-billion-parameter model on a single 48 GB GPU.
Neither LoRA nor QLoRA removes the activation term. The gradient has to travel from the loss back to the adapters in the earliest layers, and every layer on the way needs the intermediate results it computed in the forward pass. Those results grow with batch size and sequence length: at the same batch size, going from 512 to 4,096 tokens per sequence multiplies activation memory by at least 8. When memory runs out, shorten sequences or shrink the batch first, and use gradient accumulation to recover the effective batch size. If that is not enough, turn on activation checkpointing: it discards intermediate results in the forward pass and recomputes them in the backward pass, spending extra compute to save memory.
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 samples and whether the same user or document appears across splits 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.