Gradient Descent and Modern Optimizers

You already know the core loop: compute a gradient, step against it, repeat. What that loop leaves unanswered is why almost nothing at production scale actually uses it in that plain form. GPT-class models, Llama, and the Indian-language models coming out of labs like Sarvam AI are all trained with a variant called Adam — and that single choice does two things at once. It changes what a single update step actually computes, and it roughly doubles the GPU memory a training run needs, independent of the model's own weights. Both consequences follow from the same piece of algebra, and neither is visible if you only ever look at the high-level "loss goes down" curve. This chapter opens the optimizer up: the exact update Adam performs, why it needs a correction term that has nothing to do with the loss function, how that correction was later shown to break convergence in a specific, constructible case, and why a production engineer sizing a GPU cluster has to count optimizer state as carefully as they count parameters.

Inside Adam: two moving averages and a cold-start problem

Plain gradient descent uses only the current gradient g_t. Adam (Kingma & Ba, "Adam: A Method for Stochastic Optimization," ICLR 2015) instead tracks two running statistics of the gradient over time, for every parameter independently:

m_t = β1 · m_(t-1) + (1 - β1) · g_t        # exponential moving average of the gradient itself
v_t = β2 · v_(t-1) + (1 - β2) · g_t²       # exponential moving average of the squared gradient

m_t is a smoothed estimate of the gradient's direction (it plays a role similar to momentum). v_t is a smoothed estimate of the gradient's typical squared magnitude — an uncentered variance. The defaults from the original paper, still the near-universal default today, are β1 = 0.9 and β2 = 0.999. Both m_0 and v_0 are initialized to zero.

That zero initialization is the problem. Because β2 = 0.999 is so close to 1, v_t barely moves away from 0 in the first few steps — at t = 1, v_1 = (1 - 0.999) · g_1² = 0.001 · g_1², which is a wildly biased underestimate of the true squared-gradient scale. If Adam divided by the raw v_t, the first several updates would be enormous (dividing by something near zero), which is exactly the instability the algorithm is supposed to prevent. Kingma and Ba's fix is a bias-correction step applied at every iteration:

m̂_t = m_t / (1 - β1^t)
v̂_t = v_t / (1 - β2^t)
θ_t = θ_(t-1) - α · m̂_t / (√v̂_t + ε)

Here α is the base learning rate and ε (typically 1e-8) exists purely to stop division by zero when v̂_t is tiny. The correction terms 1 - β1^t and 1 - β2^t both start small (near 0 when t is small) and grow toward 1 as t increases, which inflates m_t and v_t exactly enough to cancel the zero-initialization bias, then fades out and does almost nothing once training has run for a few hundred steps.

Tracing Adam by hand

Take a one-parameter toy loss, f(θ) = (θ - 3)², so g(θ) = 2(θ - 3), starting at θ_0 = 1.0 with β1 = 0.9, β2 = 0.999, α = 0.1, ε = 1e-8.

Step 1. g_1 = 2(1.0 - 3) = -4.0.

m_1 = 0.9·0 + 0.1·(-4.0)        = -0.4
v_1 = 0.999·0 + 0.001·(-4.0)²   = 0.016
m̂_1 = -0.4 / (1 - 0.9^1)  = -0.4 / 0.1  = -4.0
v̂_1 = 0.016 / (1 - 0.999^1) = 0.016 / 0.001 = 16.0
update = 0.1 · (-4.0) / (√16.0 + 1e-8) = 0.1 · (-4.0/4.0) ≈ -0.1
θ_1 = 1.0 - (-0.1) = 1.1

Step 2. g_2 = 2(1.1 - 3) = -3.8.

m_2 = 0.9·(-0.4) + 0.1·(-3.8)   = -0.74
v_2 = 0.999·0.016 + 0.001·(-3.8)² = 0.030424
m̂_2 = -0.74 / (1 - 0.81)   = -0.74 / 0.19   ≈ -3.8947
v̂_2 = 0.030424 / (1 - 0.998001) = 0.030424 / 0.001999 ≈ 15.2196
√v̂_2 ≈ 3.9012
update = 0.1 · (-3.8947) / (3.9012 + 1e-8) ≈ 0.1 · (-0.9983) ≈ -0.0998
θ_2 = 1.1 - (-0.0998) = 1.1998

Here is the same trace as runnable code — every value it prints matches the hand computation above:

def adam_step(theta, m, v, t, g, beta1=0.9, beta2=0.999, alpha=0.1, eps=1e-8):
    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 = theta - alpha * m_hat / (v_hat ** 0.5 + eps)
    return theta, m, v

theta, m, v = 1.0, 0.0, 0.0
for t in range(1, 3):
    g = 2 * (theta - 3)
    theta, m, v = adam_step(theta, m, v, t, g)
    print(f"t={t}: g={g:.4f}, theta={theta:.4f}")

# t=1: g=-4.0000, theta=1.1000
# t=2: g=-3.8000, theta=1.1998

Why the first step is always about α

Look at what happened at t = 1: m̂_1 came out to exactly g_1, and v̂_1 came out to exactly g_1². This is not a coincidence of the numbers chosen — it is an algebraic identity. At t = 1, m_1 = (1 - β1)·g_1, and dividing by the correction term (1 - β1^1) = (1 - β1) cancels exactly, leaving m̂_1 = g_1 for any valid β1. The same cancellation gives v̂_1 = g_1² for any valid β2. So the very first update is

m̂_1 / √v̂_1 = g_1 / |g_1| = sign(g_1)

— a pure sign, with the gradient's magnitude divided completely out. The first step Adam ever takes has size α (up to the negligible ε), no matter how steep or shallow the loss surface is at the starting point. Contrast that with plain gradient descent on the same problem: θ_1 = θ_0 - α·g_1 = 1.0 - 0.1·(-4.0) = 1.4, a jump four times larger than Adam's, purely because the gradient happened to be steep there. If the loss had instead been f(θ) = (θ-3)⁴, giving g_1 = 4(1-3)³ = -32, plain gradient descent would leap by 0.1 × 32 = 3.2, while Adam's first step would still be ≈ 0.1. This is precisely why Adam tolerates a single global learning rate across layers whose gradient scales differ by orders of magnitude — attention output projections versus embedding tables in a transformer, for instance — where plain SGD with one fixed α would either explode on the steep layers or crawl on the flat ones.

AdamW: decoupling weight decay from the adaptive scaling

Weight decay (shrinking parameters toward zero each step, to fight overfitting) predates Adam and was originally bolted on the same way it works for plain SGD: add λ·θ to the gradient before doing anything else, so the "gradient" fed into m_t and v_t becomes g_t + λ·θ_(t-1). Loshchilov and Hutter ("Decoupled Weight Decay Regularization," ICLR 2019) showed this interacts badly with Adam specifically: because the decay term now flows through v_t and gets divided by √v̂_t, a parameter with a history of large gradients has its weight decay suppressed, and a parameter with small gradients has its weight decay amplified — an effect nobody asked for and that has nothing to do with regularization strength. AdamW's fix is to keep the decay term separate from the moment estimates entirely:

θ_t = θ_(t-1) - α · ( m̂_t / (√v̂_t + ε) + λ · θ_(t-1) )

m_t and v_t are computed from the raw loss gradient only; λ·θ_(t-1) is subtracted afterward, at the same fixed rate for every parameter regardless of its gradient history. AdamW is now the default optimizer for essentially every transformer pretraining run you will read about, precisely because decay strength stops depending on gradient statistics it was never meant to depend on.

When Adam doesn't converge: AMSGrad

Reddi, Kale, and Kumar ("On the Convergence of Adam and Beyond," ICLR 2018) proved something sharper than a heuristic complaint: they constructed an explicit, simple online convex optimization problem — a periodic sequence of gradients, occasionally large and positive, mostly small and negative — on which Adam's iterate provably drifts to the wrong answer and never recovers, even though the problem is convex, the setting Adam is supposed to handle safely. The mechanism is that v_t is a plain exponential moving average, so it can decrease again after a large gradient fades out of the window; when it decreases, the effective step size α/√v̂_t increases at exactly the wrong moment, undoing the correction the large gradient was supposed to provide. Their fix, AMSGrad, replaces the bias-corrected second moment with a running maximum:

v̂_t = max(v̂_(t-1), v_t)

v̂_t is now non-decreasing by construction, so the effective step size can only shrink or stay flat over time, never spike back up. AMSGrad is not the default in most production training stacks today — in practice, warmup schedules and gradient clipping paper over the failure mode Reddi et al. constructed — but it is the standard citation any time someone claims Adam "always converges," and the counterexample is a clean illustration that adaptive step sizes are a heuristic, not a proof.

The other cost of "adaptive": GPU memory

Every parameter Adam trains needs its own m_t and v_t stored between steps — two extra numbers per weight, on top of the weight itself and its gradient. In a mixed-precision training run (the standard setup for any model past a few hundred million parameters), the breakdown per parameter is:

• a half-precision (fp16/bf16) copy of the weight, used for the forward and backward pass — 2 bytes
• a half-precision gradient — 2 bytes
• a full-precision (fp32) master copy of the weight, updated by the optimizer and rounded back to fp16 — 4 bytes
• the fp32 first moment m_t — 4 bytes
• the fp32 second moment v_t — 4 bytes

That totals 16 bytes per parameter, before a single activation is stored. For a 7.5-billion-parameter model, 7.5 × 10⁹ × 16 bytes ≈ 120 GB — more than the 80 GB on a single H100, before you have run one training step. This is exactly why large-scale training cannot simply be "one big GPU with more RAM": the optimizer state itself has to be sharded across GPUs (the technique behind ZeRO — Rajbhandari, Rasley, Ruwase, and He, "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models," SC20 — and PyTorch's FSDP), with each GPU holding only a slice of m, v, and the fp32 master weights, and full copies reassembled only transiently when needed. Switching to plain SGD with momentum drops the per-parameter cost to 12 bytes (no second moment needed): fp16 weight + fp16 grad + fp32 master weight + fp32 momentum buffer. On the same 7.5B model that is 7.5 × 10⁹ × 12 = 90 GB, a 30 GB saving — one reason some very large training runs still reach for plain momentum-SGD or memory-lighter Adam variants (like 8-bit Adam) once model size collides with GPU memory limits, trading optimizer sophistication for the ability to fit the run at all.

Misconception: "adaptive" means the learning rate doesn't matter

A common assumption is that because Adam scales each parameter's step individually, the base learning rate α is a minor knob you can leave at a library default. This is wrong on two counts. First, mechanically: α is not divided out by the adaptive machinery — it is the multiplier standing outside m̂_t/(√v̂_t + ε), so doubling α still roughly doubles every step, exactly as it would for plain gradient descent. The adaptivity rescales relative to each parameter's own gradient history, not relative to the value you chose for α itself. Second, empirically: Wilson, Roelofs, Stern, Srebro, and Recht ("The Marginal Value of Adaptive Gradient Methods in Machine Learning," NeurIPS 2017) showed that models trained with Adam can reach lower training loss than SGD with momentum yet generalize worse on held-out data, on several standard benchmarks — meaning the choice of optimizer, and the learning rate schedule paired with it, is a genuine hyperparameter decision with measurable downstream consequences, not a solved problem you can ignore once you've picked "Adam" from a list.

Active recall

1. In the worked trace above, θ_0 = 1.0 under f(θ) = (θ-3)² gave g_1 = -4.0 and a first Adam step of ≈ -0.1. If the loss were instead f(θ) = (θ-3)⁴ (so g_1 = -32), what would Adam's first step be, and why?
2. Suppose ε is changed from 1e-8 to 1.0 in the same trace (θ_0=1.0, g_1=-4). Recompute θ_1. Does the "first step ≈ α" property still hold exactly?
3. Plain gradient descent with α=0.1 and g_1=-4 moves θ_0=1.0 to θ_1=1.4 in one step; Adam moved it only to ≈1.1. Why does this matter for training stability when different layers of a network have very different gradient magnitudes?
4. A team trains a model with plain "Adam + L2 regularization" (weight decay folded into the gradient before the moment estimates) and then switches to AdamW with the same λ. Does the optimizer-state memory change? What actually changes in the update?
5. If β2 is lowered from 0.999 to 0.9, does the exact cancellation m̂_1 = g_1, v̂_1 = g_1² at t=1 still hold? What changes from t=2 onward instead?
6. Using the 16-bytes-per-parameter breakdown, compute the optimizer-plus-weight memory for a 13-billion-parameter model trained with plain SGD+momentum (12 bytes/parameter: fp16 weight, fp16 grad, fp32 master weight, fp32 momentum — no second moment). How much memory is saved compared to training the same model with Adam?

Worked answers

1. Still ≈ -0.1. At t=1, m̂_1 = g_1 and v̂_1 = g_1² for any nonzero g_1, so m̂_1/√v̂_1 = g_1/|g_1| = sign(g_1) = -1 regardless of magnitude. The update is α · (-1) = -0.1; plain gradient descent, by contrast, would take a step of 0.1 × 32 = 3.2 — 32 times larger, purely from the steeper loss.

2. With ε=1.0: v̂_1 = 16.0, √v̂_1 = 4.0, update = 0.1 · (-4.0)/(4.0+1.0) = 0.1·(-0.8) = -0.08, so θ_1 = 1.08, not 1.1. The property breaks because it relied on ε being negligible relative to √v̂_t; with ε comparable in size to the gradient scale, the step is damped below α.

3. Adam's per-parameter normalization caps the first (and typically early) update near α regardless of that parameter's local gradient steepness, so one global α can be used safely across layers whose gradients differ by orders of magnitude (e.g. embedding tables versus attention output projections in a transformer). Plain gradient descent with one shared α is at the mercy of the single steepest gradient in the model — too large an α for the steep layers explodes them, too small an α for the flat layers stalls them.

4. Memory is unchanged — both track the same m_t and v_t shapes, so the 16-bytes-per-parameter footprint is identical. What changes is only where λθ enters the arithmetic: Adam+L2 adds it into g_t before computing m_t/v_t, so the effective decay strength gets divided by √v̂_t and ends up parameter-dependent (suppressed for parameters with a history of large gradients); AdamW applies -α·λ·θ_(t-1) directly, after the adaptive term, at the same rate for every parameter.

5. Yes — the cancellation at t=1 is algebraic ((1-β)·g / (1-β) = g) and holds for any β1, β2 ∈ [0,1), so it is unaffected by lowering β2. What changes is the behavior from t=2 onward: β2=0.9 makes v_t weight the most recent squared gradient much more heavily (it forgets older ones faster). In this specific trace, where |g_t| decreases only slowly step to step (g_1=-4.0, g_2=-3.8), that faster forgetting actually keeps v̂_t tracking closer to the still-large recent g_t² than the slower-forgetting β2=0.999 does, so with β2=0.9 the step stays nearer to α for longer, not farther from it: deviation from α is only 0.000033 at t=2 and 0.000092 at t=3 for β2=0.9, versus 0.000166 and 0.000457 for β2=0.999. The direction of this effect depends on whether the gradient magnitude is growing or shrinking over the relevant steps, not on β2 alone — a faster-forgetting v_t tracks whatever the recent gradient scale is doing more tightly, for better or worse.

6. SGD+momentum: 13×10⁹ × 12 bytes = 156 GB. Adam: 13×10⁹ × 16 bytes = 208 GB. Savings: 13×10⁹ × 4 bytes = 52 GB.

Think About It

Think about this: How would you explain gradient descent and modern optimizers to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind gradient descent and modern optimizers, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.