AI & ML — INTERMEDIATE
Gradient Descent: The Update Rule, Learning Rate, and Common Loss Functions
The Introduction level covered the core idea: measure how wrong a prediction is, nudge the weights downhill, repeat. This level gets concrete about the actual update rule, why the learning rate is arguably the single most consequential number in the entire training process, and the two loss functions you’ll actually run into in practice.
The Quick Answer
Every gradient descent step follows one exact rule: w_new = w_old − learning_rate × gradient. The gradient is the partial derivative of the loss with respect to that weight — it tells you which way increases the loss, so subtracting a step in that direction moves the weight toward lower loss. The learning ratescales how big that step is, and getting it wrong in either direction has real, distinct failure modes: too high overshoots and can diverge, too low wastes enormous amounts of training time. The exact shape of “loss” itself also depends on the task — mean squared error (MSE) for continuous outputs, cross-entropy for classification and next-token prediction.
The Update Rule, Applied Repeatedly
Training a model is this update rule applied to every weight, on every training example or batch, over and over:
w_new = w_old − learning_rate × gradient
The gradient here is ∂L/∂w — the partial derivative of the loss L with respect to that specific weight w. It answers one question: if this weight increased slightly, would the loss go up or down, and by how much? A positive gradient means increasing the weight would increase loss, so the rule subtracts a positive amount — decreasing the weight. A negative gradient means the opposite. Watch this play out step by step on a real (if simplified) loss curve:
Learning Rate: Three Trajectories, Same Starting Point
The learning rate is often called the most important hyperparameter in the entire training process, and it’s easy to see why once you watch it change the actual path a weight takes. Same starting weight, same loss curve, same number of steps — only the learning rate changes:
Notice the “too high” case doesn’t just converge slower — it can genuinely diverge, landing on a point steeper than where it started, so the next step overshoots even further. This isn’t a hypothetical edge case; it’s one of the most common reasons a training run’s loss curve suddenly spikes or turns into noise instead of smoothly decreasing.
One Weight, Computed by Hand
Before scaling up to many weights, it’s worth confirming the update rule by computing it yourself on numbers small enough to check with a calculator. This uses the simplest model shape possible — prediction = weight × input — with squared-error loss:
From One Weight to Several
Real models have far more than one weight, but the mechanism doesn’t change — it just runs in parallel. Each weight gets its own partial derivative of the same loss value, and its own update, computed independently:
Scale that up from two weights to the billions in a modern language model, and the idea is identical — only the bookkeeping changes. Efficiently computing every one of those partial derivatives at once, layer by layer, is exactly what backpropagation does, which is the natural next topic in this series once you’re comfortable with the update rule itself.
MSE vs. Cross-Entropy: Picking the Right Loss for the Task
Every example so far used squared-error loss, which fits regression tasks — predicting a continuous number, like a price or a temperature. But a huge share of machine learning, including language models, is classification — picking one option out of a fixed set of possible classes. That needs a different loss function.
| Loss function | Fits this kind of output | Example task |
|---|---|---|
| Mean squared error (MSE) | A single continuous number | Predicting a house price |
| Cross-entropy | A probability distribution over classes | Predicting the next token out of tens of thousands of possible tokens |
Cross-entropy loss measures how far off the model’s predicted probability distribution is from the correct answer, and it specifically punishes confidently wrong predictions much harder than a model that was appropriately unsure. This is the loss function actually used to train language models, since “predict the next token” is a classification problem over the entire vocabulary at every single step of generation.
Whichever loss function fits the task, the training loop around it never changes: compute the loss, compute its gradient with respect to every weight, take a step, repeat.
When the Loss Landscape Has More Than One Dip
One more real complication: gradient descent only ever reacts to the local slope. If the loss landscape has more than one dip, where you start can determine which one you end up in — even if a better, deeper dip exists somewhere else entirely.
In practice, this is one reason training runs sometimes use random weight initialization, multiple training runs, or momentum-based variants of gradient descent — all ways of giving the process a better chance of not getting permanently stuck in a shallow dip early on.
Fun Fact
The Adam optimizer — one of the most widely used training algorithms in deep learning today — is itself just a more sophisticated variant of this same update rule. Instead of a single fixed learning rate for every weight, Adam keeps a running estimate of each weight’s own gradient history and adapts its effective step size individually, per weight, over the course of training.
Source: Kingma & Ba, “Adam: A Method for Stochastic Optimization,” arXiv:1412.6980
Test Yourself
What is the exact gradient descent update rule for a single weight?
Based on the learning-rate comparison activity, what actually happens with a learning rate that is set far too high?
When would you use cross-entropy loss instead of mean squared error (MSE)?
In the multi-weight example on this page, why does each weight get its own gradient value?
Brainstorm
No right or wrong answers here — just questions worth pausing on before you expand the model answer.
The update rule subtracts learning_rate × gradient, rather than adding it. Why subtract specifically?
If two weights in the same model have very different gradient magnitudes at the same training step, what does that suggest about how the loss depends on each of them right now?
Cross-entropy loss is built around probabilities the model assigns to each possible class. What problem would you run into if you used plain MSE for a task like next-token prediction instead?
Suppose you trained a model and the loss dropped extremely fast for the first few steps, then barely moved for thousands of steps after that. What are two different explanations for that shape, and how might they call for different responses?
If you doubled the learning rate partway through training instead of keeping it fixed, what would you expect to happen to the size of each subsequent weight update, and why might a real training setup deliberately change the learning rate over time?
Ready for adaptive optimizers like Adam and momentum in more depth, the vanishing/exploding gradient problem, and how learning rate schedules are actually chosen in production training runs? Continue to the Advanced level →