JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Stochastic Gradient Descent: How ML Trains

Every method so far assumed you could compute the full gradient. But a modern model has a loss summed over millions of examples, and one such gradient is too expensive to take even once. The fix is almost embarrassingly simple — peek at a tiny random handful each step — and it is the engine that trains essentially every neural network you have ever used.

The wall: a gradient you cannot afford

Every method in this rung so far has rested on one quiet assumption: that you can compute the gradient of your objective at the current point. Steepest descent needs the full gradient to pick its descent direction; Newton's method needs the gradient and the Hessian on top. For the smooth little test functions of the earlier guides, that cost nothing. But step into machine learning and the bill explodes.

Here is why. Training a model means minimizing a loss that is an average over data: L(w) = (1/N) * sum over i of loss_i(w), where w are the model's parameters and each loss_i measures how badly the model does on training example i. The full gradient is therefore also an average, grad L(w) = (1/N) * sum over i of grad loss_i(w). To take ONE such gradient, you must run every one of the N examples through the model and back. With N = 10^7 examples (small by modern standards) and millions of parameters, a single full gradient is a colossal computation — and gradient descent wants thousands of them.

So the methods you have mastered hit a brutal scaling wall. Newton is doubly out: forming and factoring a Hessian over millions of parameters is O(n^3) in the parameter count and not even storable. Plain full-gradient descent is merely impossibly slow rather than impossible, but 'impossibly slow' still means a model nobody can train. We need a new idea, and it comes from noticing what that average really is.

The idea: estimate the gradient from a tiny sample

The full gradient is an average over N examples. An average is exactly the kind of thing you can ESTIMATE from a small random sample — this is the same statistical instinct behind a poll, or behind Monte Carlo integration from the integration rung. Pick a handful of examples at random, average just their gradients, and you get a noisy but unbiased guess of the true gradient. On average it points the same way; it just wobbles. Stochastic gradient descent (SGD) is nothing more than running gradient descent on these cheap, noisy estimates.

In its purest form you use a single random example per step. The update is w_{n+1} = w_n - eta * grad loss_i(w_n), where i is one freshly drawn index and eta is the step size (in ML, called the learning rate). Notice the cost: instead of touching all N examples to take one step, you touch ONE example and take one step. In the time the full method took a single careful stride, SGD has taken a million scrappy ones. Even though each step is in a slightly wrong direction, a million slightly-wrong steps beat one perfect step you could never afford to compute.

Mini-batches: the practical middle

One example per step is very noisy; all N per step is very slow. The workhorse lives in between: the mini-batch. Each step draws a small random batch of, say, B = 32 or 256 examples and averages their gradients. This is still SGD — the estimate is still noisy and unbiased — but the noise is smaller. The key statistical fact, the same 1/sqrt(N) law from Monte Carlo, is that averaging B independent samples cuts the standard deviation of the estimate by a factor of sqrt(B). So a batch of 256 has gradient noise about 16x smaller than a single example, while costing 256x the work.

That trade-off — 16x less noise for 256x more work — looks bad on paper, so why are mini-batches universal? Because of hardware, not mathematics. A GPU processes a batch of 256 examples almost as fast as a single one: the work is a big matrix multiply, and the chip is built to do those in parallel. As you saw in the foundations rung, a real kernel is often limited by memory traffic, not raw flops — and a fat batch reuses the loaded model weights across all B examples, turning a memory-bound trickle into a compute-bound flood. The batch size is thus tuned to the machine: large enough to saturate the hardware, small enough that each step is still cheap and noisy enough to explore.

given parameters w, learning rate eta, batch size B:
  shuffle the N training examples
  repeat for each epoch (one full pass over the data):
    for each mini-batch of B examples:
      g = (1/B) * sum of grad loss_i(w) over the batch   # noisy gradient
      w = w - eta * g                                     # one cheap step
  (optionally decay eta after each epoch)
Mini-batch SGD in pseudocode. One 'epoch' is a full sweep through the data, made of many cheap noisy steps rather than one expensive exact one.

Momentum: borrowing the cure for zig-zag

Recall the zig-zag problem from guide 2: on an ill-conditioned objective — a long, thin valley — plain steepest descent bounces across the narrow walls while creeping along the valley floor, and its speed is throttled by the condition number of the Hessian. SGD inherits this pathology AND adds gradient noise on top, so its raw path is even more jittery. The classic, cheap remedy is momentum: instead of stepping along the current noisy gradient alone, keep a running average of recent gradients and step along that.

The mechanics are a two-line update. Maintain a velocity v, blend the new gradient into it, and step along the velocity: v_{n+1} = mu * v_n + grad, then w_{n+1} = w_n - eta * v_{n+1}, with mu a number like 0.9. The physical picture is exact and worth holding onto: instead of a weightless particle that obeys the gradient instantaneously, you now have a heavy ball rolling downhill. In the narrow direction, successive gradients point back and forth and cancel in the average, so the oscillation is damped. In the valley direction, successive gradients all point the same way and ADD UP, so the ball accelerates. Momentum quietly fixes both halves of the zig-zag at once.

There is a second, subtler gift in the stochastic setting. Because momentum averages many recent gradients, it also averages out a chunk of the SGD noise — the running velocity is smoother than any single mini-batch gradient. So momentum does double duty: it tames the deterministic ill-conditioning AND it quiets the stochastic jitter, for the price of storing one extra vector the same size as w. That is an extraordinarily good deal, which is why almost no one runs bare SGD.

Adam: a per-parameter learning rate

There is one more layer in the workhorse most people actually run, the Adam optimizer. Recall what made Newton so powerful: by using the Hessian, it gave each direction its own scaling, so a long thin valley got stretched back into a round bowl. Adam reaches for a shadow of that benefit without ever forming a Hessian. It keeps, for EACH parameter separately, a running average of recent squared gradients — a cheap, diagonal estimate of how steep that coordinate has been — and divides each parameter's step by the square root of that quantity.

The effect is a per-coordinate adaptive step size, the step-size choice made automatic and individual. Parameters whose gradients have been large and erratic get small, cautious steps; parameters whose gradients have been small and steady get large, confident steps. Combine that with the momentum-style averaging of the gradient itself, and you have Adam: one running mean of gradients (for direction, the momentum part) and one running mean of squared gradients (for per-parameter scale). It is not Newton — it only ever uses a diagonal, never the true cross-coupling in the Hessian — but it captures the cheap part of Newton's medicine and asks for almost nothing in return.

Why first-order methods won

Step back and see the strange inversion. In the earlier guides, the smartest method was the one that used the most information: Newton, with its full curvature, beat plain gradient descent handily. The quasi-Newton methods like BFGS were a brilliant compromise, approximating that curvature from gradient history. Yet at the scale of modern ML, the hierarchy flips. The method that wins is the DUMBEST one — first-order, noisy, curvature-blind SGD — precisely because it is the only one cheap enough to run.

This is a deep lesson about computation at scale, not a quirk of ML. When N is enormous, the right currency is not 'accuracy per step' but 'progress per unit of compute'. A cheap, sloppy step you can take a billion times beats an exact, expensive step you can take a hundred times. The same logic drove the iterative-solvers rung — iterate cheaply rather than factor exactly — and it is exactly why second-order curvature, the crown jewel of small-scale optimization, gets left at the door of large-scale learning. The noise is not even purely a cost: it helps the iterate skip past shallow saddle points and bad minima that a deterministic method might get stuck in.

One thread remains. Everything here optimized FREELY — w could be anything. But real problems often come with constraints: weights that must stay non-negative, a budget that must not be exceeded, a probability that must sum to one. SGD on its own knows nothing of walls. The final guide of this rung takes up exactly that question — Lagrange multipliers and the KKT conditions — and shows how optimization learns to respect a boundary. The humble, scalable algorithm you met here is the engine; the next guide gives it a fence.