024 — Learning Rate and the Shape of the Loss Surface

Where We Are

In Chapter 023 we finally gave the model a rule for improving itself:

\theta_{t+1}=\theta_t-\eta\nabla L(\theta_t)

The gradient tells us which direction lowers the loss.

But the learning rate $\eta$ decides how far we trust that direction.

That sounds like a small implementation detail. It is not. It is the difference between learning, wasting compute, and divergence.

1. The Problem: Downhill Can Still Go Wrong

Imagine standing on a valley floor in fog.

You can feel the local slope beneath your feet, so you know which way is downhill. But you do not know how far the hill keeps bending the same way.

Three choices:

tiny step → safe, but painfully slow,
sensible step → fast progress,
giant step → jump across the valley and climb the other wall.

Gradient descent has exactly the same problem.

Suppose

L(w)=(w-3)^2

Then

\frac{dL}{dw}=2(w-3)

Start at $w_0=0$ .

Case A — $\eta=0.1$

w_1=0-0.1(-6)=0.6

We moved toward 3.

Case B — $\eta=1$

w_1=0-1(-6)=6

We jumped past the answer to the opposite side.

One more step:

w_2=6-1(6)=0

Now we bounce forever: $0\rightarrow6\rightarrow0\rightarrow6$ .

Case C — $\eta=1.2$

w_1=7.2

The next step overshoots even farther. The oscillation grows. Training diverges.

The gradient was correct in all three cases. The difference was the step size.

2. Why Curvature Matters

A steeply curved loss changes direction quickly.

A gently curved loss changes direction slowly.

That means a learning rate that is safe on a flat surface may be too aggressive on a sharply curved one.

For

L(w)=\frac{a}{2}(w-w^*)^2

we have

\frac{dL}{dw}=a(w-w^*)

Gradient descent becomes

w_{t+1}-w^*=(1-\eta a)(w_t-w^*)

The error shrinks only when

|1-\eta a|<1

which gives

\boxed{0<\eta<\frac{2}{a}}

The number $a$ is curvature.

Higher curvature → smaller safe learning rate.

This is our first deep connection between geometry and optimization.

3. One Parameter Is Easy. Neural Networks Are Not.

A real model does not have one parameter.

It may have millions or billions.

Then the loss is not a curve. It is a high-dimensional surface.

Near a point, the curvature is captured by the Hessian:

H=\begin{bmatrix} \frac{\partial^2L}{\partial\theta_1^2} & \cdots \\ \vdots & \ddots \end{bmatrix}

Its eigenvalues tell us how sharply the surface curves along different directions.

For a quadratic loss, the safe constant learning rate must satisfy approximately

\boxed{0<\eta<\frac{2}{\lambda_{\max}}}

where $\lambda_{\max}$ is the largest eigenvalue of the Hessian.

So the steepest direction controls the maximum safe step for every direction.

4. Why Narrow Valleys Make Learning Zig-Zag

Consider a bowl that is wide in one direction and narrow in another.

        steep wall
          /\
         /  \
--------/----\-------- flat direction
       /      \

Along the narrow direction, curvature is large. Along the long valley, curvature is small.

A single learning rate must be small enough not to explode across the narrow direction.

But that same small step now moves slowly along the flat direction.

This creates the classic zig-zag path:

\     /\     /\
 \   /  \   /  \
  \ /    \ /    → minimum

The model wastes many steps bouncing across the valley instead of moving directly along it.

This is why feature scaling, normalization, momentum, RMSProp and Adam matter later: they are all, in different ways, attempts to make optimization less hostage to badly shaped geometry.

5. Tiny Numerical Experiment

Let

L(x,y)=50x^2+y^2

Then

\nabla L=[100x,2y]

The curvature in $x$ is 50 times larger than in $y$ .

Start at $(1,1)$ .

With $\eta=0.01$ :

x_1=1-0.01(100)=0

y_1=1-0.01(2)=0.98

One direction is solved almost instantly. The other barely moves.

That mismatch is the optimization problem.

6. Code It From Scratch

import numpy as np

def loss(x, y):
    return 50*x**2 + y**2

def grad(x, y):
    return np.array([100*x, 2*y])

def run(lr, steps=20):
    p = np.array([1.0, 1.0])
    history = []
    for _ in range(steps):
        history.append((p.copy(), loss(*p)))
        p = p - lr * grad(*p)
    return history

for lr in [0.001, 0.01, 0.021]:
    h = run(lr)
    print(lr, h[-1])

Before running it, predict:

which learning rate is painfully slow?
which one is useful?
which one becomes unstable?

That prediction matters more than the code.

7. What the Loss Curve Tells You

A training-loss graph is not just a scorecard. It is a diagnostic tool.

Smooth downward curve

Likely healthy learning rate.

Almost flat

Learning rate may be too small, gradients may be tiny, or optimization may be stuck in a flat region.

Wild oscillation

Learning rate may be too large.

Loss suddenly becomes `nan`

Possible exploding activations/gradients, numerical overflow, or an excessively large update.

Falls fast then plateaus

Could be normal convergence, poor conditioning, saturation, insufficient model capacity, or a learning-rate schedule that needs adjustment.

Optimization debugging begins by learning to read the shape of training.

8. Why We Change the Learning Rate During Training

Early in training, we are far from the solution. Large steps are useful.

Near the minimum, large steps cause bouncing. Smaller steps help settle.

So modern training often uses a learning-rate schedule.

A simple decay might be

\eta_t=\frac{\eta_0}{1+kt}

Other common ideas include:

step decay,
cosine decay,
exponential decay,
warmup,
one-cycle schedules.

Do not memorize them yet.

Remember the principle:

far away → exploration can tolerate larger steps; near the bottom → precision needs smaller steps.

9. Warmup: Why Start Smaller?

Transformers are often trained with learning-rate warmup.

At the beginning, parameters and optimizer statistics are poorly calibrated. Large updates can destabilize the network before useful representations form.

Warmup gradually increases the learning rate:

0\rightarrow\eta_{target}

and then a decay schedule takes over.

This may look like a hack until you see the geometry: the model begins in a region whose scale and curvature are not yet well behaved.

10. PyTorch View

import torch

w = torch.tensor([0.0], requires_grad=True)
optimizer = torch.optim.SGD([w], lr=0.1)

for step in range(20):
    optimizer.zero_grad()
    loss = (w - 3).pow(2).mean()
    loss.backward()
    optimizer.step()
    print(step, w.item(), loss.item())

Change lr to 0.01, 0.5, 1.0, and 1.2.

Do not merely observe whether it works. Explain the behaviour from the update equation.

11. Shape Check

The parameter vector has shape

[P]

The gradient must have the same shape:

[P]

The update

parameter - scalar_learning_rate * gradient

therefore returns

[P]

A scalar learning rate changes magnitude, not dimensionality.

Later optimizers will effectively create different step sizes for different parameters, but the shape rule remains the same.

12. Break It

Try these experiments:

Make the curvature in the first direction 5000 instead of 50.
Keep the same learning rate.
Observe divergence.
Reduce the learning rate until training becomes stable.
Rescale the variable so the two directions have similar curvature.
Compare the number of steps needed.

You are experimentally discovering conditioning.

13. Research Lens

The learning-rate problem is not a beginner-only issue.

Large-scale training still spends enormous effort on:

optimizer choice,
learning-rate schedules,
warmup length,
gradient clipping,
normalization,
curvature-aware updates,
stability at extreme batch sizes.

The same question survives from a two-parameter parabola to frontier models:

How large a trustworthy update can we take from the information we have right now?

14. What You Should Now Understand

You should be able to explain:

why the gradient alone is not enough,
why curvature determines safe step size,
why too-small learning rates waste time,
why too-large learning rates oscillate or diverge,
why ill-conditioned valleys cause zig-zagging,
why schedules and warmup exist,
why optimization is fundamentally a geometric problem.

Bridge to Chapter 025

We can now move parameters through a loss surface.

But neural networks do not contain one scalar weight. They contain vectors and matrices of weights.

Y=XW+b

what exactly does

\frac{\partial L}{\partial W}

mean?

And how can we derive it without differentiating every matrix entry one by one?

Next → Matrix Calculus for Neural Networks.

Level 1 — Explain It

In one paragraph, explain why a correct gradient can still produce bad training.
Explain the difference between direction and step size.
Why does a sharply curved loss require a smaller safe learning rate?
What does an oscillating loss curve suggest?

Level 2 — Calculate It

Let

L(w)=(w-4)^2

\frac{dL}{dw}=2(w-4)

Start at $w_0=0$ .

Compute $w_1,w_2,w_3$ for $\eta=0.1$ .
Repeat for $\eta=0.5$ .
Repeat for $\eta=1.0$ .
For each case, calculate the loss after every step.
Which case converges fastest? Which oscillates?

Level 3 — Predict It

Consider

L(x,y)=100x^2+y^2

with gradient

\nabla L=[200x,2y].

Before running code, predict what happens from $(1,1)$ for:

$\eta=0.001$
$\eta=0.005$
$\eta=0.011$

Explain your prediction using the largest curvature direction.

Level 4 — Break It

Write gradient descent for

def loss(x, y):
    return 100*x*x + y*y

Then deliberately:

choose an unstable learning rate,
log the parameters and loss,
identify the first point at which the loss grows,
reduce the learning rate until it converges,
describe the zig-zag path in words.

Level 5 — Build It

Build a small experiment that compares at least five learning rates on the same quadratic loss.

Your program should produce:

final loss,
number of steps to reach a chosen tolerance,
whether training diverged,
a plot of loss vs. step,
a 2D trajectory over loss contours.

Do not use an optimizer library for the update.

Level 6 — Investigate It

Experiment: conditioning

Compare these two losses:

L_1(x,y)=x^2+y^2

and

L_2(x,y)=1000x^2+y^2.

Use the same starting point.

Investigate:

What is the largest stable learning rate for each?
How many steps are needed to reduce the loss by a factor of $10^4$ ?
What does each optimization trajectory look like?
How does scaling one coordinate change the result?
Connect your observation to the condition number.

Write a short research-style conclusion:

Hypothesis → Experiment → Observation → Explanation → Limitation

Challenge

For

L(w)=\frac{a}{2}(w-w^*)^2

derive from first principles why convergence requires

0<\eta<\frac{2}{a}.

Then explain why the multidimensional analogue depends on the largest eigenvalue of the Hessian.