Adam’s second moment is a scale, not a curvature

Adam tracks v_t, an exponential average of g_t ⊙ g_t, and divides by sqrt(v_t). It is routinely described as an approximate curvature preconditioner, and that description is loose in three specific ways.

First, g ⊙ g built from gradients of the true labels is the diagonal of the empirical Fisher, which is a known-poor proxy for curvature: near a minimum the gradients shrink toward zero regardless of how sharp the loss surface is there. Second, Adam squares the mini-batch mean gradient, so its magnitude scales like 1/B relative to per-example curvature and drifts as you change batch size. Third, and most simply, sqrt(v) is a half power. A Newton step divides by curvature to the first power; dividing by its square root is a normalizer that equalizes gradient magnitudes across coordinates, not one that equalizes progress toward the minimum. Sophia’s starting point is to fix all three.

Advertisement

The update rule

Sophia keeps two states per parameter, exactly like Adam: a momentum m and a curvature estimate h. There is no memory penalty for going second-order this way.

m_t = β1 · m_(t-1) + (1 - β1) · g_t
h_t = β2 · h_(t-1) + (1 - β2) · hhat_t     [only when t mod k == 0]
h_t = h_(t-1)                                [otherwise]
θ_(t+1) = θ_t - η · clip( m_t / max(ρ · h_t, ε), 1 )

Read the last line right to left. m_t / h_t is a per-coordinate Newton step: for a quadratic L = ½ h θ^2, the step θ ← θ - g/h lands on the minimum in one move, whatever h is. ρ is a scalar that sets how aggressive that division is, max(·, ε) keeps the denominator positive, and the element-wise clip(·, 1) bounds every coordinate’s update by η. All shapes are [P], one entry per parameter.

Advertisement

Estimator one: Hutchinson

The first way to get hhat is a classic randomized trace-style estimator. Draw u with independent, zero-mean, unit-variance entries and form the element-wise product of u with the Hessian-vector product Hu:

u ~ N(0, I),   hhat = u ⊙ (H u)
E[hhat_i] = Σ_j H_ij · E[u_i u_j] = H_ii

The cross terms vanish because E[u_i u_j] = 0 for i ≠ j, leaving exactly the Hessian diagonal in expectation. Hu never requires materializing H: it is one extra backward pass through the gradient’s inner product with u, so the estimator costs roughly two gradient computations. Two problems follow. The variance of a single probe is large, and — decisively — the true Hessian of a deep network is not positive semi-definite, so individual hhat_i come back negative. Dividing by a negative curvature points the step uphill.

Estimator two: Gauss-Newton-Bartlett, and why it fits language modelling

The second estimator sidesteps the sign problem by targeting the Gauss-Newton matrix rather than the Hessian. Gauss-Newton drops the term involving second derivatives of the network output and keeps only the part built from the loss curvature composed with first derivatives — which makes it positive semi-definite by construction. For a softmax cross-entropy head it coincides with the Fisher information, and Bartlett’s identity gives it to you from gradients alone, provided the labels are drawn from the model:

yhat_b ~ softmax(f(θ, x_b))       [sampled, NOT the true label]
ghat  = ∇_θ (1/B) Σ_b loss(f(θ, x_b), yhat_b)
hhat  = B · (ghat ⊙ ghat)

The factor B is not a fudge: ghat is a mean of B independent, zero-mean per-example score vectors, so E[ghat_i^2] = (1/B) · E[g_i^2], and multiplying back by B recovers the per-example Fisher diagonal. Language modelling makes this nearly free: the model already emits a categorical distribution over the vocabulary at every position, so sampling a label is one draw from a softmax you computed anyway, and the estimator is one ordinary forward-backward.

Why the estimate is refreshed only every k steps

Neither estimator is free, so Sophia amortizes. Write c_step for the cost of a normal training step and c_est for one curvature estimate. Over a window of k steps you pay k · c_step + c_est, so the relative overhead is c_est / (k · c_step). With the Gauss-Newton-Bartlett estimator, c_est ≈ c_step (one extra forward-backward), so k = 10 costs about 10%. Hutchinson’s c_est ≈ 2 · c_step doubles that, which is why implementations also compute the estimate on a subsample of the batch — a quarter of the batch cuts the overhead fourfold.

The reason this is legitimate rather than merely cheap is that curvature moves slowly. Gradients change every step because the batch changes; the local geometry of the loss surface does not reorganize between step 4,000 and step 4,010. The β2 average over refreshes stretches the effective window further still: β2 = 0.99 sampled every tenth step averages over roughly 100 refreshes, about a thousand training steps.

Clipping is what makes a bad estimate survivable

A stale, subsampled, single-probe diagonal estimate is wrong in every coordinate and occasionally wrong in sign. Sophia does not try to make it accurate; it makes the update insensitive to the inaccuracy. Two guards do the work.

max(ρ · h, ε) floors the denominator. For Hutchinson this catches genuinely negative entries; for Gauss-Newton-Bartlett, which is PSD by construction, it only catches near-zero ones. Then the element-wise clip(·, 1) caps the ratio, so the effective step on coordinate i is η · min(1, |m_i| / (ρ h_i)). The worst case is now bounded and interpretable: when the curvature estimate is useless, the coordinate takes a step of size η in the direction of sign(m_i) — signSGD with momentum, a stable if unexciting optimizer. Note this is a per-coordinate guarantee, stronger and more local than a global gradient-norm clip, which lets one coordinate dominate the budget. The clipped fraction is also a free diagnostic.