Sophia Optimizer

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.

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.

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.

A worked two-coordinate example

Take two parameters with identical gradients but wildly different curvature, a situation transformers produce routinely — a LayerNorm gain and a rare embedding row do not live on the same scale. Set m_1 = m_2 = 1, h_1 = 100 (sharp), h_2 = 0.1 (flat), and take ρ = 0.05 as a representative value.

coord 1:  1 / max(0.05 · 100, ε) = 1/5   = 0.2   → step = 0.2η
coord 2:  1 / max(0.05 · 0.1, ε) = 1/0.005 = 200   → clipped → step = η

Adam, whose ratio |g|/sqrt(v) sits near 1 whenever the gradient is steady, would move both by about η. Sophia matches it on the flat coordinate and throttles the sharp one by 5×. That is the honest characterization of the method: it is not a Newton method that takes giant strides through flat valleys — clipping forbids that — it is signSGD that backs off exactly where a full-length step would overshoot.

Steps versus wall clock: the break-even condition

The claim that sold the method was step efficiency: on GPT-2-scale pre-training runs the original work reported reaching a target loss in roughly half the steps of a tuned AdamW baseline. Independent replications have been mixed, and the result is sensitive to how hard the baseline was tuned, so treat the headline number as setting-specific rather than universal.

The part you can reason about without trusting anyone’s benchmark is the arithmetic. At k = 10 with the cheap estimator you pay about 10% more per step, so Sophia must cut steps-to-target by more than 10% merely to tie on wall clock. There is a second, subtler drag: early in training almost every coordinate clips, meaning Sophia is signSGD-with-momentum during the phase where you are paying the overhead. The curvature term only starts earning its keep once |m_i| / (ρ h_i) falls below one for a meaningful fraction of parameters. Measure loss against seconds, never against steps.

Practical notes and pitfalls

Hyperparameters are coupled. Once most coordinates clip, ρ stops mattering and η alone sets the step; once few clip, they trade off directly. Tuning one while holding the other fixed will mislead you. Track the clipped fraction as the state variable that tells you which regime you are in.

Hutchinson needs double backpropagation. Many small-scale and CPU-oriented stacks — quantized graphs, exported or traced models, custom kernels — do not support taking a gradient of a gradient. The Gauss-Newton-Bartlett route needs only a second ordinary backward pass with sampled labels, which is why it is the practical default for language models.

On CPU the overhead is pure latency. There is no spare device to overlap the extra pass with, and small models are already memory-bandwidth bound, so raise k and subsample the estimation batch. Richer non-diagonal preconditioners exist and cost considerably more; the diagonal is the point here.