Stochastic Gradient Descent (SGD)

Stochastic gradient descent (SGD) is an optimization algorithm that trains a model by repeatedly updating its weights with the gradient of the loss on a small random sample of the training data, instead of on the whole dataset. Each sample's gradient is a noisy but unbiased estimate of the full-dataset gradient, so SGD makes many cheap updates in the time full-batch gradient descent would make one expensive one. The method descends from Robbins and Monro's 1951 stochastic approximation method, and with momentum it remains a standard optimizer for convolutional neural networks (CNNs).

How SGD Works#

Every SGD step takes a mini-batch B, computes the gradient of the loss on it with backpropagation, and moves each weight a small step against that gradient: w ← w − η·∇L_B(w), where η is the learning rate. Strictly, SGD uses one example per step, but deep learning frameworks apply it to mini-batches of tens to thousands of examples, so "SGD" and "mini-batch gradient descent" are synonyms in practice. A larger batch size gives a less noisy gradient and better hardware utilization, at the cost of fewer updates per epoch.

Training shuffles the data at the start of each epoch, walks through it in mini-batches, and takes one step per batch. Fitting a line to 1,000 noisy points shows the whole procedure:

import torch

torch.manual_seed(0)

# 1,000 noisy samples of y = 3x + 2
x = torch.randn(1000, 1)
y = 3 * x + 2 + 0.1 * torch.randn(1000, 1)

model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)

for epoch in range(5):
    # Reshuffle every epoch, then update the weights once per 32-sample mini-batch
    for idx in torch.randperm(1000).split(32):
        loss = torch.nn.functional.mse_loss(model(x[idx]), y[idx])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: w = {model.weight.item():.3f}, b = {model.bias.item():.3f}, loss = {loss.item():.4f}")

The weight reaches 3.081 and the bias 2.012 after two epochs, 64 updates in total. Over the next three epochs the weight keeps moving between 2.967 and 3.003 instead of settling, because every mini-batch gradient carries sampling noise, and the last-batch loss stays near 0.01, the variance of the noise added to y.

Learning Rate, Noise, and Batch Size#

With a constant learning rate, SGD fluctuates in a region around the minimum whose size grows with the learning rate and the gradient noise. Stochastic approximation theory guarantees convergence only when the learning rate decays toward zero, so practical training lowers it over time with step decay or cosine schedules, usually after a short warmup. The gradient noise of small batches also steers SGD toward flat minima that generalize better than the sharp minima large-batch training tends to find, an implicit form of regularization.

Batch size and learning rate interact. Raising the learning rate in proportion to the batch size, with a five-epoch warmup, let ResNet-50 train on ImageNet in one hour with mini-batches of 8,192 images and no loss of accuracy.

Momentum and Nesterov Momentum#

In a long, narrow valley of the loss surface, the gradient points mostly across the valley, so plain SGD zigzags between the walls while making slow progress along the floor. Momentum keeps a running velocity, v ← μ·v + g, and steps along it, w ← w − η·v, so consistent gradient directions build up speed and oscillating directions cancel out. A coefficient of μ = 0.9 is the usual default. Nesterov momentum evaluates the gradient after a look-ahead step along the velocity, which damps overshoot further. PyTorch's SGD implements both and applies weight_decay by adding λ·w to the gradient, which for SGD is equivalent to an L2 penalty on the weights.

SGD vs. Other Optimizers#

• Batch gradient descent: Computes the exact gradient over the full dataset for every update, so each epoch yields a single step and memory needs grow with the dataset.
• Adam: Gives each parameter its own step size from running averages of its gradients and squared gradients, and stores two state values per parameter against one for SGD with momentum. Adam makes faster early progress and is less sensitive to the learning rate, while well-tuned SGD with momentum often matches or beats its final accuracy on CNNs. AdamW, its variant with decoupled weight decay, is the standard optimizer for transformers.
• Muon: Orthogonalizes the momentum update of each weight matrix as a whole, and has been scaled up to large language model training.

Ultralytics YOLO training builds PyTorch SGD with Nesterov momentum when optimizer="SGD" is set, using lr0=0.01 and momentum=0.937 by default. The default optimizer="auto" ignores lr0 and picks MuSGD, a hybrid that blends SGD and Muon updates, for runs longer than 10,000 iterations, and AdamW for shorter ones.

Where SGD Is Used#

SGD with momentum trained most of the landmark ImageNet CNNs. The original ResNet models used momentum 0.9, weight decay 0.0001, and mini-batches of 256 images, with a learning rate that started at 0.1 and was divided by 10 whenever the error plateaued. Outside deep learning, scikit-learn's SGDClassifier and SGDRegressor fit linear models such as logistic regression and linear support vector machines with SGD, and their partial_fit method trains on datasets too large to fit in memory one chunk at a time.