JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Beyond Adam: Preconditioning and the Geometry of an Update

Every optimizer is a choice of geometry. Learn to read SGD, momentum, AdamW, and mirror descent as different answers to one question: in which metric should a gradient step be a step?

An optimizer is a metric in disguise

You already know gradient descent: subtract a scaled gradient from the weights. But the gradient is not a direction in weight space — it is a covector, an object that eats a direction and returns a slope. To turn it into a move you must pick a metric that says how far apart two parameter settings are. Plain SGD silently assumes the Euclidean metric: every coordinate is equally expensive to change. That assumption is almost never true for a deep network, where one layer's weights might be a thousand times more sensitive than another's.

Once you see it this way, the whole zoo of advanced optimizers becomes one idea with many faces: precondition the gradient by an estimate of the local geometry, then step. The update is `w ← w − η P⁻¹ g`, where `g` is the gradient and `P` is a preconditioner that rescales (and rotates) the space. SGD takes `P = I`. Adam takes `P` diagonal from squared gradients. Second-order methods take `P` from curvature. The art is choosing a `P` that is both informative and cheap.

w \;\leftarrow\; w - \eta\, P^{-1} g

Every optimizer in this guide is this one update — the matrix P encodes the geometry in which a gradient step is taken.

AdamW: the most important one-line fix

The single most consequential refinement of the last decade is almost embarrassingly small. Classic Adam with L2 regularization adds `λw` into the gradient before the adaptive rescaling. That means the effective decay on each weight gets divided by its own gradient-magnitude estimate — large-gradient weights barely decay, small-gradient weights decay hard. The regularizer you thought you set is not the one you got. AdamW fixes this by decoupling decay: apply the adaptive step, then shrink the weights by a separate multiplicative term `w ← (1 − ηλ) w`.

# AdamW: decay is NOT mixed into the gradient
m = b1*m + (1-b1)*g
v = b2*v + (1-b2)*(g*g)
mh, vh = m/(1-b1**t), v/(1-b2**t)
w = w - lr*(mh/(vh.sqrt()+eps) + wd*w)   # wd*w sits OUTSIDE the 1/sqrt(v) rescale
Decoupled weight decay: the `wd*w` term is not divided by the second-moment estimate.

Why care so much about one term? Because weight decay is one of the few knobs that reliably improves generalization in over-parameterized networks, and coupling it to the optimizer turned a clean hyperparameter into a tangled one. AdamW is now the default for nearly every transformer. If you take one practical thing from this guide, it is: use AdamW, not Adam-plus-L2, and tune the decay independently of the learning rate.

w_t \leftarrow w_{t-1} - \eta\,\dfrac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} - \eta\,\lambda\, w_{t-1}

AdamW's one-line fix: the weight-decay term lambda-w is subtracted directly, never run through the adaptive second-moment rescaling.

Momentum done right: heavy ball vs Nesterov

Momentum accumulates a velocity so the iterate keeps rolling through small gradients and damps oscillation across a narrow valley. The classic heavy-ball update evaluates the gradient at the current point and then adds it to the velocity. Nesterov's accelerated gradient makes one surgical change: evaluate the gradient at a look-ahead point — where the velocity is about to carry you — and correct from there. The intuition is a ball that can see the upcoming slope and brakes before overshooting, rather than after.

Watch an iterate roll down the loss surface — momentum lets it carry speed through small gradients and damp oscillation across a narrow valley.

Interactive ball rolling down a loss surface toward the minimum.

On smooth convex problems this is not just a heuristic: Nesterov's method provably improves the convergence rate from `O(1/t)` to `O(1/t²)`, an accelerated rate that matches a lower bound for first-order methods. Deep networks are nonconvex, so the theorem does not apply literally — but the look-ahead correction still buys a real, measurable reduction in oscillation, which is why most frameworks expose a `nesterov=True` flag on SGD.

Mirror descent: when the geometry is not Euclidean

Sometimes the right metric is not a rescaling of Euclidean space at all. If your parameters live on the probability simplex (they must be nonnegative and sum to one), a Euclidean step can walk you straight off the set. Mirror descent generalizes gradient descent to a chosen geometry defined by a Bregman divergence — a distance built from a strictly convex potential. Pick the squared-Euclidean potential and you recover ordinary SGD; pick negative entropy and you recover the exponentiated-gradient update, which stays on the simplex by construction and rescales multiplicatively rather than additively.

w_{t+1} = \arg\min_{w}\;\langle g_t,\, w\rangle + \tfrac{1}{\eta}\, D_\psi\!\left(w,\, w_t\right)

Mirror descent: minimize the linearized loss plus a Bregman distance, so 'closeness' is measured in the non-Euclidean geometry you chose.

Mirror descent matters here for two reasons. First, it is the cleanest formal statement that the geometry is a design choice, not a law of nature — the same insight AdamW and natural gradient cash out numerically. Second, several modern tricks (multiplicative weight updates, some optimizers for attention logits, online-learning algorithms) are mirror descent in disguise, so recognizing the pattern lets you transfer guarantees between them.

A working mental model

  1. Write any update as `w ← w − η P⁻¹ g`. Ask: what is `P`, and is it telling the truth about local geometry?
  2. Decouple anything that is conceptually separate from the gradient step — weight decay above all (AdamW).
  3. If oscillation across valleys is the problem, add momentum, and prefer the Nesterov look-ahead variant.
  4. If the parameters live on a constrained set or a non-Euclidean space, ask whether mirror descent gives a more natural step.

Hold this `P⁻¹ g` template in your head. Everything that follows — Fisher metrics, Kronecker factors, sharpness penalties, layerwise scaling — is a more ambitious answer to the same question of what `P` should be. The rest of the track is just better and better `P`.