SGD Variants - Momentum, RMSprop, Adam

By the end of this lesson, you will be able to...

• trace how SGD, momentum, RMSprop, and Adam turn a gradient into a parameter update;

• distinguish smoothing a direction from adapting its coordinate scale;

• diagnose a training symptom without assuming an optimizer can repair every problem.

Idea in one sentence: Optimizers do not create gradients; they keep state about past gradients so that the next parameter move can be smoother, differently scaled, or both.

Core Insight

After lesson 008, a learner has a gradient check that says the backward pass is locally consistent. In a notebook, the training log prints these first four y updates: -0.40, +0.40, -0.40, +0.40. The training loss falls slowly, the first coordinate moves steadily, and the second coordinate jumps left, right, left, right. The gradient is not necessarily wrong. It is giving a noisy or badly scaled direction.

The first reasonable model is plain stochastic gradient descent: take the negative gradient and multiply it by a learning rate. This is exactly the right baseline when the gradient is informative and one step size fits the coordinates. It becomes insufficient when consecutive mini-batches disagree in one direction or when one coordinate's gradients are much larger than another's.

Momentum, RMSprop, and Adam add state to the update rule. They do not change the loss, labels, architecture, or backpropagation. They decide how the already-computed gradient becomes movement. The numerical values below are an illustrative gradient stream, not measurements from a trained model.

The Two Coordinates That Behave Differently

Call the parameters x and y. Four consecutive mini-batches produce these gradients:

step 1: g = [0.2,  4.0]
step 2: g = [0.2, -4.0]
step 3: g = [0.2,  4.0]
step 4: g = [0.2, -4.0]

The x component consistently asks to move negative. The y component alternates. A picture of the implied pressure is:

x:  ->  ->  ->  ->       persistent direction
y:  up  down up  down    alternating direction

This is a teaching model for a narrow, noisy region of a loss surface. It does not prove that all alternating gradients come from the same geometry; data order, batch size, and the current parameters can all contribute. But it gives us a stable trace for comparing update rules.

Every rule in this lesson uses a learning rate eta = 0.1. Keep the distinction visible:

gradient: information from the current loss and batch
optimizer state: memory retained from earlier gradients
update: the change applied to the parameters

Plain SGD Has No Memory

SGD uses only the current gradient:

theta_new = theta - eta * g

For the first two steps, its updates are:

After two steps, x moved -0.04, while y returned to its original value. SGD faithfully followed each mini-batch, but the alternating y signal caused a wide zig-zag rather than durable progress.

SGD is not the “bad” optimizer in this story. Its rule is transparent, it adds no per-parameter history, and it gives a dependable baseline for a controlled comparison. The trade-off is that its simplicity leaves direction noise and unequal gradient scales untreated. A future lesson on schedules will address changing eta over time; the question here is what state to add between steps.

Momentum Remembers Direction

Momentum keeps one velocity-like buffer per parameter. One common sign convention is:

v = mu * v + g
theta_new = theta - eta * v

Start with v = [0, 0] and use mu = 0.9. Step one matches SGD because there is no earlier state:

v1 = [0.2, 4.0]
update1 = [-0.02, -0.40]

At step two, the old direction remains part of the state:

v2 = 0.9 * [0.2, 4.0] + [0.2, -4.0]
   = [0.38, -0.40]

update2 = -0.1 * v2
        = [-0.038, +0.040]

The persistent x component builds from 0.2 to 0.38, so it moves more decisively. The alternating y component is damped: its second movement is +0.04, not SGD's +0.40. Momentum is smoothing a direction over time. It is not measuring that y has a large scale and then dividing by it.

This helps when several consecutive gradients contain a useful common direction. It can still overshoot when the loss landscape changes quickly, and a large momentum coefficient can make a stale direction persist too long. The signal to watch is not “momentum is enabled”; it is whether loss and parameter updates oscillate, explode, or continue moving after the current gradient has changed sign.

RMSprop Adapts Coordinate Scale

RMSprop keeps a moving average of squared gradients for each coordinate:

r = rho * r + (1 - rho) * g^2
theta_new = theta - eta * g / (sqrt(r) + epsilon)

Use r = [0, 0], rho = 0.9, and a tiny epsilon only to avoid division by zero. After the first gradient:

r1 = 0.9 * [0, 0] + 0.1 * [0.2^2, 4.0^2]
   = [0.004, 1.6]

update1 ≈ [-0.316, -0.316]

The raw y gradient was twenty times larger than x, but RMSprop's denominator was also much larger for y. At step two, its state becomes approximately [0.0076, 3.04], yielding an update near [-0.229, +0.229].

RMSprop is adapting per-coordinate scale, not deciding whether the recent direction is persistent. It can reduce the domination of a coordinate with repeatedly large gradients, but it does not by itself accumulate a long-term direction the way momentum does. The initial updates in this toy trace look large because a new squared-gradient accumulator starts small; this is a reason to inspect the complete rule and its hyperparameters rather than treating one two-step table as a universal prediction.

Adam Keeps Both Kinds of State

Adam combines an average of gradients with an average of squared gradients. Its usual form also corrects the bias introduced by initializing both averages at zero:

m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * g^2
m_hat = m / (1 - beta1^t)
v_hat = v / (1 - beta2^t)
theta_new = theta - eta * m_hat / (sqrt(v_hat) + epsilon)

For an inspectable trace, choose beta1 = beta2 = 0.9, start with zero state, and keep eta = 0.1. On the first step, bias correction restores the current gradient scale:

m_hat1 = [0.2, 4.0]
v_hat1 = [0.04, 16.0]
update1 ≈ [-0.1, -0.1]

On step two, the corrected first-moment estimate is approximately [0.2, -0.2105]; the corrected squared estimate remains [0.04, 16.0]. The update is roughly:

update2 ≈ [-0.1, +0.0053]

The x direction stays strong because it is consistent. The y update is small because its direction has reversed and its recent squared magnitude is large. This one trace makes the combined mechanism visible: m_hat smooths direction; v_hat rescales each coordinate.

Adam's extra state does not make it a universal repair tool. An incorrect loss, broken data pipeline, bad gradient, or unsuitable learning rate remains a problem. Different tasks can favor different optimizer and learning-rate combinations, so describe Adam as a useful candidate to evaluate under named constraints, not as a default that proves training is healthy.

Optimizer State Is Evidence, Not a Verdict

The same raw gradient can lead to different updates because each optimizer carries a different history into the current step. This is why comparing only grad.norm() is incomplete. Log the update norm too, and, when a run behaves strangely, inspect the momentum buffer or moment estimates for the affected parameter group.

For example, a small gradient after a long run of large gradients can produce a smaller RMSprop or Adam update than raw SGD would make. That is not automatically a bug: the second-moment state is still saying that coordinate has recently been volatile. Conversely, a sudden large update may reflect a reset optimizer state, a changed learning rate, or a parameter that entered an optimizer group with different settings.

This also explains a reproducibility boundary. Two runs with identical model weights at step 100 are not necessarily in the same training state if one has momentum buffers and the other begins with zero buffers. Their next parameters may diverge even with the same batch. The useful operational question is therefore not “which optimizer name is best?” but “what state and update evidence did this run actually use?”

Read the Update Trace Before Switching Methods

The optimizer owns state as surely as a layer owns weights. When resuming training, losing momentum or moment estimates changes the next updates even if model weights are restored. This lesson does not cover checkpoint design in depth, but it explains why optimizer state belongs in a reproducible training artifact.

Checks Before the Next Batch

Check: In the two-step trace, why does momentum reduce the second y update from +0.40 to +0.04?

Think first, then reveal.

Answer: The prior positive velocity in y remains in the buffer and partly cancels the new negative gradient. Momentum is averaging directional history, so a one-step reversal does not completely determine the next update.

Check: Which state distinguishes RMSprop from momentum?

Answer: RMSprop stores a moving average of g^2 for each coordinate. That state estimates recent magnitude and becomes a denominator; momentum stores a signed direction buffer.

Check: Does a passed gradient check imply Adam will find a useful solution?

Answer: No. It says the local derivative matches the forward loss at a tested point. Optimizer behavior still depends on learning rate, data, initialization, objective, and the training path.

Trace It Yourself

Start with v = [0, 0] for momentum and use mu = 0.8, eta = 0.1. Apply the two gradients [1, 3] and [1, -3].

1. Calculate the velocity and update after each step.
2. Compare the second y update with plain SGD.
3. State whether the persistent x direction receives more or less movement than in plain SGD.
4. Name one boundary: when would this two-step trace be insufficient evidence for choosing momentum?

Model answer: v1 = [1, 3], so the first update is [-0.1, -0.3]. Then v2 = 0.8*[1, 3] + [1, -3] = [1.8, -0.6], giving [-0.18, +0.06]. Plain SGD would make the second y update +0.3, so momentum damped the reversal. It also accelerates the consistent x direction. This trace alone is insufficient if gradients change for another reason—such as a broken batch pipeline, an overly large learning rate, or a different validation outcome—so compare controlled runs rather than selecting from one toy sequence.

Resources

• [DOCUMENTATION] PyTorch: Optimizers — Focus: optimizer state and the parameters it updates.
• [DOCUMENTATION] PyTorch: SGD — Focus: learning rate, momentum, and the update interface.
• [PAPER] Adam: A Method for Stochastic Optimization — Focus: first- and second-moment estimates with bias correction.
• [TUTORIAL] CS231n: Learning and Evaluation — Focus: update traces and optimization diagnostics.
• [BOOK] Deep Learning, Chapter 8 — Focus: optimization limits and gradient-based methods.

Key Takeaways

• SGD applies the current gradient directly; it is a clear baseline with no per-parameter history.
• Momentum smooths direction across steps, while RMSprop adapts updates to recent coordinate scale.
• Adam combines both kinds of state and uses bias correction, but it does not replace a valid gradient, loss, or evaluation process.
• Update norms, moment state, loss, and validation evidence reveal more than an optimizer name alone.