Sophia (Second-Order Clipped Stochastic Optimization)

A scalable second-order stochastic optimization algorithm (Algorithm 1) designed for large-scale language model pre-training that preconditions an exponential moving average (EMA) of gradients m_t = beta_1 m_{t-1} + (1 - beta_1) g_t by an EMA of diagonal Hessian estimates h_t = beta_2 h_{t-1} + (1 - beta_2) hat{h}_t updated lazily once every k steps (typically k = 10), combined with coordinate-wise positive curvature thresholding max(gamma h_t, epsilon) and element-wise update clipping clip(z, rho) = max(min(z, rho), -rho) to update parameters via theta_{t+1} = theta_t - eta_t lambda theta_t - eta_t clip(m_t / max(gamma h_t, epsilon), rho). Sophia instantiates the diagonal Hessian estimator hat{h}_t via either the Gauss-Newton-Bartlett (GNB) estimator (Sophia-G, Algorithm 2: sampling mini-batch labels hat{y}_b ~ softmax(f(theta, x_b)) and computing B hat{g} odot hat{g} where hat{g} = nabla_theta (1/B) sum_{b=1}^B l_ce(f(theta, x_b), hat{y}_b)) or the Hutchinson estimator (Sophia-H, Algorithm 3: u odot nabla(langle nabla L(theta), u rangle) with u ~ N(0, I_d)).