Stochastic Gradient Descent: Mini-Batches, LR Schedules

In the previous chapter on backpropagation, we learned how to compute the gradient of the loss with respect to every parameter in a neural network. We have the gradient. Now what do we do with it?

The straightforward answer is gradient descent: repeatedly move parameters in the direction opposite to the gradient. But this naive formulation hides a critical practical problem. Computing the exact gradient requires evaluating the loss over your entire training dataset. If you have 10 million training examples, you need to run all 10 million through the network before taking a single parameter update step. This is batch gradient descent, and on modern datasets it is painfully slow.

Stochastic gradient descent (SGD) is the solution. Instead of computing the exact gradient over the full dataset, SGD estimates the gradient using a single training example or a small random batch. The estimate is noisy, but it comes fast, and with careful tuning, that noise turns out to be a feature rather than a bug.

SGD is the engine that makes deep learning work in practice. Every major language model you interact with today, from BERT to GPT-4, was trained using SGD or one of its descendants (Adam, AdamW). Understanding SGD at a deep level means understanding why neural networks can be trained at all, how learning rate choices shape the quality of what you train, and why certain tricks like mini-batching and learning rate schedules are so universally adopted. We will explore all of this here. The next two chapters will build on SGD to cover momentum and the Adam optimizer, which add adaptive mechanisms on top of the core ideas introduced here.

To understand why SGD matters, we first need to see what goes wrong with full-batch gradient descent on large datasets.

In standard gradient descent, we compute the loss averaged over all $N$ training examples and then update the parameters:

where:

• $\boldsymbol{\theta}$: the full parameter vector of the network
• $\ell(\boldsymbol{\theta}; \mathbf{x}_i, y_i)$: the loss on a single training example $(\mathbf{x}_i, y_i)$
• $N$: the total number of training examples
• $\eta$: the learning rate, a positive scalar controlling step size
• $t$: the iteration index

The gradient $\nabla_{\boldsymbol{\theta}} \mathcal{L}$ is an exact average over the full training set. This is mathematically clean: the gradient points precisely downhill in the loss landscape. But it has three serious problems in practice.

Computational cost. Modern NLP datasets contain millions of examples. Each forward-backward pass is expensive. Waiting until you have seen all $N$ examples before updating means you update parameters just once per epoch. With a dataset of 100 million sentences, you might wait hours between updates.

Memory cost. Full-batch gradient descent requires holding the gradients from all examples in memory before computing the average. For large datasets, this can exceed the available GPU memory by orders of magnitude.

Escape from saddle points. Neural network loss landscapes are high-dimensional and filled with saddle points (more on this later in this chapter). The exact gradient at a saddle point is exactly zero, which means batch gradient descent stalls. A little noise can kick the optimizer free.

The key insight of SGD is that we don't need the exact gradient. We need a gradient estimate that points, on average, in the right direction. A single randomly chosen training example provides exactly that:

The left side is the gradient on example $i$, a cheap computation. It is a noisy estimate of the true gradient on the right, but its expectation equals the true gradient:

where the expectation is taken over a uniformly random choice of example index $i$. This is the mathematical foundation of SGD: stochastic gradients are unbiased estimators of the true gradient. Each individual step is noisy, but over many steps the noise averages out.

The SGD update rule is simple. At iteration $t$, we randomly select example index $i_t$ and update:

where:

• $\boldsymbol{\theta}_{t+1}$: the updated parameters after step $t$
• $\boldsymbol{\theta}_t$: the current parameters before step $t$
• $\eta$: the learning rate (a fixed positive scalar, or one that changes according to a schedule)
• $i_t$: a uniformly random index drawn at iteration $t$
• $\nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t; \mathbf{x}_{i_t}, y_{i_t})$: the gradient of the loss on the single selected example

In practice, we almost never sample a single example at a time. We use mini-batches.

Mini-batch SGD averages the gradient over a small random subset of the training data, called a batch of size $B$:

where:

• $\mathcal{B}_t$: a randomly sampled subset of $B$ training indices at iteration $t$
• $B$: the batch size, typically 32, 64, 128, 256, or 512
• The $\frac{1}{B}$ averaging gives a more stable gradient estimate than a single example

Mini-batch SGD sits between the two extremes. It is noisier than full-batch gradient descent (which uses all $N$ examples) but much cheaper per update. It is less noisy than pure single-example SGD, which tends to produce erratic updates that can destabilize training.

In the deep learning literature, the term "SGD" almost always refers to mini-batch SGD in practice. When you see "SGD with batch size 256", that is mini-batch SGD with $B = 256$.

The choice of batch size involves an engineering tradeoff:

Gradient quality. Larger batches produce more accurate gradient estimates. With $B = N$ (full batch), variance in the gradient estimate is zero. With $B = 1$, variance is maximized. Mini-batches balance variance with cost.

Hardware efficiency. Modern GPUs are designed to process matrices in parallel. A single example cannot saturate a GPU's compute units. Batches of 64-512 examples typically utilize hardware much more efficiently than single-example updates, so larger batches often give more gradient updates per second, not fewer.

Generalization. This is the surprising part. Networks trained with small batches often generalize better to held-out data than networks trained with large batches, even when the training loss is similar. The reasons are still debated, but the leading hypothesis is that small-batch noise acts as a regularizer, pushing the optimizer toward flatter minima that generalize better. We discuss this further in the section on SGD noise below.

Practical observation. Across most NLP tasks, batch sizes in the range of 32 to 512 perform well. Very large batches (above 4096) often require special learning rate warmup strategies (covered below) to maintain convergence quality.

The learning rate $\eta$ is the most important hyperparameter in SGD. Choose it badly and your model will diverge, oscillate, or converge too slowly to be useful. Understanding the mechanics of learning rate selection is one of the most transferable skills in deep learning.

Consider what different learning rates do to the optimization trajectory.

Too large. If $\eta$ is too large, each step overshoots the local minimum. The loss oscillates or increases instead of decreasing. In extreme cases, the gradient explodes: parameter values reach infinity or NaN within a few steps.

Too small. If $\eta$ is too small, the optimizer makes tiny steps and takes a very long time to converge. Training becomes inefficient. With a decaying learning rate schedule, an initial rate that is too small can prevent you from ever reaching a good region of parameter space.

Just right. A good learning rate produces steady, consistent loss reduction without oscillation. The loss curve slopes downward smoothly.

Practitioners have developed several heuristics for choosing a learning rate:

• Start around 0.01 to 0.1 for standard SGD. These values work well across many tasks.
• Use the learning rate range test. Train for a few hundred steps while increasing the learning rate exponentially from $10^{-7}$ to 1, and plot the loss versus learning rate. The optimal learning rate is usually just before the point where the loss starts to increase.
• If using larger batches, scale the learning rate. A commonly used rule is linear scaling: if you double the batch size, double the learning rate. The intuition is that with twice as many examples per step, each step produces a more accurate gradient, so you can afford to take larger steps.
• Try 3e-4 for Adam-based optimizers. Adam is generally less sensitive to learning rate than plain SGD, and 3e-4 (also written 0.0003) works surprisingly often across language model pre-training tasks.

For plain SGD on NLP tasks, learning rates between 0.01 and 0.1 are common. For pre-training transformers, much smaller rates are typical (1e-4 to 1e-3) because the gradient landscape is sharper.

The theoretical convergence guarantees of SGD depend on properties of the loss function and on how the learning rate is set.

For convex loss functions, classical optimization theory provides clean convergence results. The key insight is that we need the learning rates to sum to infinity (so the optimizer can travel as far as needed), but the sum of their squares must be finite (so step sizes shrink fast enough to prevent endless oscillation). Formally, if the loss is convex and the learning rates $\{\eta_t\}$ satisfy the Robbins-Monro conditions:

where:

• $\eta_t$: the learning rate at iteration $t$
• $\sum_{t=1}^{\infty} \eta_t = \infty$: the divergence condition, which ensures the optimizer is not trapped in any finite region of parameter space
• $\sum_{t=1}^{\infty} \eta_t^2 < \infty$: the square-summability condition, which ensures that accumulated gradient noise does not prevent convergence

Then SGD converges to the global minimum with probability one. The classic schedule satisfying both conditions is:

where:

• $c > 0$: a positive constant that sets the scale of the learning rate
• $t_0 \geq 0$: an optional offset that prevents the learning rate from being excessively large at $t = 0$

This polynomial decay satisfies the divergence condition (the harmonic series $\sum_t 1/t$ diverges) and the square-summability condition (the series $\sum_t 1/t^2$ converges).

In practice, neural networks have non-convex losses. The Robbins-Monro theory does not directly apply. But the spirit carries over: a decaying learning rate helps SGD settle into a minimum rather than bouncing around it.

For non-convex losses (the realistic case for neural networks), the convergence theory is weaker. SGD is guaranteed to converge to a stationary point, a point where the gradient is zero. But a stationary point might be a local minimum, a saddle point, or even a local maximum (though the last case is rare and unstable in practice).

The surprising insight from empirical and theoretical work is that for overparameterized neural networks, almost all local minima have similar loss values. The loss landscape has many valleys, but they are all roughly equally deep. This means that SGD almost always finds a solution with good test performance, even though it is not guaranteed to find the globally optimal parameters.

The noise in mini-batch gradient estimates can help optimization and generalization rather than merely disrupting updates.

Escaping shallow minima. The random fluctuations in SGD gradient estimates can push the optimizer out of sharp, narrow minima. Wide, flat minima are more robust to noise (the gradient stays small even with perturbations), so SGD naturally gravitates toward them over time. This "implicit regularization" toward flat minima is thought to be a major reason why SGD generalizes well.

Saddle point escape. In high-dimensional loss landscapes, saddle points are far more common than local minima. At a saddle point, the gradient is zero, but the loss can be decreased by moving in at least one direction. SGD noise provides random perturbations that break the symmetry at saddle points, allowing the optimizer to continue making progress. We revisit saddle points in more detail later.

Noise as regularization. The SGD noise is equivalent to a form of regularization on the loss function. Specifically, training with SGD of batch size $B$ and learning rate $\eta$ approximately minimizes a regularized version of the loss with regularization strength proportional to $\eta / B$. Smaller batches or larger learning rates introduce more noise and thus more implicit regularization.

Training with a constant learning rate has a fundamental tension. Early in training, you want a large learning rate to explore parameter space and make rapid progress. Late in training, you want a small learning rate to converge precisely to a minimum without oscillating around it. A learning rate schedule resolves this tension by changing $\eta$ over the course of training.

This section covers the most important schedules. All of them modify the base learning rate $\eta_0$ over training steps or epochs.

Step decay reduces the learning rate by a fixed factor every $k$ epochs. If $\gamma$ is the decay factor (typically 0.1 or 0.5) and $e$ is the current epoch:

where:

• $\eta_e$: the learning rate at epoch $e$
• $\eta_0$: the initial learning rate
• $\gamma$: the decay factor per drop, a value in $(0, 1)$ (typically 0.1 or 0.5)
• $k$: the number of epochs between drops
• $\lfloor e / k \rfloor$: the number of drops so far (the floor of $e/k$)

Step decay is simple and predictable. The loss usually decreases smoothly between drops and then jumps to a lower value immediately after each drop, as the optimizer can now make more precise adjustments in the current neighborhood.

The main drawback of step decay is that the sharp transitions can be undesirable. Hyperparameter tuning requires choosing both $\gamma$ and $k$, and the performance can be sensitive to these choices.

Instead of discrete drops, exponential decay reduces the learning rate continuously at every step:

where:

• $t$: the current training step
• $\lambda > 0$: the decay rate, controlling how quickly the learning rate shrinks
• $e^{-\lambda t}$: the continuous decay factor, which approaches zero as $t$ increases

Exponential decay produces a smooth learning rate curve. However, if $\lambda$ is too large, the learning rate decays too quickly and training stalls before convergence.

Cosine annealing, introduced by Loshchilov and Hutter (2017) in their "SGDR" paper, is one of the most popular learning rate schedules for modern NLP. It smoothly decreases the learning rate from $\eta_{\max}$ to $\eta_{\min}$ following a cosine curve:

where:

• $\eta_t$: the learning rate at step $t$
• $\eta_{\min}$: the minimum learning rate (often 0 or a very small value like 1e-6)
• $\eta_{\max}$: the maximum learning rate (the "base" learning rate)
• $T$: the total number of training steps
• $\cos(\pi t / T)$: the cosine term, which starts at 1 (giving $\eta_{\max}$) and ends at $-1$ (giving $\eta_{\min}$) over the interval $[0, T]$

At $t = 0$, the cosine term is 1, so $\eta_0 = \eta_{\max}$. At $t = T$, the cosine term is $-1$, so $\eta_T = \eta_{\min}$.

Cosine annealing has several attractive properties: it starts large, decays slowly at first, then accelerates, and finally decays slowly again near the end. This matches intuitions about optimization: you want large steps early on, more cautious steps in the middle, and very fine-tuning at the end.

The "SGDR" variant adds warm restarts: at the end of a cosine cycle, the learning rate jumps back to $\eta_{\max}$ and the cycle repeats, often with a longer period. This allows the optimizer to escape any local minimum it has settled into and explore other regions of parameter space.

For large language models and transformers, training often begins with a warmup phase where the learning rate is linearly increased from 0 (or a tiny value) to the target learning rate over a number of steps $T_{\text{warmup}}$:

where:

• $T_{\text{warmup}}$: the number of warmup steps (typically a few hundred to a few thousand)
• $\eta_{\max}$: the target learning rate after warmup completes
• $t / T_{\text{warmup}}$: the linear ramp factor, going from 0 to 1

After warmup, the learning rate typically follows a decay schedule (cosine annealing is most common).

Why is warmup necessary? Early in training, the parameters are far from their optimal values. Gradients can be very large in magnitude, and starting with a large learning rate can cause catastrophic parameter updates that destabilize the entire network. By starting small and ramping up, warmup gives the network time to initialize into a reasonable part of parameter space before taking large steps.

In the original transformer (Vaswani et al., 2017), the following combined warmup-decay schedule was used:

where:

• $d_{\text{model}}$: the model dimension (e.g., 512 in the original transformer), which scales the overall learning rate to be smaller for larger models
• $t^{-0.5}$: the decay term active after warmup (polynomial decay)
• $t \cdot T_{\text{warmup}}^{-1.5}$: the warmup term active during the warmup phase
• The $\min$ selects the warmup term early in training and the decay term after

This formula automatically transitions from warmup to decay without requiring separate configuration of the phase boundary.

For most modern transformer training, the practical recommendation is: linear warmup for 1-5% of total training steps, followed by cosine annealing to $\eta_{\min} \approx 0$.

The chapter on Momentum will cover momentum in depth, but Nesterov momentum is closely enough tied to SGD that we introduce it here. Standard SGD with momentum maintains a velocity vector $\mathbf{v}$ that accumulates past gradients:

where:

• $\mathbf{v}_t$: the velocity vector at step $t$ (initialized to zero)
• $\mu \in [0, 1)$: the momentum coefficient, controlling how much past gradients influence the current step (typically 0.9)
• $\eta$: the learning rate
• $\nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t)$: the gradient evaluated at the current position $\boldsymbol{\theta}_t$

Momentum accumulates gradient directions over time, which helps SGD move faster in directions with consistent gradients and slows updates in directions where the gradient oscillates.

Nesterov momentum makes one key improvement: instead of computing the gradient at the current position $\boldsymbol{\theta}_t$, it computes the gradient at a "lookahead" position $\boldsymbol{\theta}_t + \mu \mathbf{v}_t$:

where:

• $\boldsymbol{\theta}_t + \mu \mathbf{v}_t$: the "lookahead" position, approximately where we will be after applying the accumulated momentum
• $\nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t + \mu \mathbf{v}_t)$: the gradient evaluated at the lookahead position, not at $\boldsymbol{\theta}_t$

Intuitively, Nesterov momentum first takes a step in the direction of the accumulated velocity, then computes the gradient at the resulting position, and corrects the step based on the resulting position. This "correct then move" approach is more responsive to changes in the gradient landscape than standard momentum's "move then correct" approach.

In practice, Nesterov momentum tends to converge faster and with better stability than standard momentum on convex problems. For non-convex neural networks, the empirical differences are smaller but Nesterov is still generally preferred. PyTorch's SGD implementation supports Nesterov momentum via the nesterov=True parameter.

Understanding the geometry of neural network loss landscapes reveals why SGD with noise often works surprisingly well, even without convergence guarantees.

In a neural network with $D$ parameters, the loss surface is a function in $D$-dimensional space. At any stationary point (where the gradient is zero), the local geometry is described by the Hessian matrix $H$, the $D \times D$ matrix of second derivatives:

The eigenvalues of the Hessian classify the stationary point:

• Local minimum: all eigenvalues are positive. The loss increases in every direction from this point.
• Local maximum: all eigenvalues are negative. The loss decreases in every direction. These are highly unstable and the optimizer neededly never rests here.
• Saddle point: some eigenvalues are positive (uphill directions) and some are negative (downhill directions). The loss can be decreased by moving along any negative eigenvalue direction.

Here is the main insight: in $D$ dimensions, a saddle point requires only that at least one eigenvalue is negative, while a local minimum requires that all $D$ eigenvalues are positive. With $D$ in the billions (for large language models), local minima become exponentially rarer than saddle points. The loss landscape of a deep network is dominated by saddle points.

This property helps SGD. Unlike local minima, saddle points have escape routes: directions in which the loss decreases. If SGD can find and follow these directions, it makes progress even in the absence of a true local minimum.

At a saddle point, the exact gradient is zero, so standard gradient descent stops. Mini-batch SGD has a natural escape mechanism: the random noise in the gradient estimate provides a random perturbation to the update direction. With high probability, this perturbation has a component along the negative-curvature direction (the downhill direction from the saddle), and the optimizer escapes.

The escape time is related to the magnitude of the noise relative to the saddle's curvature. Larger batch sizes mean less noise and slower escape. Smaller batch sizes provide more noise and faster (though less controlled) escape. This is one reason why very large-batch training sometimes gets stuck in poor regions of parameter space that small-batch training avoids.

Among the minima that SGD does find, not all are equally good for generalization. A local minimum can be:

Sharp: the loss surface curves steeply in many directions away from the minimum. A sharp minimum is easy to find but sensitive to small perturbations, which means the model's performance can degrade significantly when the inputs shift slightly from the training distribution.

Flat (wide): the loss surface is nearly flat around the minimum. Many nearby parameter configurations have similar loss. Flat minima tend to generalize better because small distributional shifts in the test data do not dramatically change the loss.

SGD with small batches and appropriate learning rates naturally gravitates toward flat minima. The intuition is that noise in the gradient estimate acts as a random perturbation to the parameters. At a sharp minimum, these perturbations frequently push the optimizer over the steep walls and out of the minimum. At a flat minimum, the perturbations stay within the low-loss region. Over time, SGD spends more time in flat regions, which correspond to models that generalize better.

Visualizing the loss landscape directly is difficult: it lives in $D$-dimensional space, where $D$ can be billions. But we can create informative 2D projections.

The technique introduced by Li et al. (2018) in "Visualizing the Loss Landscape of Neural Nets" creates a 2D slice of the loss surface by choosing two random directions $\boldsymbol{\delta}$ and $\boldsymbol{\eta}$ in parameter space and evaluating the loss at:

where:

• $\boldsymbol{\theta}^*$: the final trained parameters (the center of the visualization)
• $\alpha, \beta$: scalars that vary over a grid (e.g., both from $-1$ to $1$ in steps of 0.1)
• $\boldsymbol{\delta}, \boldsymbol{\eta}$: random unit vectors in parameter space (often constructed by sampling each component from a Gaussian and then normalizing)

The result is a 2D "slice" through the loss landscape centered on the solution. Sharp solutions produce highly curved, steep contours. Flat solutions produce gently curving, wide contours.

This visualization technique has been used to show that residual connections (skip connections, as in ResNets) dramatically smooth the loss landscape, making optimization much easier. Similar studies have shown that batch normalization also smooths the landscape. These empirical findings help explain why these architectural choices work so well.

To build intuition for how SGD moves across a loss surface, we can visualize optimizer trajectories on a simple 2D quadratic loss function. While real neural network losses are far more complex, the basic dynamics of learning rate effects and batch noise are captured clearly in this simplified setting.

The trajectory plot on the left demonstrates a key challenge with SGD on ill-conditioned problems: on the elongated ellipsoidal loss surface, a large learning rate causes the optimizer to overshoot and oscillate across the steep dimension while making slow progress along the shallow one. A well-chosen learning rate finds the right balance. The right plot illustrates why flat minima are preferred: SGD noise (represented by the shaded band) easily pushes the optimizer out of the narrow sharp minimum, but stays contained within the flat minimum's wide basin.

Let us implement SGD from scratch and apply it to a practical NLP task: training a logistic regression classifier on sentiment analysis, then extending to a small multilayer perceptron.

We will train on the 20 Newsgroups dataset, classifying posts between two sports categories using TF-IDF features. This lets us see SGD in action on a real text classification task without requiring any external downloads.

The training set has several hundred thousand parameters to learn across the TF-IDF feature space. SGD processes mini-batches of this data to iteratively improve the model.

Before building a multi-layer network, let us implement logistic regression with pure SGD. This isolates the optimization algorithm from architecture complexity.

The model learns quickly and reaches high accuracy. Notice how the training and validation losses move together, indicating the model is not badly overfitting with this batch size and learning rate.

One of the most illuminating experiments is to compare how different batch sizes affect the training trajectory.

Under this fixed 30-epoch budget, batch size 1 makes enough updates to reduce loss quickly and reaches strong validation accuracy first. Batch size 32 improves more gradually and eventually approaches the same accuracy range. The larger batches make far fewer updates per epoch, so batch size 256 and full-batch training remain close to chance accuracy in this run. The smooth epoch-level curve for batch size 1 reflects averaging across all of its noisy per-example updates.

Next, let us compare different learning rate schedules on the same classification task.

The learning-rate schedules produce visibly different update profiles. In this run, the constant rate continues reducing validation loss and finishes lowest. Cosine and warmup-cosine converge to nearly the same intermediate loss, while step decay finishes slightly higher. This result is specific to the shared training budget and illustrates why a schedule must be evaluated on the task rather than assumed to outperform a constant rate.

To develop intuition for SGD's path, we can visualize a low-dimensional slice of the loss surface around the trained solution.

The broad color bands indicate low curvature along these sampled directions. Their steady slope also shows that the held-out validation subset would continue improving away from the training solution along this particular slice, so the center should not be interpreted as the validation-loss minimum.

Understanding the main hyperparameters of SGD helps you tune it effectively for NLP tasks:

• learning_rate: The most critical hyperparameter. Start around 0.01-0.1 for plain SGD, or 1e-3 to 1e-4 for transformers. Use a learning rate finder to narrow the range.
• batch_size: Controls the noise-efficiency tradeoff. Typical values are 32-512. Smaller batches add regularizing noise; larger batches are faster per epoch but may generalize slightly worse.
• momentum: Coefficient for momentum (used when SGD is configured with momentum). Typical value is 0.9. Nesterov momentum is generally preferred over standard momentum.
• weight_decay: L2 regularization applied to weights (separate from learning rate). Typical values are 1e-4 to 1e-2. Helps prevent overfitting, especially in language model fine-tuning.
• schedule: The learning rate schedule. Cosine annealing with linear warmup is the current default for transformer training.
• T_warmup: Number of warmup steps before the main schedule begins. Typically 1-5% of total training steps.

Training language models has revealed several practical rules that work well across tasks:

Pre-training transformers. Use AdamW (covered in its own chapter), not plain SGD. AdamW combines adaptive learning rates (per-parameter) with proper weight decay. For learning rates, use cosine decay from 1e-4 to 3e-4 with linear warmup for 1-2% of steps.

Fine-tuning transformers. Smaller learning rates than pre-training, typically 1e-5 to 1e-4. Cosine decay with warmup still applies. Larger batches (256-2048) are common when fine-tuning on labeled datasets.

Classical NLP models. Plain SGD with momentum 0.9 works well for simpler models like logistic regression and shallow networks. Learning rates of 0.01-0.1, batch sizes of 32-128. Step decay (halve every 10 epochs) or cosine annealing both work.

Monitoring training. Always track both training and validation loss. If validation loss stops improving while training loss continues to decrease, you are overfitting and may want to reduce the learning rate more aggressively, increase weight decay, or stop early.

Gradient norms. Monitor the gradient norm (the Euclidean norm of all gradients stacked into a vector) throughout training. A gradient norm that grows without bound signals divergence. A gradient norm that drops to near zero too quickly signals premature convergence. The next chapter on Gradient Clipping covers this in detail.

Despite its central role in deep learning, SGD has real limitations that explain why practitioners have developed more sophisticated optimizers on top of it.

Learning rate sensitivity. Plain SGD is highly sensitive to the learning rate. A rate that is slightly too large causes divergence; a rate that is slightly too small wastes training compute. In contrast to adaptive optimizers, there is no mechanism in plain SGD to automatically compensate for a poor initial choice. You must rely on manual tuning, learning rate range tests, and experience. This sensitivity is particularly acute early in training, when gradient magnitudes can vary wildly across layers. This is why adaptive optimizers like Adam and AdamW, which adjust the effective learning rate per parameter based on gradient history, have largely replaced plain SGD for language model training.

Isotropic updates. SGD applies the same learning rate to all parameters, regardless of how much each one is used or how frequently it receives signal. In natural language processing, this creates a serious problem: token embeddings for common words like "the" or "is" receive gradient updates on nearly every batch, while embeddings for rare tokens like technical jargon may be updated only a handful of times per epoch. When the same learning rate governs both, you face an impossible tradeoff. The rate that works for common tokens overshoots on rare tokens; the rate that works for rare tokens is too slow for common ones. Adaptive optimizers like AdamW address this by scaling updates by the estimated curvature (second moment) of each parameter, effectively giving rare parameters a larger effective learning rate and common parameters a smaller one.

Convergence speed on ill-conditioned problems. If the loss function has very different curvatures in different directions, forming a "ravine" in the loss landscape, SGD oscillates back and forth across the narrow (steep) dimension while making slow progress along the long (shallow) dimension. Mathematically, the condition number of the Hessian measures this ill-conditioning. High condition numbers lead to slow convergence for SGD. Momentum helps with this by accumulating velocity that smooths out oscillations (covered in the next chapter), and adaptive optimizers help even more by effectively pre-conditioning each parameter independently.

Stale gradients in distributed training. Modern language model training runs across hundreds or thousands of GPUs. In distributed training, gradient synchronization takes time: each worker computes its gradient on its local mini-batch, then gradients must be aggregated across all workers before updating the parameters. In asynchronous SGD, the gradient computed by a worker may be based on parameters that have already been updated multiple times by other workers. These "stale" gradients can point in the wrong direction relative to the current parameter values, slowing or destabilizing convergence. Modern distributed training frameworks balance these tradeoffs with techniques like gradient accumulation, local SGD, and carefully tuned communication patterns.

Despite these limitations, SGD remains foundational. Even when using Adam or AdamW, you are using a first-order optimizer that fundamentally estimates the gradient from mini-batches and uses it to update parameters. The core ideas here, batch size, learning rate, schedules, noise, and landscape geometry, all remain directly relevant. Adam and AdamW, which we cover in the next chapters, are best understood as SGD with additional machinery for adapting the learning rate and handling weight decay correctly.

Stochastic gradient descent addresses the core impracticality of batch gradient descent on large datasets by replacing the exact gradient with a cheap, unbiased stochastic estimate. The key ideas from this chapter are:

• Mini-batch SGD averages gradients over small random subsets of data, balancing noise against computational efficiency. Batch sizes of 32-512 work well for most NLP tasks.
• Learning rate is the most critical hyperparameter. Too large causes divergence; too small causes slow convergence. Learning rate range tests and standard heuristics (0.01-0.1 for plain SGD, 1e-4 to 3e-4 for transformers) provide starting points.
• Learning rate schedules decay the learning rate over training to transition from exploration to fine-tuning. Cosine annealing with linear warmup is the current default for transformer training.
• SGD noise acts as implicit regularization, preferring flat minima over sharp ones. This helps generalization and allows escape from saddle points.
• High-dimensional loss landscapes are dominated by saddle points rather than local minima. SGD noise provides the random perturbations needed to escape these saddle points.
• Nesterov momentum computes the gradient at a lookahead position rather than the current position, enabling more efficient updates. The next chapter covers momentum in depth.

The chapters that follow build directly on SGD. Chapter 7 (Momentum) shows how to accumulate velocity across steps to accelerate convergence. Chapter 8 (Adam) introduces per-parameter adaptive learning rates that largely solve the sensitivity problem. Chapter 9 (AdamW) fixes Adam's interaction with weight decay, making it the default optimizer for modern language model training.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about stochastic gradient descent.