Optimization Algorithms in Deep Learning: SGD, Momentum and Adam

Training a neural network means finding millions of weights that make a loss function as small as possible. We cannot solve for them with algebra, so we walk downhill, one small step at a time. The rule that decides the size and direction of each step is the optimizer, and the choice between plain SGD, momentum and Adam can decide whether a model trains in an hour, a week or not at all.

The Optimization Problem

Given weights \( \mathbf{w} \) and a training set of \( N \) examples, the loss is an average of per-example losses:

$$ L(\mathbf{w}) = \frac{1}{N} \sum_{i=1}^{N} \ell(\mathbf{w}; \mathbf{x}_i, y_i). $$

The gradient \( \nabla L(\mathbf{w}) \) points in the direction of steepest increase, so moving against it reduces the loss. Every optimizer below is a variation on that idea.

Gradient Descent and SGD

Plain (batch) gradient descent updates

$$ \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \nabla L(\mathbf{w}_t), $$

where \( \eta \) is the learning rate. Computing the exact gradient requires a pass over all \( N \) examples, which is very expensive when \( N \) is in the millions.

Stochastic gradient descent (SGD) replaces the full gradient with an estimate from a random mini-batch \( B \) of, say, 32 or 256 examples:

$$ \mathbf{g}_t = \frac{1}{|B|} \sum_{i \in B} \nabla \ell(\mathbf{w}_t; \mathbf{x}_i, y_i), \qquad \mathbf{w}_{t+1} = \mathbf{w}_t - \eta\, \mathbf{g}_t . $$

The estimate is noisy but unbiased: on average it equals the true gradient. The noise even helps by jostling the weights out of poor flat regions.

The learning rate is the most important setting. Too small and training crawls; too large and the loss oscillates or diverges. For the simple function \( f(w) = a w^2 \), gradient descent converges only if \( 0 < \eta < 1/a \), since each step multiplies \( w \) by \( (1 - 2a\eta) \).

Momentum

Picture a ball rolling down a long, narrow valley. Plain SGD zigzags across the steep walls while making slow progress along the gentle floor. Momentum gives the ball inertia by keeping a running "velocity":

$$ \mathbf{v}_t = \beta \mathbf{v}_{t-1} + \mathbf{g}_t, \qquad \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \mathbf{v}_t , $$

with \( \beta \) typically 0.9. Gradients that consistently point the same way accumulate and speed up, while gradients that flip sign cancel out, damping the zigzag. In a steady slope the velocity approaches \( \mathbf{g}/(1 - \beta) \), so with \( \beta = 0.9 \) the effective step is up to ten times larger.

Adam

Different weights often need different step sizes: a rarely active feature produces small gradients and deserves bigger steps. Adam (adaptive moment estimation), introduced by Diederik Kingma and Jimmy Ba in 2014, tracks two moving averages for each weight:

$$ \mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1)\mathbf{g}_t, \qquad \mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2)\mathbf{g}_t^2 . $$

Here \( \mathbf{m}_t \) estimates the mean gradient (like momentum) and \( \mathbf{v}_t \) estimates the mean squared gradient; the square is taken element by element. Both start at zero, which biases them towards zero early on, so Adam corrects them:

$$ \hat{\mathbf{m}}_t = \frac{\mathbf{m}_t}{1 - \beta_1^t}, \qquad \hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_2^t}, \qquad \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}. $$

Default values are \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \) and \( \epsilon = 10^{-8} \). Dividing by \( \sqrt{\hat{v}} \) normalises each weight's step, so the step size is roughly \( \eta \) regardless of how large or small the raw gradients are.

Worked Example 1: Gradient Descent versus Momentum

Minimise \( f(w) = (w - 3)^2 \), whose minimum is at \( w = 3 \). The gradient is \( f'(w) = 2(w - 3) \). Start at \( w_0 = 0 \) with \( \eta = 0.1 \).

Plain gradient descent.

• Step 1: \( g = 2(0 - 3) = -6 \). \( w_1 = 0 - 0.1(-6) = 0.6 \).
• Step 2: \( g = 2(0.6 - 3) = -4.8 \). \( w_2 = 0.6 + 0.48 = 1.08 \).
• Step 3: \( g = 2(1.08 - 3) = -3.84 \). \( w_3 = 1.08 + 0.384 = 1.464 \).

Each step closes 20% of the remaining gap, since \( w - 3 \) is multiplied by \( 1 - 2(0.1) = 0.8 \).

Momentum with \( \beta = 0.9 \), starting from \( v_0 = 0 \).

• Step 1: \( g = -6 \). \( v_1 = 0.9(0) - 6 = -6 \). \( w_1 = 0 - 0.1(-6) = 0.6 \).
• Step 2: \( g = -4.8 \). \( v_2 = 0.9(-6) - 4.8 = -5.4 - 4.8 = -10.2 \). \( w_2 = 0.6 + 1.02 = 1.62 \).
• Step 3: \( g = 2(1.62 - 3) = -2.76 \). \( v_3 = 0.9(-10.2) - 2.76 = -9.18 - 2.76 = -11.94 \). \( w_3 = 1.62 + 1.194 = 2.814 \).

After three steps momentum is at 2.814, very close to 3, while plain gradient descent is at 1.464. Sanity check: momentum keeps its velocity, so it will overshoot past 3 on the next step and then oscillate inward. That is the typical behaviour of a heavy ball, and the reason \( \beta \) is kept below 1.

Worked Example 2: Two Steps of Adam

Use the same function and start, with \( \eta = 0.1 \), \( \beta_1 = 0.9 \), \( \beta_2 = 0.999 \), and \( m_0 = v_0 = 0 \). We ignore \( \epsilon \), which is negligible here.

Step 1 (\( t = 1 \)): \( g_1 = -6 \).

• \( m_1 = 0.9(0) + 0.1(-6) = -0.6 \)
• \( v_1 = 0.999(0) + 0.001(36) = 0.036 \)
• Bias correction: \( \hat{m}_1 = -0.6/(1 - 0.9) = -6 \), \( \hat{v}_1 = 0.036/(1 - 0.999) = 36 \)
• Update: \( w_1 = 0 - 0.1 \times (-6)/\sqrt{36} = 0 + 0.1 = 0.1 \)

The first Adam step always has size exactly \( \eta \), whatever the gradient's magnitude.

Step 2 (\( t = 2 \)): \( g_2 = 2(0.1 - 3) = -5.8 \).

• \( m_2 = 0.9(-0.6) + 0.1(-5.8) = -0.54 - 0.58 = -1.12 \)
• \( v_2 = 0.999(0.036) + 0.001(33.64) = 0.035964 + 0.03364 = 0.069604 \)
• Bias correction: \( 1 - 0.9^2 = 0.19 \), so \( \hat{m}_2 = -1.12/0.19 \approx -5.8947 \). Also \( 1 - 0.999^2 = 0.001999 \), so \( \hat{v}_2 = 0.069604/0.001999 \approx 34.8194 \) and \( \sqrt{\hat{v}_2} \approx 5.9008 \).
• Update: \( w_2 = 0.1 + 0.1 \times 5.8947/5.9008 \approx 0.1 + 0.0999 = 0.1999 \)

Adam takes steps of almost exactly 0.1 here. On this simple problem that is slower than plain gradient descent, because the large, consistent gradient would have allowed bigger steps. Adam shines instead on huge, badly scaled problems where some weights see gradients of \( 10^{-5} \) and others of \( 10 \); there its normalisation keeps every weight learning at a sensible pace.

Practical Guidance

• Adam is the usual default for transformers and many new architectures, often with \( \eta \) around \( 10^{-4} \) to \( 10^{-3} \).
• SGD with momentum remains popular for image classification CNNs and is sometimes reported to generalise slightly better when carefully tuned.
• AdamW separates weight decay from the gradient-based update and is standard for large language models.
• Learning-rate schedules matter as much as the optimizer: a short warm-up followed by gradual decay (step, cosine or linear) is common.
• If the loss becomes NaN or explodes, lower the learning rate first, then consider gradient clipping.

Did You Know?

• Stochastic approximation, the theory behind SGD, dates back to Herbert Robbins and Sutton Monro in 1951, long before deep learning.
• Adam stores two extra numbers per weight, so for a model with a billion parameters, the optimizer state alone needs two billion additional values.
• Nesterov momentum evaluates the gradient at a "look-ahead" point \( \mathbf{w} - \eta\beta\mathbf{v} \), which often reduces overshooting.

Key Takeaways

• SGD uses noisy mini-batch gradients, trading accuracy per step for far cheaper steps.
• Momentum accumulates past gradients, accelerating along consistent directions and damping oscillations.
• Adam combines momentum with per-weight scaling by the root mean squared gradient, plus bias correction.
• Adam's first step always has size \( \eta \); its strength is robustness on badly scaled problems.
• The learning rate and its schedule are usually more important than the exact optimizer chosen.

Struggling with an optimization derivation or a deep learning project? Amathify's experts and PhD scholars provide clear, step-by-step solutions; contact us on WhatsApp or through the Ask a Question form.