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.

A worked numeric example

Make it concrete with a two-dimensional quadratic. Suppose the Hessian has eigenvalues λ_max = 10 and λ_min = 1, so the condition number is κ = 10. The stability ceiling is:

η < 2/λ_max = 2/10 = 0.2

Pick η = 0.25 and the sharp direction multiplies error by |1 - 0.25×10| = 1.5 each step — error grows 1.5× per step, so ~12 steps turns an error of 1 into over 100 and the run diverges. Pick η = 0.19, just under the wall, and the sharp direction contracts by |1 - 1.9| = 0.9 — stable, but the flat direction contracts only by |1 - 0.19×1| = 0.81, so it still converges. The theoretical best is η* = 2/(1+10) ≈ 0.182, giving a worst-case rate of (κ-1)/(κ+1) = 9/11 ≈ 0.818 per step. Notice how narrow the usable band is — a 30% overshoot from 0.19 to 0.25 flips convergence into explosion. That razor edge is why learning-rate tuning feels so unforgiving.

Implications for CPU and small language models

On small models trained on modest hardware, the learning rate matters more, not less. You often cannot afford large batches, so gradients are noisier and the safe step size is smaller — the batch-coupling logic works against you. You also cannot afford many exploratory runs, so a single divergence to NaN three hours in is expensive wasted compute. Two habits pay off. First, find the peak learning rate with a short range test — ramp η up over a few hundred steps and watch where the loss first turns up; the stability wall announces itself clearly. Second, respect that smaller and shallower models generally have less extreme curvature than giant ones, so a slightly higher relative η is often both safe and faster. The scalar you are choosing is a direct claim about the sharpest curvature your network will encounter — on a tight compute budget, guessing it well is the highest-leverage decision you make.

Common pitfalls

The classic mistake is treating the learning rate as a knob to nudge instead of a quantity bounded by geometry. If a run diverges, the fix is almost always a lower η (or gradient clipping to tame a transient λ_max spike), not more epochs. A second pitfall is porting a learning rate across a change that alters effective curvature — a different optimizer, precision, normalization, initialization, or parameterization — and expecting it to hold; it silently won’t, because you changed the landscape the number was calibrated for. A third is confusing ‘stable’ with ‘good’: a step size well below the wall will never diverge and will also waste most of your compute crawling. And a fourth is judging the learning rate by the very first steps — gradients at initialization can be huge, which is exactly the problem warmup exists to handle, but that shaping belongs to a different discussion. The fundamentals here are the fixed points every schedule and trick is ultimately negotiating with.