PyTorch Loss Functions and Optimizers: Adam vs SGD

Every PyTorch training loop combines a loss function with an optimizer. The loss function produces a scalar that measures how far the model's predictions are from the targets, and the optimizer uses the gradient of that scalar to update the model parameters. For most projects, the practical decision is between cross-entropy or MSE for the loss, and Adam or SGD for the optimizer. This article explains how python pytorch loss functions optimizers adam and sgd fit together in a training loop, and what actually differs between Adam and SGD at runtime.

Selecting the Right Loss Function

PyTorch provides loss functions as modules under torch.nn. The choice depends on the task:

• nn.CrossEntropyLoss for multi-class classification. It expects raw logits, not probabilities, because it applies log-softmax internally. Passing already-softmaxed outputs double-applies the softmax and produces incorrect gradients.
• nn.MSELoss for regression. It computes the mean squared error between predictions and targets.
• nn.BCEWithLogitsLoss for binary classification. It combines a sigmoid with binary cross-entropy, which is numerically more stable than applying sigmoid and then BCE separately.

The loss function and the optimizer are independent. You can pair any loss with any optimizer. The optimizer only sees the gradient of the loss with respect to the parameters; it does not care which loss produced it.

Configuring the Optimizer

PyTorch optimizers live in torch.optim. The two most common are optim.SGD and optim.Adam. Both take model.parameters() as the first argument and a learning rate.

import torch.optim as optim

sgd_optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
adam_optimizer = optim.Adam(model.parameters(), lr=0.001)

SGD without momentum is plain gradient descent: each parameter is moved in the negative gradient direction by the learning rate. Adding momentum=0.9 makes the update accumulate a velocity term, which smooths the path and helps escape small local variations.

Adam is an adaptive method. It maintains a moving average of the gradient (first moment) and a moving average of the squared gradient (second moment) for every parameter. The update step scales the learning rate per parameter based on the ratio of these two moments. This is why Adam typically works with a smaller default learning rate and requires less tuning.

How SGD Updates Parameters

Plain SGD applies the update:

param = param - lr * grad

With momentum, the update becomes:

velocity = momentum * velocity - lr * grad
param = param + velocity

The momentum term carries information from previous steps. A common value is 0.9, which means the velocity retains 90% of its previous value at each step. Higher momentum values make the optimizer more aggressive and can overshoot, while lower values make it behave closer to plain gradient descent.

SGD stores one extra tensor per parameter when momentum is enabled (the velocity). Without momentum, it stores nothing beyond the parameters themselves. This makes SGD the most memory-efficient optimizer in PyTorch.

How Adam Updates Parameters

Adam tracks two state tensors per parameter: the first moment m and the second moment v. At each step, it computes:

m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * grad^2

Then it applies bias correction to both moments, because they start at zero and would otherwise be underestimated early in training. The final update is:

param = param - lr * m_hat / (sqrt(v_hat) + epsilon)

The default values are beta1=0.9, beta2=0.999, and epsilon=1e-8. The epsilon term prevents division by zero when the second moment is very small.

Because Adam tracks two moving averages per parameter, it uses roughly twice the memory of SGD with momentum. For a model with millions of parameters, this difference is measurable and can matter on memory-constrained hardware.

Adam vs SGD: What Actually Differs

The practical differences between Adam and SGD come down to learning rate behavior, memory, and training dynamics.

SGD with momentum is often preferred when you have a well-tuned learning rate schedule and a model that benefits from careful annealing. Adam is a strong default for getting a model to converge quickly without much hyperparameter work.

A common pattern is to start with Adam for initial experiments, then switch to SGD with momentum for final training if the model is small enough and the schedule can be tuned.

A Minimal Training Loop

Here is a complete training step that combines a loss function with an optimizer:

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

The order matters. zero_grad() clears gradients from the previous step; without it, gradients accumulate across steps. backward() computes the gradient of the loss with respect to every parameter. step() applies the optimizer's update rule.

Common Failure Modes

Several issues appear frequently when combining loss functions and optimizers in PyTorch.

Using CrossEntropyLoss with softmaxed outputs is a common mistake. The loss applies log-softmax internally, so feeding it probabilities produces incorrect gradients. Pass raw logits.

Forgetting zero_grad() causes gradients to accumulate. The model updates in the wrong direction because each step uses the sum of all previous gradients. The loss may appear to decrease while the model actually diverges.

A learning rate that is too high with Adam can cause early divergence, especially with small datasets. The default lr=0.001 is a reasonable starting point, but it is not universal.

SGD without momentum converges slowly on many problems. If you switch from Adam to SGD, expect to add momentum and possibly lower the learning rate.

Weight decay behaves differently between the two optimizers. optim.Adam applies L2 regularization to the gradient before the adaptive scaling, which is not the same as true decoupled weight decay. For that, use AdamW (optim.AdamW), which subtracts the decay term directly from the parameter. This distinction matters when you rely on weight decay for regularization.

When to Choose SGD Over Adam

The choice is not about which optimizer is better in general, but which fits the training setup.

Use SGD with momentum when you have a learning rate schedule that was tuned for it, or when memory is constrained and you cannot afford Adam's two state tensors per parameter.

Use Adam when you want a robust default that converges quickly without much tuning, or when your model has parameters at very different scales.

If you are unsure, start with Adam at lr=0.001 and CrossEntropyLoss for classification or MSELoss for regression. Once the model trains reliably, you can experiment with SGD and a schedule if you need the memory savings or want to test whether SGD generalizes better for your specific data.