How Models Learn

1. The one-sentence intuition

"Training" is just an optimization loop: you measure how wrong the model is with a single number (the loss), figure out which direction to nudge each of the model's numbers (its parameters) to make that number smaller, take a small step in that direction, and repeat millions of times until the loss stops dropping. The direction is the gradient; stepping against it is gradient descent.

If you've ever written a feedback controller, a hill-climbing search, or a Newton-Raphson root-finder, you already have the mental model: define an objective, compute its slope, step toward the optimum, loop until convergence. A neural network is the same idea with a few billion knobs instead of one, and a clever trick (backpropagation) for computing the slope of all of them in a single pass.

2. Why a software engineer needs this

You will rarely write a training loop from scratch on the job — but every concept in this lesson is load-bearing vocabulary for the rest of AI engineering, and interviews assume you own it cold.

• Fine-tuning (/finetuning) is this loop, run on your data with a pretrained model as the starting point. LoRA, QLoRA, DPO, and GRPO are all variations on "what loss do we descend, and which parameters do we update." You cannot reason about them without the loop.
• Loss and gradients show up constantly in discourse: "the model collapsed," "loss spiked," "we lowered the learning rate," "gradients exploded," "we warmed up then decayed the LR." These are not metaphors — they're literal events in the loop you're about to understand.
• Evals (/evals) exist precisely because the training loss is not the thing you care about. Knowing the difference between the loss you optimize and the metric you ship is half of practical ML.
• Transformers (/transformers) are "just" a particular architecture plugged into this exact loop. Once the loop is second nature, the transformer is the only new thing left to learn.

Interviewers use this topic as a filter. If you can explain why cross-entropy and not MSE for classification, what the learning rate trades off, and what backprop actually computes, you sound like someone who has trained a model. If you can't, no amount of API experience hides it.

3. Build it up from scratch

3.1 The loss function: turning "wrong" into one number

A model is a function with tunable parameters: it takes an input x and produces a prediction ŷ ("y-hat"). To improve it, we first need to score it. A loss function maps a prediction and the true answer y to a single non-negative number: 0 means perfect, larger means worse. Training = making the average loss over your data as small as possible.

Two losses cover almost everything you'll meet.

Regression → Mean Squared Error (MSE). When the target is a continuous number (price, temperature), penalize the squared gap between prediction and truth:

Here n is the number of examples, i indexes them, ŷ_i is the model's prediction for example i, and y_i is the true value. Squaring makes every error positive and punishes big misses far more than small ones (being off by 4 costs 16, off by 2 costs 4). It measures distance.

Classification → Cross-Entropy. When the target is a category (cat/dog, or one of 50,000 possible next tokens), the model outputs a probability for each class, and we want to reward it for putting high probability on the correct class. Cross-entropy for a single example is:

where p_correct is the probability the model assigned to the true class. If the model is confident and right (p = 0.99), the loss is -log(0.99) ≈ 0.01 — tiny. If it's confident and wrong (p = 0.01 on the true class), the loss is -log(0.01) ≈ 4.6 — huge. It measures surprise: how shocked the model should be by the right answer.

Two terms you'll hear constantly live right here. Logits are the raw, unbounded scores the network outputs per class (e.g. [2.0, 1.0, 0.1]). To turn logits into probabilities that sum to 1, you apply softmax: exponentiate each, then normalize.

z_k is the logit for class k; the denominator sums over all classes j so the outputs form a valid probability distribution. Worked micro-example: logits [2.0, 1.0, 0.1] exponentiate to [7.39, 2.72, 1.11], sum 11.2, giving probabilities [0.66, 0.24, 0.10]. If the true class is the first, cross-entropy is -log(0.66) ≈ 0.42. If the true class were the last, it'd be -log(0.10) ≈ 2.3 — the same prediction is "good" or "bad" depending only on which class was correct.

Why not MSE for classification? MSE on probabilities barely punishes confident wrong answers and produces tiny, flat gradients when the model is very wrong (the place you most need a strong push). Cross-entropy's gradient stays large exactly when the model is badly wrong, so it learns faster and more stably. This pairing — softmax + cross-entropy — is the workhorse loss of every LLM, where "classification" means "predict the next token out of the vocabulary."

3.2 Gradient descent: which way is downhill

Now we have a number (loss) we want to minimize by adjusting parameters. Picture the loss as a landscape: the parameters are your coordinates, and the height is the loss. You're standing somewhere in the fog and want to reach the lowest valley. The strategy: feel the slope under your feet and step downhill.

The gradient, written ∇L ("grad L"), is the vector of partial derivatives of the loss with respect to every parameter. Each component answers: "if I increase this one parameter a hair, does the loss go up or down, and how steeply?" The gradient points in the direction of steepest increase, so to go downhill we step in the negative gradient direction. The update rule for a parameter θ ("theta") is:

θ is a parameter, ∂L/∂θ is the gradient component for it (the slope), and η ("eta") is the learning rate — how big a step you take. The ← means "assign." That's the entire algorithm: subtract a scaled slope from every parameter, repeat.

Worked micro-example. Fit ŷ = w·x + b to three points (1,2), (2,4), (3,6) (the true rule is y = 2x). Start at w = 0, b = 0, so every prediction is 0 and the errors (ŷ − y) are (−2, −4, −6). The MSE gradients are:

Plugging in: ∂L/∂w = (2/3)[(−2)(1)+(−4)(2)+(−6)(3)] = (2/3)(−28) ≈ −18.7 and ∂L/∂b = (2/3)(−12) = −8. Both gradients are negative — meaning increasing w and b would lower the loss. With learning rate η = 0.1:

w ← 0 − 0.1·(−18.7) = 1.87 and b ← 0 − 0.1·(−8) = 0.8.

One step moved w from 0 to 1.87, already close to the true value of 2. Repeat a few dozen times and it converges. That's gradient descent.

3.3 The learning rate: the one knob you must feel

The learning rate η controls step size, and it's the single most consequential hyperparameter. Too small and training crawls — thousands of wasted steps inching toward the valley. Too large and you overshoot the bottom on every step, bouncing up the far wall and diverging (loss climbs to infinity or NaN). There's a Goldilocks band, and in practice people don't keep it fixed: they warm up (start tiny so early steps don't blow up) then decay it over training (smaller steps near the minimum for fine settling).

Play with it directly — drag the learning rate and watch the ball either glide into the valley, crawl, or fly off the surface:

3.4 SGD, minibatches, and the vocabulary of a training run

The gradient formulas above sum over all n examples. With millions of examples, computing the full-dataset gradient for every single step is absurdly slow. The fix is Stochastic Gradient Descent (SGD): estimate the gradient from a small random minibatch of examples (say 32 or 256) instead of the whole dataset. Each estimate is noisy but unbiased, and you get to take many cheap steps instead of one expensive one.

The noise is a feature, not just a tax. The jitter from random batches helps the optimizer skip past shallow bad spots and tends to find flatter minima that generalize better — a mild, free regularization effect. (Regularization = anything that fights overfitting, i.e. memorizing the training data instead of learning the pattern.)

This gives you the three timing words that trip up beginners:

So "we trained for 3 epochs with batch size 256 on 1M examples" means 3 × (1,000,000 / 256) ≈ 11,700 parameter updates. LLM pretraining is usually described in steps or tokens, not epochs, because it sees each example roughly once.

3.5 Backpropagation: getting every gradient cheaply

In the toy example we had two parameters and wrote their gradients by hand. A real network has millions to billions, arranged in layers where each layer's output feeds the next. How do you get the gradient for a weight buried five layers deep? The honest answer is the chain rule from calculus — but the insight of backpropagation is the bookkeeping that makes it cheap.

A network's prediction is a composition of functions: loss(layer_N(...layer_2(layer_1(x)))). The chain rule says the derivative of a composition is the product of the derivatives of each step. Backprop runs in two passes:

1. Forward pass: feed x through the layers to compute the prediction and the loss, caching each layer's intermediate output.
2. Backward pass: start from the loss and walk backward, multiplying local derivatives layer by layer, propagating the "error signal" from the output back to every parameter.

The key efficiency point — the thing interviewers want — is that naively, computing each parameter's gradient separately would re-derive shared sub-expressions over and over (exponential blowup). Backprop computes the gradient of all parameters in one backward pass, at roughly the same cost as one forward pass, by reusing each intermediate result exactly once (this is dynamic programming on the computation graph). That O(network size) cost — rather than O(network size × parameters) — is why deep learning is computationally feasible at all. Modern frameworks call this autograd: they record the graph of operations during the forward pass and replay it backward automatically, so you never write a derivative by hand.

3.6 Optimizers: SGD → momentum → Adam/AdamW

Plain SGD takes a step proportional to the raw gradient. It works, but it's twitchy: it zig-zags across narrow valleys and stalls on flat plateaus. Optimizers are smarter recipes for turning gradients into updates.

You don't need the update equations memorized for IC3. You need the story: momentum smooths the path using gradient history; Adam additionally gives each parameter its own adaptive step size; AdamW is the LLM default. When someone says "we used AdamW with a cosine schedule and linear warmup," you now know that means: the adaptive-momentum optimizer, learning rate ramped up then smoothly decayed.

3.7 The canonical loop

Every framework, every fine-tune, every LLM pretrain is this five-line rhythm repeated until the loss flattens:

forward   →   loss   →   backward   →   step   →   repeat
(predict)   (score)    (get grads)   (update)

That's it. Everything else — architecture, data, schedules — is detail bolted onto this skeleton.

4. See it in code

First, the whole loop in pure NumPy so nothing is hidden — linear regression by gradient descent, exactly the math from §3.2:

import numpy as np
 
# tiny dataset; the true rule is y = 2x (the model has to discover it)
X = np.array([1.0, 2.0, 3.0, 4.0])
y = np.array([2.0, 4.0, 6.0, 8.0])
 
w, b = 0.0, 0.0      # parameters we will learn (start at zero)
lr = 0.01            # learning rate η
n = len(X)
 
for step in range(1000):
    y_hat = w * X + b                     # 1. FORWARD: predict
    error = y_hat - y                     #    residual (ŷ − y)
    loss  = np.mean(error ** 2)           # 2. LOSS: mean squared error
    grad_w = (2 / n) * np.sum(error * X)  # 3. BACKWARD: ∂L/∂w
    grad_b = (2 / n) * np.sum(error)      #            ∂L/∂b
    w -= lr * grad_w                      # 4. STEP: descend the gradient
    b -= lr * grad_b
    if step % 200 == 0:
        print(f"step {step:4d}  loss {loss:7.4f}  w {w:.3f}  b {b:.3f}")

Line by line: the forward pass computes predictions; error and loss score them; grad_w/grad_b are the hand-derived gradients (this is what backprop automates for big networks); the two -= lines are the step. Run it and the loss falls toward 0 while w → 2 and b → 0. Bump lr to 1.0 and watch it diverge to NaN — that's an oversized learning rate, live.

Now the same loop in PyTorch, where autograd computes the gradients for you:

import torch
import torch.nn as nn
 
X = torch.tensor([[1.], [2.], [3.], [4.]])
y = torch.tensor([[2.], [4.], [6.], [8.]])
 
model   = nn.Linear(1, 1)                              # ŷ = wx + b
opt     = torch.optim.SGD(model.parameters(), lr=0.01) # optimizer holds the params
loss_fn = nn.MSELoss()
 
for step in range(1000):
    opt.zero_grad()              # clear last step's gradients (they accumulate!)
    loss = loss_fn(model(X), y)  # FORWARD + LOSS in one line
    loss.backward()              # BACKWARD: autograd fills every param's .grad
    opt.step()                   # STEP: optimizer updates params using .grad

The structure is identical — zero_grad → forward → loss → backward → step — but loss.backward() walks the computation graph and computes ∂loss/∂θ for every parameter automatically, and opt.step() applies the update rule. Swap SGD for torch.optim.AdamW and you've upgraded optimizers in one token. One subtlety worth internalizing: PyTorch accumulates gradients across backward() calls, so you must zero_grad() each step or your gradients pile up and training breaks — a classic first-day bug. The full runnable version, with both implementations and a divergence demo, lives in examples/ml/train_loop.py.

5. Mental models & SWE analogies

• Gradient ≈ the compiler error that tells you which direction to fix. It doesn't hand you the answer; it tells you, for every knob, which way reduces the problem and by how much.
• Training loop ≈ gradient-based hill climbing / a feedback controller. Measure error, compute correction, apply a fraction of it (the learning rate is your control gain), repeat. Too much gain and the system oscillates and blows up.
• Backprop ≈ memoized recursion / dynamic programming. The chain rule is the recurrence; backprop is the memoization that computes all derivatives in one pass instead of recomputing shared sub-expressions exponentially.
• Learning rate ≈ a retry/backoff or simulated-annealing temperature. Big early, small late: explore boldly, then settle precisely.
• Loss vs. eval metric ≈ the proxy you can differentiate vs. the KPI you actually ship. You optimize cross-entropy because it's smooth and differentiable; you report accuracy/F1/win-rate because that's what users feel. Conflating them is a real bug (/evals).

6. Common confusions

• "The loss is the metric I care about." No — loss is a smooth, differentiable proxy chosen so gradients exist. Accuracy, BLEU, and win-rate are what you actually want; they're often non-differentiable, so you can't descend them directly.
• "Logits are probabilities." Logits are raw unbounded scores; softmax turns them into probabilities. A logit of 8 isn't "80%."
• "Backprop is a different algorithm from the chain rule." It is the chain rule — applied with smart caching so the whole gradient costs one backward pass. The cleverness is the bookkeeping, not new calculus.
• "Bigger learning rate = faster learning." Up to a point. Past it, you overshoot every step and the loss diverges to NaN. Faster only until it's catastrophic.
• "Epoch = step." An epoch is one full pass over the data; a step is one update on one minibatch. One epoch is usually thousands of steps.
• "SGD uses one example at a time." In modern practice "SGD" means minibatch SGD — a small batch per step, not a single example.
• "Adam is always better than SGD." Adam converges faster and is the default for transformers, but well-tuned SGD+momentum often generalizes as well or better in vision. "Default," not "universally optimal."
• "More epochs always help." Past a point the model starts memorizing the training set (overfitting) and held-out performance degrades. You watch validation loss, not training loss.

7. Check yourself

8. Go deeper

• Stanford CS229 — Machine Learning (cs229.stanford.edu): the canonical derivation of gradient descent, MSE, and logistic regression from first principles.
• Stanford CS231n — Optimization & Backprop notes (optimization-1, optimization-2): the clearest written treatment of SGD and backprop as a computation graph.
• 3Blue1Brown — Neural Networks series (3blue1brown.com): the visual intuition for gradient descent and backpropagation; watch before or after this lesson.
• Dive into Deep Learning — Linear Regression & Optimization (d2l.ai): runnable notebooks of exactly the loop in §4.
• Adam (arxiv 1412.6980) and AdamW (arxiv 1711.05101): the two optimizer papers you'll see cited everywhere — skim the intros now that you have the intuition.

Next: Neural Networks → — stack these trainable layers into a network and see what the forward pass and backprop actually run over, then on to Fine-tuning, which is this loop applied to a pretrained model.