Learning Rate Math for LLM

The step-size role in gradient descent

Every gradient-based optimizer is a rule for turning a gradient into a parameter update. The simplest, plain gradient descent, is:

θ_{t+1} = θ_t - η · ∇L(θ_t)

The gradient ∇L is a direction and a magnitude in parameter space — it points uphill on the loss, so we subtract it to go down. But ∇L is only a local slope; it says nothing about how far that slope stays valid. The learning rate η is the scalar that converts ‘which way is downhill’ into ‘how big a step to commit.’ That is its entire job. Everything else people tune — momentum, adaptive scaling, schedules — is a way of making that step size smarter or time-varying. If you understand η as the trust you place in the local gradient extrapolating to a finite distance, the rest of learning-rate practice stops being folklore and starts being consequences.

The loss landscape and why curvature sets the limit

To see what bounds η, zoom in on the loss near a point and approximate it with a second-order Taylor expansion — a quadratic bowl:

L(θ + Δ) ≈ L(θ) + ∇L·Δ + ½ Δ^T H Δ

Here H is the Hessian, the matrix of second derivatives. It encodes curvature: how fast the slope itself changes as you move. A gentle valley has small curvature — the gradient stays roughly constant, so a big step lands where you expected. A steep, narrow ravine has large curvature — the slope flips quickly, so a step sized for the gentle direction overshoots the wall of the ravine and bounces higher than it started. The learning rate does not know which kind of terrain it is in; it applies the same η to every direction. That mismatch between a single global step size and a landscape with wildly different curvatures in different directions is the source of almost every learning-rate headache in deep networks.

The maximum stable step size: two over lambda-max

Take the cleanest case: a quadratic loss with Hessian H. Because H is symmetric, gradient descent decouples along its eigenvectors. Along an eigenvector with eigenvalue λ, the error multiplies each step by a factor:

e_{t+1} = (1 - ηλ) · e_t

The step shrinks the error only if |1 - ηλ| < 1, i.e. 0 < η < 2/λ. Every direction imposes its own ceiling, and the tightest ceiling comes from the largest eigenvalue. So the whole run is stable only when:

η < 2 / λ_max

Cross that line and the sharpest direction has |1 - ηλ_max| > 1: the error grows geometrically and the loss explodes, usually within a few steps, straight to NaN. This is the hard wall. It explains why diverging runs blow up so fast rather than drifting off gently — instability is exponential, not gradual.

The condition number: why one scalar is always a compromise

The stability limit is set by the largest eigenvalue, but the speed of convergence is set by the smallest. Along a direction with eigenvalue λ_min, the error shrinks by (1 - ηλ_min) per step — and since η is capped near 2/λ_max, that factor is close to 1 when curvatures differ a lot. The ratio that governs everything is the condition number:

κ = λ_max / λ_min

The optimal fixed step for a quadratic is η* = 2 / (λ_min + λ_max), and even at that best choice convergence slows as (κ - 1)/(κ + 1) per step. When κ is huge — and in deep networks it routinely spans many orders of magnitude — you are forced to pick a step small enough not to explode along the sharp directions, which leaves the flat directions crawling. A single scalar simply cannot be right for both. This is precisely the gap that momentum, preconditioning, and adaptive optimizers exist to close.

Adam changes what the step size means

Modern transformers are almost never trained with plain SGD; they use Adam or AdamW. This matters for the learning rate because Adam does not step by the raw gradient. It divides each coordinate’s momentum by the square root of that coordinate’s squared-gradient average:

θ ← θ - η · m̂ / (√v̂ + ε)

Because m̂/√v̂ is a normalized quantity of order 1 regardless of the gradient’s absolute size, the actual displacement per step is roughly η in each coordinate. Adam is a crude per-coordinate preconditioner: it partially cancels the curvature mismatch that wrecks SGD, which is why it tolerates the pathological conditioning of transformer landscapes. The practical consequence is that an Adam learning rate like 3e-4 is not comparable to an SGD learning rate — it sets a step measured in normalized units, not in gradient units. Copying a learning rate across optimizers is a category error, even though people do it constantly.

Learning rate and batch size are coupled

You cannot reason about η in isolation from the batch size, because a minibatch gradient is a noisy estimate of the true gradient. A larger batch averages more samples, so its gradient has lower variance — a cleaner signal you can trust for a bigger step. Two regimes are worth naming.

For SGD-style training, the linear scaling rule holds up to a point: double the batch, double the learning rate, and the expected update stays consistent. This works only below a critical batch size — the gradient-noise-scale story — beyond which the gradient is already nearly noise-free and further scaling stops helping and starts hurting. For Adam, the common heuristic is closer to square-root scaling (η ∝ √B), which follows from how the √v̂ denominator responds to reduced variance. The honest summary: batch and learning rate move together, the correct exponent depends on the optimizer and the regime, and no single rule is universal — know which regime you are in before you scale.

The edge of stability: networks train near the wall

The clean 2/λ_max picture assumes a fixed quadratic, but a real network reshapes its own loss surface as it trains. A striking empirical finding is that full-batch gradient descent does not settle comfortably below the stability limit. Instead the sharpness — the top Hessian eigenvalue λ_max — rises during training until it reaches roughly 2/η, then hovers there:

λ_max  →  ≈ 2/η   (and stays)

This is the edge of stability. The loss still decreases overall, but non-monotonically, jittering as it rides the boundary. The causation runs from the learning rate to the geometry: η effectively selects the sharpness the network is allowed to occupy — a smaller η permits sharper minima, a larger one forces the trajectory into flatter regions. That is a large part of why learning rate acts as an implicit regularizer, and why the flat, generalizing solutions people prize tend to come from training warm and high rather than tiny and timid.