Stochastic Gradient Descent (SGD): The Optimizer That Started It All

Every deep learning optimizer in use today — Adam, RMSProp, AdaGrad, and dozens of lesser-known variants — is a descendant of one simple, elegant idea: stochastic gradient descent. If you understand SGD deeply, you understand the backbone of nearly all modern machine learning training. In this article, I want to break SGD down completely, from the intuition to the math to the practical code, so you walk away with a genuinely solid grasp of it.

The Core Problem SGD Solves

Training a neural network means finding parameters $\theta$ that minimize a loss function $J(\theta)$, which measures how wrong the model’s predictions are on the training data. In principle, you could compute the exact gradient of the loss with respect to $\theta$ by averaging over every single training example:

$$J(\theta) = \frac{1}{N}\sum_{i=1}^{N} L(f(x_i;\theta), y_i)$$

$$\nabla_\theta J(\theta) = \frac{1}{N}\sum_{i=1}^{N} \nabla_\theta L(f(x_i;\theta), y_i)$$

This is called batch gradient descent — using the entire dataset for every single parameter update. The problem is obvious once you consider real-world dataset sizes: if $N$ is a million or a billion, computing a single gradient means passing the entire dataset through the model. That’s expensive, slow, and often doesn’t even fit in memory.

The Stochastic Insight

SGD’s insight is refreshingly simple: instead of computing the gradient over the full dataset, estimate it using just a small, randomly sampled subset — often even a single example. The update rule becomes:

$$\theta_{t+1} = \theta_t – \eta \nabla_\theta L(f(x_i; \theta_t), y_i)$$

where $(x_i, y_i)$ is a single randomly chosen training example (or, in the mini-batch version, a small batch of them).

This gradient estimate is noisy — it’s not the true gradient of the full dataset, just an unbiased estimate of it. But that noise turns out to be a feature, not just a bug. It lets the optimizer take many more steps per unit of computation, and the randomness itself can help the optimizer escape shallow local minima and saddle points that a smoother, exact gradient might get trapped in.

Mini-Batch SGD: The Practical Standard

In practice, almost nobody uses pure single-example SGD or full-batch gradient descent. The standard is mini-batch SGD, which strikes a balance: compute the gradient over a small batch of examples (commonly 32, 64, 128, or 256):

$$\nabla_\theta J(\theta_t) \approx \frac{1}{B}\sum_{i \in \text{batch}} \nabla_\theta L(f(x_i;\theta_t), y_i)$$

$$\theta_{t+1} = \theta_t – \eta \cdot \frac{1}{B}\sum_{i \in \text{batch}} \nabla_\theta L(f(x_i;\theta_t), y_i)$$

This gives you a gradient estimate that’s less noisy than single-example SGD (because you’re averaging over $B$ samples) but far cheaper to compute than full-batch gradient descent. It also maps efficiently onto GPU parallelism, since a batch of examples can be processed simultaneously.

Batch vs. Mini-Batch vs. Stochastic: A Comparison

Why the Noise Helps

Loss landscapes for deep neural networks are notoriously non-convex, filled with saddle points, plateaus, and local minima. A smooth, exact gradient (from full-batch descent) will happily walk straight into a nearby saddle point and get stuck, since the gradient there is essentially zero in every direction that matters. The noisy gradient of mini-batch SGD, by contrast, rarely points in exactly zero direction across every batch — the randomness effectively “jitters” the optimizer out of these bad flat regions. This is one reason SGD variants, even simple ones, tend to generalize surprisingly well in deep learning compared to methods that converge to sharper (but often worse-generalizing) minima.

The Learning Rate: SGD’s Most Sensitive Hyperparameter

Everything about SGD’s behavior hinges on the learning rate $\eta$:

• Too large: The updates overshoot the minimum, causing the loss to oscillate or even diverge.
• Too small: Training becomes painfully slow, and the optimizer may get stuck in poor regions for a very long time.
• Just right: Convergence is steady, and given enough iterations, the loss decreases reliably.

This sensitivity is exactly why learning rate scheduling (discussed in a separate article) and adaptive optimizers like Adam were developed — to reduce the burden of hand-tuning a single fixed learning rate for the entire training run.

SGD with Momentum

Vanilla SGD is often improved with a momentum term (covered in detail in my momentum article), which accumulates a moving average of past gradients:

$$v_t = \gamma v_{t-1} + \eta \nabla_\theta J(\theta_t)$$

$$\theta_{t+1} = \theta_t – v_t$$

Momentum tends to smooth out the noisy trajectory that plain SGD produces and often speeds up convergence considerably, especially on ill-conditioned loss surfaces.

Implementing SGD From Scratch

import numpy as np

def sgd(grad_fn, theta_init, data, lr=0.01, batch_size=32, epochs=10):
    theta = theta_init.copy()
    n = len(data)

    for epoch in range(epochs):
        np.random.shuffle(data)
        for start in range(0, n, batch_size):
            batch = data[start:start + batch_size]
            grad = grad_fn(theta, batch)  # average gradient over the batch
            theta = theta - lr * grad

    return theta

import numpy as np

def sgd(grad_fn, theta_init, data, lr=0.01, batch_size=32, epochs=10):
    theta = theta_init.copy()
    n = len(data)

    for epoch in range(epochs):
        np.random.shuffle(data)
        for start in range(0, n, batch_size):
            batch = data[start:start + batch_size]
            grad = grad_fn(theta, batch)  # average gradient over the batch
            theta = theta - lr * grad

    return theta

Using PyTorch, mini-batch SGD is built directly into the training loop through DataLoader and the SGD optimizer class:

import torch
from torch.utils.data import DataLoader

model = torch.nn.Linear(20, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.MSELoss()

for epoch in range(10):
    for x_batch, y_batch in DataLoader(dataset, batch_size=64, shuffle=True):
        optimizer.zero_grad()
        preds = model(x_batch)
        loss = loss_fn(preds, y_batch)
        loss.backward()
        optimizer.step()

import torch
from torch.utils.data import DataLoader

model = torch.nn.Linear(20, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.MSELoss()

for epoch in range(10):
    for x_batch, y_batch in DataLoader(dataset, batch_size=64, shuffle=True):
        optimizer.zero_grad()
        preds = model(x_batch)
        loss = loss_fn(preds, y_batch)
        loss.backward()
        optimizer.step()

Visualizing the Training Loop

flowchart TD
    A[Shuffle training data] --> B[Split into mini-batches]
    B --> C[For each mini-batch: forward pass]
    C --> D[Compute loss]
    D --> E[Backward pass: compute gradients]
    E --> F[Update parameters: theta = theta - lr * grad]
    F --> G{More batches?}
    G -- Yes --> C
    G -- No --> H{More epochs?}
    H -- Yes --> A
    H -- No --> I[Training complete]

Advantages of SGD

1. Computational efficiency — updates happen far more frequently than full-batch descent, so training progresses faster in wall-clock time.
2. Memory efficiency — you only need to hold a small batch in memory at once, rather than the entire dataset.
3. Better generalization in many cases — the inherent noise can act as an implicit regularizer, often steering the optimizer toward flatter, more generalizable minima.
4. Simplicity — extremely easy to implement, understand, and debug.
5. Scalability — works well for datasets too large to fit in memory or on a single machine, and is trivially parallelizable across GPUs.

Disadvantages and Limitations

1. Sensitive to learning rate — poor choices can cause slow convergence or divergence.
2. Noisy convergence path — the loss curve during training can look jagged and erratic, especially with small batch sizes.
3. Same learning rate for all parameters — plain SGD applies the identical step size to every parameter, even though different parameters (e.g., in the first layer vs. the last layer) might need very different step sizes.
4. Can struggle with ill-conditioned loss surfaces — narrow ravines cause zig-zagging (this is exactly what momentum was designed to fix).
5. No adaptivity to gradient magnitude — a parameter whose gradient is chronically small will update very slowly compared to one with a large gradient, unless you add adaptive techniques.

Real-World Use Cases

• Large-scale image classification: SGD with momentum remains a standard choice in training ResNets and other CNNs, often outperforming Adam in final test accuracy when properly tuned with a learning rate schedule.
• Online learning: When data arrives continuously (e.g., ad click-through prediction), single-example or small-batch SGD lets the model update incrementally as new data streams in.
• Federated learning: Distributed devices compute local mini-batch SGD updates that get aggregated centrally, since it’s simple, low-bandwidth, and doesn’t require the full dataset in one place.

The Statistical View: SGD as an Unbiased Estimator

It helps to think about mini-batch SGD through a statistical lens. The full-dataset gradient is a fixed quantity, but each mini-batch gradient is a random variable — a sample drawn from the “population” of per-example gradients. Crucially, the expected value of a mini-batch gradient equals the true full-dataset gradient:

$$\mathbb{E}\left[\frac{1}{B}\sum_{i \in \text{batch}} \nabla_\theta L(f(x_i;\theta), y_i)\right] = \nabla_\theta J(\theta)$$

This is why SGD is described as using an unbiased estimator of the true gradient. The batch size $B$ controls the variance of this estimator: larger batches produce lower-variance (smoother) gradient estimates, while smaller batches produce higher-variance (noisier) ones. This single fact explains almost every practical trade-off you’ll encounter when tuning batch size — it’s fundamentally a bias-variance style trade-off between computational efficiency, gradient noise, and generalization behavior.

How Batch Size Affects Training Dynamics

A well-known empirical finding is that when you increase batch size, you often need to proportionally increase the learning rate to maintain similar training dynamics — this is sometimes called the “linear scaling rule,” and it’s frequently paired with a warmup period to avoid early instability at the higher learning rate.

Convergence Guarantees and Theoretical Behavior

For convex loss functions, SGD with a properly decaying learning rate (satisfying conditions such as $\sum \eta_t = \infty$ and $\sum \eta_t^2 < \infty$, the classic Robbins-Monro conditions) is guaranteed to converge to the global minimum in expectation. For non-convex loss functions — which describes essentially every deep neural network — such strong guarantees don’t hold. Instead, SGD is generally only guaranteed to converge to a stationary point (a point where the gradient is zero), which could be a local minimum, a saddle point, or in rare cases a local maximum. In practice, the noise inherent in mini-batch gradients tends to help SGD avoid getting permanently stuck at saddle points, which is one reason it performs far better empirically on deep networks than the theoretical worst-case guarantees might suggest.

Debugging Common SGD Training Issues

• Loss oscillates wildly or diverges to NaN: almost always a learning rate that’s too high relative to your batch size and model scale. Try reducing it by a factor of 10 as a first diagnostic step.
• Loss decreases extremely slowly: learning rate may be too low, or your batch size may be too small relative to your chosen learning rate, producing gradient estimates too noisy to make consistent progress.
• Loss plateaus early and never improves further: consider whether you need momentum, a learning rate schedule, or whether the model itself lacks sufficient capacity for the task.
• Validation loss diverges from training loss (overfitting): this is a model capacity/regularization issue, not something SGD’s learning rate alone can fix — consider dropout, weight decay, early stopping, or more data.

SGD vs. Other Optimizers

Best Practices

• Always shuffle your data before each epoch to avoid the model learning spurious patterns from data ordering.
• Use mini-batches (commonly 32–256) rather than pure single-sample SGD or full-batch descent.
• Pair SGD with momentum (0.9 is a strong default) unless you have a specific reason not to.
• Use a learning rate schedule (step decay, cosine annealing, or warmup) rather than a single fixed rate for the whole run.
• Normalize or standardize your input features — SGD converges much more reliably when features are on comparable scales.
• Monitor validation loss, not just training loss, to catch overfitting early.

Frequently Asked Questions

Is “SGD” in frameworks like PyTorch actually pure stochastic gradient descent? Not quite — despite the name, torch.optim.SGD and its equivalents in other frameworks almost always implement mini-batch gradient descent, since you supply batches of data rather than single examples through your data loader. The name “SGD” has stuck around historically to refer to the entire family of stochastic (sampling-based) gradient methods, of which true single-example SGD is just one extreme case.

Why does my model train fine with SGD but the loss curve looks noisier than with Adam? This is expected and, to a degree, normal — SGD’s gradient estimates aren’t scaled or smoothed the way Adam’s are, so its loss curve typically has more visible short-term fluctuation even when the overall trend is healthy. As long as the underlying trend is decreasing and validation performance is improving, this jaggedness on its own isn’t necessarily a problem, especially with momentum added to smooth it out further.

Should I always shuffle my data before each epoch? Yes, essentially always, unless you have a very specific reason not to (such as certain time-series or curriculum-learning setups where order matters intentionally). Without shuffling, SGD can learn spurious patterns tied to the order in which data was originally collected or stored, and consecutive batches with correlated examples produce gradient estimates that are far noisier and less representative than properly shuffled batches.

How do I know if my batch size is too small or too large? A batch size that’s too small typically shows up as unusually noisy loss curves and slow, unstable convergence even after tuning the learning rate carefully. A batch size that’s too large typically shows up as fast per-step progress in wall-clock terms but a plateau in final model quality (or a need for a much larger learning rate and longer warmup to match the quality achieved with smaller batches). In practice, most practitioners pick the largest batch size that comfortably fits their hardware’s memory, then tune the learning rate accordingly, rather than treating batch size as a primary independent knob.

A Quick Mental Model to Keep

SGD’s entire philosophy can be summarized in one sentence: trade a small amount of gradient accuracy for a large amount of computational efficiency, and let the resulting noise work partly in your favor rather than purely against you. Every major optimizer built on top of SGD — momentum-based methods, adaptive methods like Adam — exists to manage or exploit that trade-off more intelligently, but the underlying stochastic sampling idea never goes away; it’s the foundation everything else is built on.

Summary

Stochastic gradient descent is the foundational optimization algorithm behind virtually all of deep learning. Instead of computing an exact but expensive gradient over the whole dataset, it estimates the gradient from small, randomly sampled batches, trading a bit of precision for a massive gain in computational efficiency and, often, better generalization. Its simplicity is deceptive — nearly every advanced optimizer used today, from momentum-based methods to Adam, is fundamentally an extension or refinement of this same core stochastic idea. Mastering SGD, including its sensitivity to learning rate and batch size, gives you the intuition needed to understand and debug virtually any deep learning training pipeline.

References and Further Reading

• Robbins, H., & Monro, S. (1951). “A Stochastic Approximation Method.” Annals of Mathematical Statistics.
• Bottou, L. (2010). “Large-Scale Machine Learning with Stochastic Gradient Descent.”
• Goodfellow, I., Bengio, Y., & Courville, A. (2016). “Deep Learning,” Chapter 8: Optimization for Training Deep Models.
• PyTorch SGD documentation: https://pytorch.org/docs/stable/generated/torch.optim.SGD.html
• TensorFlow SGD documentation: https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD