Stochastic Gradient Descent

Every modern neural network — the language model answering your questions, the image recognizer on your phone — was trained by minimizing a loss function. Loss measures how wrong the model is; training means finding the model weights that make loss as small as possible.

The textbook answer is gradient descent: compute the slope of the loss at your current weights, step downhill, repeat. The slope is the gradient, a vector pointing in the direction of steepest increase, so stepping in the opposite direction decreases the loss.

The catch: computing the exact gradient means summing over the entire training dataset — millions of examples — every single step. For even a modest deep network that is ruinously expensive.

Stochastic Gradient Descent (SGD) makes one radical substitution: instead of the full dataset, use a random mini-batch of, say, 32 or 256 examples. The mini-batch gradient is a noisy estimate of the true gradient, but it is cheap, and with enough steps the noise averages out. The algorithm arrived — in statistical form — with Herbert Robbins and Sutton Monro in 1951, and its modern deep-learning incarnation was systematized by Léon Bottou through the 1990s and 2000s.

The surprising truth: the noise is not merely tolerated — it actively helps. Noisy updates escape shallow local minima that would trap the exact gradient, and they implicitly regularize the model by preventing it from memorizing every quirk of the training set.

The canvas below shows a 2D loss surface (warm = high loss, cool = low loss). Each optimizer starts from the same point and tries to reach the minimum. Press Run to start, then adjust the learning rate and batch size to see how trajectory and speed change.

<div class="controls">
  <label>{{lbl_lr}} <span id="lrVal">0.05</span>
    <input type="range" id="lr" min="0.005" max="0.15" step="0.005" value="0.05">
  </label>
  <label>{{lbl_bs}}
    <select id="bs">
      <option value="1">{{opt_pure_sgd}}</option>
      <option value="8" selected>{{opt_mini_batch}}</option>
      <option value="200">{{opt_full_batch}}</option>
    </select>
  </label>
  <div class="btn-row">
    <button id="runBtn">{{btn_run}}</button>
    <button id="resetBtn" class="ghost">{{btn_reset}}</button>
  </div>
</div>
<canvas id="cv" width="380" height="280"></canvas>
<div id="legend" class="legend">
  <span class="dot" style="background:#e63946"></span> {{legend_full_batch}}
  <span class="dot" style="background:#457b9d"></span> {{legend_mini_batch}}
  <span class="dot" style="background:#2a9d8f"></span> {{legend_pure_sgd}}
</div>
<div id="info" class="info">{{info_initial}}</div>

*{box-sizing:border-box;margin:0;padding:0}
body{font-family:system-ui,sans-serif;color:#222;background:#f8f9fa;padding:.6rem}
.controls{display:flex;flex-wrap:wrap;gap:.5rem .9rem;align-items:center;margin-bottom:.5rem}
label{display:flex;align-items:center;gap:.35rem;font-size:.85rem;font-weight:600}
input[type=range]{width:100px;cursor:pointer}
select{font-size:.85rem;padding:.2rem .4rem;border:1px solid #adb5bd;border-radius:6px;cursor:pointer}
.btn-row{display:flex;gap:.4rem}
button{font:600 13px system-ui;padding:.35rem .8rem;border:1px solid #1d3557;
       background:#1d3557;color:#fff;border-radius:7px;cursor:pointer}
button.ghost{background:#fff;color:#1d3557}
canvas{display:block;border-radius:10px;border:1px solid #dee2e6;width:100%;max-width:380px}
.legend{display:flex;gap:.8rem;flex-wrap:wrap;margin-top:.4rem;font-size:.8rem;align-items:center}
.dot{display:inline-block;width:10px;height:10px;border-radius:50%;margin-right:2px}
.info{font-size:.82rem;color:#495057;margin-top:.3rem;min-height:1.2em}

// Code not found

Notice that the full-batch path is smooth but each step costs a full sweep of the data. The mini-batch path is jagged but much cheaper per step and often reaches a good solution faster in wall-clock time. Single-sample SGD bounces most — but sometimes crosses barriers that the others cannot.

SGD is deceptively simple to state; understanding why it converges is more subtle.

Convex case. When the loss surface is a bowl (convex), Robbins and Monro proved in 1951 that SGD converges to the global minimum if the learning rate schedule satisfies two conditions: the steps must sum to infinity (so the optimizer can reach anywhere) and the sum of squares must be finite (so the noise eventually dies). A harmonic schedule

$$\eta_t = \eta_0 / t$$

satisfies both.

Non-convex case. Deep networks have wildly non-convex loss landscapes — riddled with saddle points, flat plateaus, and countless local minima. Here the theory only guarantees convergence to a stationary point (where the gradient is zero), not a global minimum. Yet in practice, the minima that SGD finds in deep networks generalize remarkably well. A 2020 result by Chatterjee and others suggests that neural network loss landscapes are more benign than worst-case theory implies.

Momentum (Polyak, 1964) accumulates a velocity term: instead of stepping along the raw gradient, the optimizer builds momentum in directions of consistent gradient. This damps oscillations across steep ravines and accelerates progress along shallow valleys. PyTorch's default SGD uses momentum=0.9.

Adam (Kingma & Ba, 2014) goes further: it maintains per-parameter adaptive learning rates, dividing each gradient by a running estimate of its magnitude. Rare parameters get larger effective steps; common parameters get smaller ones. Adam is the default in almost every modern deep learning pipeline.

The cost per step is $O(B \cdot d)$, where $B$ is the batch size and $d$ is the number of parameters. Full-batch gradient descent uses $B = N$ (all examples); pure SGD uses $B = 1$; mini-batch typically uses $B \in \{32, 256, 1024\}$. Modern hardware (GPUs) parallelizes over the batch, so going from B=32 to B=256 costs little extra time while dramatically reducing gradient noise. See also non-convex optimization for the broader landscape of difficult optimization problems.

SGD and its adaptive variants are the universal training algorithm for virtually every large-scale machine-learning system:

• Large language models: GPT-4, LLaMA, Gemini and every other transformer are trained with Adam or a close variant over trillions of tokens. The sheer scale makes full-batch gradients inconceivable; mini-batches of a few thousand tokens run in parallel across thousands of GPUs.
• Image recognition: AlexNet (2012) famously used SGD with momentum to win ImageNet, triggering the modern deep-learning era. Every subsequent vision model — ResNet, ViT, CLIP — follows the same recipe.
• Recommendation systems: Netflix, YouTube and Spotify optimize embedding tables with billions of parameters using variants of Adam. Mini-batches correspond to random user sessions.
• Reinforcement learning: policy-gradient methods like PPO (used to train ChatGPT's RLHF stage) apply SGD-style updates to reward signals, not supervised labels.
• Scientific computing: neural network training for physics simulations (PINNs) and protein-structure prediction (AlphaFold) all rely on Adam to navigate high-dimensional loss surfaces with complex geometry.

The mini-batch idea has spread beyond neural networks too: online learning algorithms, Bayesian stochastic variational inference, and large-scale linear models all borrow the core SGD trick of replacing the full expectation with a cheap sample.

Stochastic gradient descent made one substitution — swap the exact gradient for a noisy mini-batch estimate — and that substitution unlocked an era. Training on millions of examples became tractable; noise that looked like a liability turned into a feature; and variants like Adam fine-tuned the idea into the universal workhorse of modern AI.

The deeper lesson is that approximate is often better than exact when the problem is large enough. Full-batch gradient descent is provably correct at each step but practically useless at scale. SGD is provably wrong at each step but practically unbeatable. Understanding why the wrong algorithm wins is one of the most productive questions in machine learning theory — and one we are still answering. Explore non-convex optimization if you want to dive into the landscape SGD is navigating.