Optimization: Adam, SGD, and Learning Rate Schedules

Open Swiggy or Zomato and place a food order. Before you have even picked a payment method, the app already shows a number: "Delivery in 34 minutes." That estimate comes from a model, and months before it became reliable, the model was wrong a great deal of the time. Every time it was wrong, a training algorithm nudged the model's internal numbers a little closer to the truth — a little less weight on this feature, a little more on that one. That repeated nudging is optimization, and this chapter is about the two questions every optimizer must answer on every single nudge: which direction should the numbers move, and how big should the step be? SGD and Adam are two carefully engineered answers to those questions, and a learning rate schedule is a plan for how the size of the step should change as training goes on.

Loss, Gradients, and the Direction of Steepest Descent

Every model starts with a loss function — a single number that measures how wrong the model's predictions are. If a model predicts a delivery time of 40 minutes for an order that actually took 34, a common choice is the squared error, (40 − 34)² = 36. Training a model means searching for the parameter values that make this number as small as possible, averaged across every example the model has seen.

Suppose, for a moment, that a model has just one learnable parameter, w. The loss is then a function of a single variable, L(w), and calculus already tells you how to search for its minimum: the derivative dL/dw measures how steeply L rises as w increases. If the derivative is positive, decreasing w decreases the loss; if it is negative, increasing w decreases the loss. Either way, moving in the direction opposite the derivative always moves you toward lower loss, at least for a small enough step. This is the whole idea behind gradient descent:

w_t = w_t−1 − η · dL/dw, with the derivative evaluated at w_t−1, the current value

Here η ("eta") is the learning rate, a small positive number that controls how big each step is. Real models don't have one parameter — a small neural network can easily have tens of thousands, and large ones have billions. The loss is then a function of many variables, L(w₁, w₂, ..., w_n), and its gradient, written ∇L, is the vector of partial derivatives (∂L/∂w₁, ∂L/∂w₂, ..., ∂L/∂w_n) — one slope per parameter. A short argument from vectors explains why the negative gradient is exactly the right direction to move in: for a small step Δw, the loss changes by approximately L(w + Δw) ≈ L(w) + ∇L(w)·Δw. Among all directions Δw of a fixed small length, this dot product is most negative — meaning the loss drops fastest — when Δw points exactly opposite to ∇L(w). Every optimizer in this chapter is, at its core, still doing exactly this: computing a gradient and stepping against it. What changes from one optimizer to the next is how that step is shaped before it is taken, and every update rule below is written for a single scalar w purely so the arithmetic stays visible — in a real model, the same equation applies independently, element by element, to every one of the millions of entries in that gradient vector at once.

From Gradient Descent to Stochastic Gradient Descent

Computing dL/dw exactly requires averaging the gradient over every training example — for a delivery-time model trained on a year of orders, that could be tens of millions of rows, recomputed before every single update. That is full-batch gradient descent, and it is correct but painfully slow: one parameter update per pass through the entire dataset.

Stochastic gradient descent (SGD) makes a trade: instead of computing the gradient over the whole dataset, estimate it from a small, randomly chosen mini-batch — say, 64 or 256 orders — and take a step using that estimate. It is noisier than the true gradient, but it is cheap, and with millions of examples available, thousands of these noisy updates fit into the time one full-batch update would have taken. In the strict historical sense, "stochastic" meant a batch of exactly one example; in modern deep learning frameworks, the optimizer everyone calls "SGD" almost always operates on mini-batches — the "stochastic" part refers to the randomness in which examples land in each batch, not to the batch size being one. One full pass through the entire training set is called an epoch; each individual mini-batch update within that pass is a step, or iteration.

The noise in SGD is not purely a cost. A perfectly smooth, deterministic gradient can walk straight into a shallow local dip or a flat saddle region and stall there; the randomness of mini-batch sampling gives each step a slightly different direction, which is often enough to jostle the parameters out of such traps. The same noise, though, means SGD rarely settles precisely at the minimum — it tends to jitter around it, which is one of the reasons learning rate schedules exist, as this chapter returns to later.

Why Plain SGD Struggles: Ravines and Ill-Conditioned Losses

Imagine a model that flags suspicious UPI transactions using two features: transaction amount, which ranges from a few rupees to several lakh, and time-since-last-transaction, measured in seconds. Because these features live on wildly different scales, the loss surface with respect to their two weights is typically not a nice round bowl — it is a long, narrow valley, steep in one direction and almost flat in the other. Plain SGD, which takes an equal-sized step along every direction, zig-zags sharply back and forth across the steep walls of the valley while creeping forward at a snail's pace along its shallow floor. Push the learning rate up to speed up that crawl, and the steep direction starts to overshoot and oscillate instead. This tension — one learning rate serving every parameter, even though different parameters may need very different step sizes — is precisely what the next two ideas, momentum and adaptive learning rates, are designed to fix.

Momentum: Carrying Velocity Downhill

Momentum borrows an idea from physics: instead of reacting only to the current slope, keep a running memory of the direction you have been moving in, and let that memory smooth out the current step. One standard way to express this is as an exponential moving average of past gradients:

v_t = β · v_t−1 + (1 − β) · g_t; then w_t = w_t−1 − η · v_t

where g_t is the gradient at step t and β (typically around 0.9) controls how much of the past is remembered. In the narrow-valley picture above, the back-and-forth components across the steep walls point in opposite directions on successive steps, so they partially cancel out in the running average v_t; the components along the shallow floor keep pointing the same way step after step, so they reinforce and accumulate. The net effect is exactly what you would want: the zig-zag gets damped, and progress along the flat direction speeds up — much like a ball rolling downhill picks up speed on a gentle, consistent slope but does not fully reverse direction every time it clips a small bump.

Adaptive Learning Rates: AdaGrad and RMSProp

Momentum smooths the direction of the step but still applies one global learning rate η to every parameter. A different fix is to give every parameter its own effective learning rate, adjusted by how large its gradients have historically been. AdaGrad (Duchi, Hazan, and Singer, 2011) does this by keeping a running total of squared gradients for each parameter and dividing the learning rate by its square root, so parameters with large, frequent gradients get smaller effective steps and parameters with small, rare gradients — like the weight on an uncommon feature — get comparatively larger ones. The catch is that AdaGrad's running total only ever grows, so the effective learning rate keeps shrinking across training and can eventually crawl to a near standstill, long before the model has finished learning.

RMSProp, introduced by Geoffrey Hinton in his Coursera lecture on neural networks rather than in a published paper, fixes exactly this. Instead of an ever-growing running total, RMSProp keeps an exponential moving average of squared gradients — the same kind of "forgetful" average used for momentum, just applied to g_t² instead of g_t:

v_t = ρ · v_t−1 + (1 − ρ) · g_t²; then w_t = w_t−1 − (η / (√v_t + ε)) · g_t

Old, stale gradients now decay out of the average instead of accumulating forever, so the effective learning rate can rise again if gradients shrink, rather than being throttled permanently. The small constant ε (often 10⁻⁸) exists purely so the denominator never touches zero.

Adam: Adaptive Moment Estimation

Adam — short for "Adaptive Moment Estimation," introduced by Diederik Kingma and Jimmy Ba in 2015 — combines both ideas at once. It keeps a first moment estimate m_t (an exponential moving average of the gradient, exactly like momentum) and a second moment estimate v_t (an exponential moving average of the squared gradient, exactly like RMSProp), and uses the first to decide the step's direction and the second to scale it, per parameter:

initialize: w = w0, m = 0, v = 0, t = 0
hyperparameters: lr (learning rate), beta1, beta2 in [0, 1), eps (small constant)
common defaults: lr = 0.001, beta1 = 0.9, beta2 = 0.999, eps = 1e-8

repeat:
    t = t + 1
    g  = gradient of L at w                # slope of the loss right now
    m  = beta1 * m + (1 - beta1) * g       # running average of the gradient
    v  = beta2 * v + (1 - beta2) * g**2    # running average of the squared gradient
    m_hat = m / (1 - beta1**t)             # bias-corrected m
    v_hat = v / (1 - beta2**t)             # bias-corrected v
    w  = w - lr * m_hat / (v_hat**0.5 + eps)
until w stops changing much

The one genuinely new ingredient here, beyond "momentum plus RMSProp," is bias correction. Both m and v start at exactly zero, and an exponential moving average that starts at zero is, for the first several steps, still mostly zero — it has not had time to move toward the true recent average yet. Dividing by (1 − β₁^t) and (1 − β₂^t) corrects for exactly this start-up bias, and the correction matters most when t is small, fading away as t grows large, since both β₁^t and β₂^t shrink toward zero as t increases. The worked example below shows precisely why this correction cannot be skipped.

Worked Example: Tracing SGD and Adam by Hand

Consider the simplest possible version of the delivery-time model: predicted time equals w times distance, ŷ = w·x (the model's one prediction), with exactly one training example — a 1 km order that actually took 4 minutes (x = 1, y = 4). The squared-error loss as a function of the single weight w is L(w) = (w − 4)², with derivative dL/dw = 2(w − 4). Starting from w₀ = 0, plain gradient descent with learning rate η = 0.1 proceeds exactly as the update rule says:

• Step 1: gradient = 2(0 − 4) = −8, so w₁ = 0 − 0.1×(−8) = 0.8000
• Step 2: gradient = 2(0.8 − 4) = −6.4, so w₂ = 0.8 − 0.1×(−6.4) = 1.4400
• Step 3: gradient = 2(1.44 − 4) = −5.12, so w₃ = 1.44 − 0.1×(−5.12) = 1.9520

Now trace Adam on the same loss, same starting point, same η = 0.1, with the standard β₁ = 0.9, β₂ = 0.999, ε = 10⁻⁸. At step 1 the gradient is again g₁ = −8. The moment estimates are m₁ = 0.9×0 + 0.1×(−8) = −0.8000 and v₁ = 0.999×0 + 0.001×64 = 0.0640. Bias-corrected: m̂₁ = −0.8 / (1 − 0.9) = −8.0000, and v̂₁ = 0.064 / (1 − 0.999) = 64.0000. Notice what just happened — because m₁ and v₁ are each built from a single term, dividing by (1 − β₁) and (1 − β₂) exactly undoes that weighting and hands back the raw gradient and its square. The update is w₁ = 0 − 0.1 × (−8) / (√64 + ε) = 0 − 0.1×(−1.0) = 0.1000.

Compare that to what would have happened without bias correction: the raw, uncorrected ratio is m₁/√v₁ = −0.8/√0.064 ≈ −3.1623, more than three times larger in magnitude than the corrected ratio of −1.0. That is not a coincidence — it follows from (1 − β₁)/√(1 − β₂) = 0.1/√0.001 ≈ 3.16, a ratio that depends only on the two β values, never on the gradient itself. Skip bias correction, and the very first Adam step on any freshly initialized parameter, in any model, would overshoot by more than 3× — not because of anything the data said, but purely because β₂ sits closer to 1 than β₁ does, and correcting for exactly that gap is the entire point of the bias-correction step.

Continuing the same four-line calculation for step 2: g₂ = 2(0.1 − 4) = −7.8, m₂ = 0.9×(−0.8) + 0.1×(−7.8) = −1.5000, v₂ = 0.999×0.064 + 0.001×7.8² = 0.124776. Bias-corrected, m̂₂ = −1.5/(1 − 0.81) = −7.8947, v̂₂ = 0.124776/(1 − 0.998001) ≈ 62.419, and √v̂₂ ≈ 7.9006, giving w₂ = 0.1 − 0.1×(−7.8947)/7.9006 ≈ 0.1999. A third step, by the same arithmetic, lands at w₃ ≈ 0.2997. The short program below runs the identical calculation and confirms it:

def adam_step(w, grad, m, v, t, lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8):
    m = beta1 * m + (1 - beta1) * grad
    v = beta2 * v + (1 - beta2) * (grad ** 2)
    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)
    w = w - lr * m_hat / (v_hat ** 0.5 + eps)
    return w, m, v

w, m, v = 0.0, 0.0, 0.0
for t in range(1, 4):
    grad = 2 * (w - 4)              # dL/dw for L(w) = (w - 4)**2
    w, m, v = adam_step(w, grad, m, v, t)
    print(f"step {t}: grad = {grad:.4f}  ->  w = {w:.4f}")

# step 1: grad = -8.0000  ->  w = 0.1000
# step 2: grad = -7.8000  ->  w = 0.1999
# step 3: grad = -7.6001  ->  w = 0.2997

Lined up, the three steps of each optimizer look like this:

• Step 1 — gradient descent: w = 0.8000. Adam: w = 0.1000
• Step 2 — gradient descent: w = 1.4400. Adam: w = 0.1999
• Step 3 — gradient descent: w = 1.9520. Adam: w = 0.2997

On this particular loss — a single smooth, well-behaved bowl — plain gradient descent is actually the faster of the two: its step size scales with the size of the gradient, and here the gradient starts large, so early steps are large. Adam, by design, normalizes its step by the recent size of the gradient, which caps its early movement at roughly η per step regardless of how large the raw gradient is. That trade is deliberate: it costs Adam some speed on a friendly, single-parameter bowl like this one, but it is exactly what keeps it stable on the real loss surfaces this chapter has already described — thousands of parameters with wildly different gradient scales, ravines, and noisy mini-batch estimates — where a step size that scales directly with a possibly huge or tiny raw gradient can just as easily blow up or stall as it can help.

The same computation is one call away in PyTorch, and starting the weight at exactly zero reproduces the hand-traced numbers precisely:

import torch

model = torch.nn.Linear(1, 1, bias=False)   # a single learnable weight w
model.weight.data.fill_(0.0)                # start at w = 0, matching the hand trace
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

x = torch.tensor([[1.0]])                    # distance = 1 km
y = torch.tensor([[4.0]])                    # actual delivery time = 4 min

for step in range(1, 4):
    prediction = model(x)
    loss = ((prediction - y) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"step {step}: w = {model.weight.item():.4f}")

Learning Rate Schedules

Every update rule so far has used a single fixed η for the whole run, but a fixed learning rate is really a compromise. Early in training, parameters are far from any good solution and large steps are usually safe and desirable. Late in training, parameters are close to a minimum, and the same large step now overshoots it on every update, keeping the loss jittering around a floor instead of settling onto it. A learning rate schedule changes η over the course of training to get the best of both.

A few schedules cover most real training runs:

• Step decay multiplies η by a fixed factor (commonly 0.1) after a set number of epochs, or whenever validation loss stops improving — simple, and still a common default for training convolutional networks with SGD.
• Cosine annealing, introduced in Loshchilov and Hutter's 2017 "SGDR" paper, smoothly decays η along one arch of a cosine curve, from η_max at the start of training to η_min at the end: η_t = η_min + 0.5(η_max − η_min)(1 + cos(t·π/T)), where T is the total number of scheduled steps. Unlike step decay's sudden drops, every step's change is small and continuous.
• Warmup does the opposite of decay at the very start: η rises from a small value up to its intended peak over the first few hundred or thousand steps, before a decay schedule takes over. The original Transformer paper, "Attention Is All You Need," is a well-known example — it increased the learning rate linearly for the first 4,000 steps, then decayed it proportionally to the inverse square root of the step number. Warmup matters most for Adam specifically because v̂_t, the adaptive scaling term, is estimated from very few gradients in the earliest steps and can be unreliable; a later refinement called Rectified Adam (RAdam) makes this link explicit by tracking the variance of that estimate directly and easing off warmup once it becomes trustworthy.

The cosine schedule above is compact enough to verify directly:

import math

def cosine_lr(step, total_steps, lr_max=0.1, lr_min=0.0):
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * step / total_steps))

for step in [0, 25, 50, 75, 100]:
    print(f"step {step}: lr = {cosine_lr(step, total_steps=100):.4f}")

# step 0: lr = 0.1000
# step 25: lr = 0.0854
# step 50: lr = 0.0500
# step 75: lr = 0.0146
# step 100: lr = 0.0000

In PyTorch, this exact behaviour is built in: torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100), with one call to scheduler.step() after every optimizer.step().

Choosing an Optimizer, and Back to the Delivery App

None of this makes SGD obsolete or Adam universally correct — real teams choose between them based on what they are training. Adam, or its close relative AdamW (which fixes a subtle interaction between Adam and weight decay, and has become a standard default for training large transformer-based models), tends to need far less learning-rate tuning to get off the ground. That is exactly why it is the default first choice for transformers, recurrent networks, GANs, and anything with sparse or noisy gradients. Plain SGD with momentum, paired with a carefully tuned schedule, is still common for training large convolutional image classifiers, where practitioners have repeatedly found that a well-tuned SGD run can generalize slightly better than Adam, even if it takes longer and more babysitting to get there.

Back to that delivery-time estimate at the top of this chapter: every "34 minutes" it shows is the end product of exactly the machinery built up here — a loss function scoring the model's past mistakes, a gradient pointing away from them, a moment-based, per-parameter step size deciding how far to move on each of millions of weights at once, and a schedule that started training with big, exploratory steps and ended it fine-tuning to the nearest minute. The estimate feels instant. The optimization behind it, run one mini-batch at a time over millions of past orders, is anything but.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where optimization: adam, sgd, and learning rate schedules is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting optimization: adam, sgd, and learning rate schedules to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind optimization: adam, sgd, and learning rate schedules, 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.