Optimization in ML: Beyond Gradient Descent

Picture a team at a Swiggy or Zomato-scale delivery platform training a model to predict order ETA. Two of the input features are distance to the customer (kilometers, roughly 0 to 15) and a live traffic-congestion index (a 0-to-1 multiplier pulled from map data). The team writes the loss function, hands it to gradient descent, and watches training crawl: thousands of epochs to shave the loss by a fraction of a percent, even though the model is nowhere near its best possible fit. Nothing is broken. The gradients are correct, the learning rate was tuned by hand, and the code has no bugs. The problem is geometric: the two features pull the loss surface into different amounts of curvature, and a single global learning rate cannot serve both directions at once. This chapter builds, from first principles and with every number checked by hand, why that happens and what momentum, RMSProp, and Adam actually do about it.

The hidden cost of feature scale: why plain gradient descent stalls

Gradient descent updates every parameter with the same rule: w ← w − η∇L(w), where η is the learning rate. For a loss that is locally well approximated by a quadratic bowl, the curvature in each direction is given by the second derivative (an eigenvalue of the Hessian matrix). A feature that swings the prediction a lot for a small change in its value produces a steep, tightly curved direction in the loss surface; a feature with a weak, diffuse effect produces a shallow, gently curved direction. When two features differ sharply in this sense, the loss surface is not a round bowl but an elongated ravine, steep walls in one direction, a nearly flat floor in the other.

The ratio of the largest to the smallest curvature is called the condition number, κ = λ_max / λ_min. For gradient descent on a quadratic, stability in any one direction requires η < 2/λ in that direction. Because η is a single global number, it must satisfy the tightest constraint, the one from λ_max. That same small η is then applied to the shallow direction too, where λ is tiny, so progress there crawls. A large condition number is exactly the geometric signature of the Swiggy ETA example: one feature (congestion index, tightly bounded, high sensitivity) creates a steep direction; another (raw distance, wide range, comparatively gentle sensitivity per unit) creates a shallow one.

Tracing gradient descent through a 10:1 ravine

To make this exact rather than hand-wavy, use a toy loss that isolates the effect: L(w1, w2) = w1² + 10·w2², standing in for a two-feature regression loss where w2 behaves like the congestion-index direction (curvature 20) and w1 behaves like the distance direction (curvature 2). The gradient is ∇L = (2w1, 20w2), condition number κ = 20/2 = 10, and the true minimum sits at (0, 0).

Start at (w1, w2) = (4, 4) with η = 0.045. Each coordinate updates independently, since the Hessian here is diagonal: w1 ← w1(1 − 2η) = 0.91·w1 and w2 ← w2(1 − 20η) = 0.10·w2.

By step 3, w2 has collapsed to essentially zero (99.9% of its distance to the minimum covered), while w1 has moved from 4 to 3.014, only 24.6% of the way. Three steps were enough to solve the steep direction and barely dented the shallow one. This is the ravine problem in numbers: the maximum stable step size, η < 2/λ_max = 2/20 = 0.1, is dictated entirely by the steep direction, and that same small step then throttles progress everywhere else. Push η past 0.1 (say η = 0.11, giving multiplier 1 − 20(0.11) = −1.2) and w2 diverges outright; the steep direction sets a hard ceiling that the shallow direction cannot escape.

The same three update rules, written once, generate every trace in this chapter:

import numpy as np

def grad(w):
    return np.array([2 * w[0], 20 * w[1]])

def gradient_descent(w0, lr, steps):
    w = np.array(w0, dtype=float)
    path = [w.copy()]
    for _ in range(steps):
        w = w - lr * grad(w)
        path.append(w.copy())
    return path

def momentum(w0, lr, beta, steps):
    w = np.array(w0, dtype=float)
    v = np.zeros_like(w)
    path = [w.copy()]
    for _ in range(steps):
        v = beta * v + lr * grad(w)
        w = w - v
        path.append(w.copy())
    return path

def rmsprop(w0, lr, beta2, steps, eps=1e-12):
    w = np.array(w0, dtype=float)
    e = np.zeros_like(w)
    path = [w.copy()]
    for _ in range(steps):
        g = grad(w)
        e = beta2 * e + (1 - beta2) * g ** 2
        w = w - lr * g / np.sqrt(e + eps)
        path.append(w.copy())
    return path

gd_path  = gradient_descent([4.0, 4.0], 0.045, 3)
mom_path = momentum([4.0, 4.0], 0.045, 0.9, 3)
rms_path = rmsprop([4.0, 4.0], 0.045, 0.9, 3)

Running gd_path reproduces the table above exactly, since w - lr*grad(w) with lr = 0.045 collapses algebraically to the 0.91 and 0.10 multipliers used by hand.

Momentum: carrying velocity through the ravine

Momentum keeps a running velocity that accumulates gradients over time instead of reacting to only the current one: v ← β·v + η·∇L(w), then w ← w − v, with β typically around 0.9. If the gradient keeps pointing the same way step after step, as it does along the shallow w1 axis, the velocity term keeps adding in that direction and grows. If the gradient keeps flipping sign, the additions partially cancel and the velocity stays small. Momentum does not change η; it changes how much of the past gradient history gets carried into the current step, and it carries more in directions of consistent sign.

Apply β = 0.9, η = 0.045 to the same start point. On the shallow w1 axis:

The two agree at step 1, since velocity starts at zero and the first momentum update is identical to plain gradient descent. From step 2 onward the accumulated velocity pulls momentum ahead: after three steps, momentum has covered (4 − 2.133)/4 = 46.7% of the distance to zero on w1, versus gradient descent's 24.6%, roughly 1.9 times the progress in the same three steps, using the same η.

But momentum is not free. Track w2, the steep axis, under the same run: 4.00 → 0.40 → −3.20 → −3.56 → −0.68. The velocity that helped w1 hurts w2: because β = 0.9 keeps a large fraction of the previous velocity even after the gradient's sign has flipped, the optimizer overshoots past zero, swings to −3.2, overshoots again to −3.56, and only starts damping back afterward. Momentum accelerates convergence in the consistent direction and can visibly worsen oscillation in the steep one if β and η are not matched to the curvature. This is a real, derivable property of the update rule, not an edge case invented for this example, and it is the reason production optimizers rarely use plain momentum alone.

Misconception: "momentum is just a bigger learning rate"

A common assumption is that momentum's speed-up on w1 could be reproduced by simply raising η in plain gradient descent. The stability bound above shows why that fails. Gradient descent is only stable while η < 2/λ_max = 0.1. Raising η toward, say, 0.09 to chase momentum's w1 speed pushes the w2 multiplier to 1 − 20(0.09) = −0.8: still technically stable, but now oscillating in sign on every step, right at the edge of the boundary where η = 0.1 causes sustained, undamped oscillation and anything beyond diverges. A single global η scales every direction by the same factor, so any increase large enough to meaningfully help the shallow axis simultaneously endangers the steep one.

Momentum achieves its 1.9x w1 speed-up while leaving η at the same safe 0.045 that keeps w2 comfortably stable (multiplier 0.10, nearly fully damped in one step when the gradient's sign stays fixed). It does this not by taking a bigger step on any single evaluation, but by letting consistent-sign gradients add up across steps while sign-flipping gradients cancel. That directional selectivity is precisely what a scalar learning rate increase cannot deliver, and it is also why momentum can still misbehave on the steep axis (as w2's overshoot showed): momentum reduces the penalty for a large η, it does not remove the underlying curvature mismatch.

RMSProp and Adam: a personal learning rate for every parameter

Momentum fixes directional consistency but leaves the per-axis curvature mismatch untouched. RMSProp attacks the mismatch directly, by dividing each parameter's step by a running estimate of that parameter's own gradient magnitude: E ← β2·E + (1 − β2)·g², then w ← w − η·g / (√E + ε), with β2 typically 0.9 to 0.999. A parameter whose gradients have consistently been large gets its step shrunk in proportion; a parameter whose gradients have been small gets a relatively larger effective step. The learning rate η is still global, but the division by √E rescales it per parameter.

Take one RMSProp step from (4, 4), β2 = 0.9, η = 0.045, E starting at 0. The gradients are g1 = 2(4) = 8 and g2 = 20(4) = 80, exactly a 10:1 ratio matching the curvature ratio. Then:

E1 = 0.1(8²) = 6.4, √E1 = 2.5298. Δw1 = 0.045(8)/2.5298 = 0.1423.
E2 = 0.1(80²) = 640, √E2 = 25.298. Δw2 = 0.045(80)/25.298 = 0.1423.

The two steps are identical, and this is not a coincidence of rounding. Whenever w1 = w2, curvature ratio 10 forces g2 = 10·g1, so g2² = 100·g1², so E2 = 100·E1 by induction on the update rule, so √E2 = 10·√E1. Substituting: g2/√E2 = 10g1/(10√E1) = g1/√E1, hence Δw2 = Δw1 exactly, at every step, for any starting point on the diagonal. RMSProp does not merely help the shallow axis; on this problem it provably erases the 10:1 imbalance and moves both coordinates by the same amount at every step, tracing a straight diagonal path from (4, 4) toward the origin: (4.000, 4.000) → (3.858, 3.858) → (3.756, 3.756) → (3.672, 3.672). Compare that to plain gradient descent's wildly unequal 24.6% (w1) versus 99.9% (w2) progress after the same three steps: RMSProp trades "fast on one axis, stalled on the other" for "fair progress on every axis," which is exactly the property a global learning rate cannot provide.

Adam combines both ideas: momentum's first-moment average of the gradient, and RMSProp's second-moment average of the squared gradient, each with its own decay rate (β1 ≈ 0.9, β2 ≈ 0.999), plus a bias correction:

m_t = β1·m_{t−1} + (1 − β1)·g_t
v_t = β2·v_{t−1} + (1 − β2)·g_t²
m̂_t = m_t / (1 − β1^t), v̂_t = v_t / (1 − β2^t)
w_t = w_{t−1} − η · m̂_t / (√v̂_t + ε)

The bias correction matters most in the first few steps. With m_0 = 0 and β1 = 0.9, a first gradient of g1 = 8 gives m1 = (1 − 0.9)(8) = 0.8, which drastically understates the true gradient. Dividing by (1 − β1¹) = 0.1 gives m̂1 = 0.8/0.1 = 8, recovering the exact gradient. Without that correction, Adam's early steps would systematically underreact, since both m and v start at zero and only "warm up" over the first several updates; because m and v use different decay rates, they warm up at different speeds, so an uncorrected Adam would not just be slow to start, it would apply a distorted ratio between numerator and denominator until both moving averages catch up.

The three trajectories through the ravine

The red path shows gradient descent's real trace from the table above: an almost vertical fall as w2 collapses, followed by a long, nearly flat crawl leftward as w1 barely moves, exactly the "steep drop, slow creep" shape a 10:1 ravine produces. The teal path shows RMSProp's first three steps sitting on the diagonal w1 = w2, the provable consequence of the equalization derived above; the dashed segment marks the direction it continues in, not additional computed points. Momentum's path is intentionally left off this diagram: its w2 overshoot (4 → 0.4 → −3.2 → −3.56) would swing far outside the plotted region and is better read from the table in the Momentum section.

Active recall

Attempt each question before reading its answer.

Q1. Why does a large condition number force gradient descent into tiny steps, even in directions where the loss is nearly flat?

Q2. For L(w1, w2) = 3w1² + 27w2², what is the largest learning rate for which plain gradient descent stays stable, and what is the condition number?

Q3. For L(w) = 5w², start w0 = 3, η = 0.05, β = 0.8. Compute w1 and w2 under plain gradient descent and under momentum (v0 = 0). What happens on the second momentum step that does not happen under plain gradient descent?

Q4. In one sentence, using the stability bound η < 2/λ_max, explain why "momentum is just a bigger learning rate" is false.

Q5. Two parameters have gradients g1 = 4 and g2 = 40 at the same step, with E1 = E2 = 0 before this update and β2 = 0.9. Show that one RMSProp step produces the same Δw for both parameters, for any η.

Q6. With β1 = 0.9 and β2 = 0.999, compare the bias-correction factors (1 − β1^t) and (1 − β2^t) at t = 1. Why does this difference matter for Adam's very first update?

Answers

A1. Stability in any direction requires η < 2/λ for that direction's curvature λ. Because η is one shared number, it must satisfy the tightest constraint, from λ_max. That same small η then applies to shallow directions too, where λ is small, so each shallow-direction step (proportional to η·λ_shallow) is tiny even though the direction itself is not close to converged.

A2. Second derivatives: ∂²L/∂w1² = 6, ∂²L/∂w2² = 54. λ_max = 54, so η_max = 2/54 ≈ 0.0370. Condition number κ = 54/6 = 9.

A3. g(w) = 10w. Plain GD: g0 = 30, w1 = 3 − 0.05(30) = 1.5; g1 = 15, w2 = 1.5 − 0.05(15) = 0.75. Momentum: step 1 matches GD exactly (v0 = 0), so w1 = 1.5. Step 2: g1 = 10(1.5) = 15, v2 = 0.8(1.5) + 0.05(15) = 1.2 + 0.75 = 1.95, w2 = 1.5 − 1.95 = −0.45. Plain GD lands at 0.75 (still approaching zero, no overshoot); momentum overshoots past zero to −0.45. The accumulated velocity from step 1 carries into step 2 and pushes the parameter past the minimum, the same overshoot mechanism seen on the w2 axis in the worked example.

A4. Raising η enough to meaningfully speed up a shallow direction raises it by the same factor everywhere, so it also pushes the steep direction's multiplier |1 − ηλ_max| toward or past the η_max = 2/λ_max boundary, while momentum accelerates only directions with consistent gradient sign and leaves η itself, and therefore the steep direction's stability margin, unchanged.

A5. E1 = 0.9(0) + 0.1(4²) = 1.6, √E1 = 1.2649. E2 = 0.9(0) + 0.1(40²) = 160, √E2 = 12.649 = 10·√E1 (since g2 = 10g1 implies g2² = 100g1², so E2 = 100E1 for any starting E1). Δw1 = η(4)/1.2649 = 3.1623η. Δw2 = η(40)/12.649 = 3.1623η. Equal for any η, because the 10x ratio in the gradients cancels exactly against the 10x ratio in √E, the same algebraic mechanism proved in the RMSProp section, here shown to hold for any gradient pair in a fixed ratio, not only the specific (8, 80) pair used there.

A6. (1 − 0.9¹) = 0.1; (1 − 0.999¹) = 0.001, a 100x difference. v's exponential moving average decays far more slowly (β2 close to 1), so it stays far closer to its zero initialization after one step than m does. Uncorrected, √v is far too small relative to m, so the ratio m/√v used in the update would be inflated well beyond its true value on the very first steps, before v has accumulated enough history; bias correction rescales each moving average by its own warm-up factor so the ratio reflects the true gradient statistics from step 1 onward.

Think About It

Think about this: How would you explain optimization in ml: beyond gradient descent 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.

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 in ml: beyond gradient descent 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 in ml: beyond gradient descent 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 in ml: beyond gradient descent, 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.