Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training — part 2/7

Cleaned-text code points [15963, 30020). Bibliography omitted here; source.md includes it.

Section 2.1 demonstrates that Adam does not sufficiently adapt to the heterogeneous curvatures. On the other hand, vanilla Newton’s method has a pre-conditioner optimal for convex functions, but is vulnerable to negative curvature and rapid change of Hessian. With these insights, we design a new optimizer, Sophia, which is more adaptive to heterogeneous curvatures than Adam, more resistant to non-convexity and rapid change of Hessian than Newton’s method, and also uses a low-cost pre-conditioner.

We use θ t \theta_{t} to denote the parameter at time step t t . At each step, we sample a mini-batch from the data distribution and calculate the mini-batch loss, denoted by L t ​ ( θ t ) L_{t}(\theta_{t}) . We denote by g t g_{t} the gradient of L t ​ ( θ t ) L_{t}(\theta_{t}) , i.e. g t = ∇ L t ​ ( θ t ) g_{t}=\nabla L_{t}(\theta_{t}) . Let m t m_{t} be the EMA of gradients, m t ← β 1 ​ m t − 1 + ( 1 − β 1 ) ​ g t m_{t}\leftarrow\beta_{1}m_{t-1}+(1-\beta_{1})g_{t} , which is the numerator of the update.

EMA of diagonal Hessian estimates. Sophia uses a diagonal Hessian-based pre-conditioner, which directly adjusts the update size of different parameter dimensions according to their curvatures. We will present two options in detail in Section 2.3 for estimating the diagonal Hessian efficiently. To mitigate the overhead, we only estimate the Hessian every k k steps ( k = 10 k=10 in our implementation). At time step t t with t ​ mod ​ k = 1 t\ \mathrm{mod}\ k=1 , the estimator returns an estimate h ^ t \hat{h}_{t} of the diagonal of the Hessian of the mini-batch loss.

Similar to the gradient of the mini-batch loss function, the estimated diagonal Hessian can also have large noise. Inspired by the EMA of moments of gradients in Adam, we also denoise the diagonal Hessian estimates with EMA across iterations. We update the EMA every k k steps, resulting in the following update rule for the diagonal Hessian estimate: h t = β 2 ​ h t − k + ( 1 − β 2 ) ​ h ^ t ​ if ​ t ​ mod ​ k = 1 ; else ​ h t = h t − 1 . \displaystyle h_{t}=\beta_{2}h_{t-k}+(1-\beta_{2})\hat{h}_{t}\ \textup{ if }t\ \mathrm{mod}\ k=1;\textup{ else }h_{t}=h_{t-1}\,. (5)

Per-coordinate clipping. As discussed in Section 2.1 , on nonconvex functions, vanilla Newton’s method, which uses Hessian as the pre-conditioner, may converge to local maxima instead of local minima. In addition, the inaccuracy of Hessian estimates and the change of Hessian along the trajectory can make the second-order information unreliable. To this end, we (1) only consider the positive entries of the diagonal Hessian and (2) introduce per-coordinate clipping to the update. For a clipping threshold ρ > 0 \rho>0 , let the clipping function be clip ​ ( z , ρ ) = max ⁡ { min ⁡ { z , ρ } , − ρ } \textup{clip}(z,\rho)=\max\{\min\{z,\rho\},-\rho\} where all operations are applied coordinate-wise. The update rule is written as: θ t + 1 ← θ t − η t ⋅ clip ​ ( m t / max ⁡ { γ ⋅ h t , ϵ } , 1 ) , \displaystyle\theta_{t+1}\leftarrow\theta_{t}-\eta_{t}\cdot\textup{clip}(m_{t}/\max\{\gamma\cdot h_{t},\epsilon\},1), (6) where ϵ > 0 \epsilon>0 is a very small constant to avoid dividing by 0 0 . Note that η t ⋅ clip ​ ( m t / max ⁡ { γ ⋅ h t , ϵ } , 1 ) = ( η t / γ ) ⋅ clip ​ ( m t / max ⁡ { h t , ϵ / γ } , γ ) \eta_{t}\cdot\textup{clip}(m_{t}/\max\{\gamma\cdot h_{t},\epsilon\},1)=\left(\eta_{t}/\gamma\right)\cdot\textup{clip}(m_{t}/\max\{h_{t},\epsilon/\gamma\},\gamma) , and thus we essentially clip the original update m t / h t m_{t}/h_{t} by γ \gamma coordinate-wise, and then re-adjust the final update size by a factor of γ \gamma . The re-adjustment makes the scale of the update less dependent on γ \gamma , because now γ \gamma only controls the fraction of clipped entries but all clipped entries will eventually be set to η t \eta_{t} in the udpate—e.g., when γ \gamma is extremely small, all entries of ( η t / γ ) ⋅ clip ​ ( m t / max ⁡ { h t , ϵ / γ } , γ ) \left(\eta_{t}/\gamma\right)\cdot\textup{clip}(m_{t}/\max\{h_{t},\epsilon/\gamma\},\gamma) will be η t \eta_{t} . In practice, the choice of γ \gamma should be tuned based on the the fraction of clipped entries (See Section 3.1 for details). We present the pseudo-code of the Sophia in Algorithm 3 .

When any entry of h t h_{t} is negative, e.g., h t ​ [ i ] < 0 h_{t}[i]<0 , the corresponding entry in the pre-conditioned gradient m t ​ [ i ] / max ⁡ { γ ⋅ h t ​ [ i ] , ϵ } = m t ​ [ i ] / ϵ m_{t}[i]/\max\{\gamma\cdot h_{t}[i],\epsilon\}=m_{t}[i]/\epsilon is extremely large and has the same sign as m t ​ [ i ] m_{t}[i] , and thus η ⋅ clip ​ ( m t ​ [ i ] / max ⁡ { γ ⋅ h t ​ [ i ] , ϵ } , 1 ) = η ⋅ sign ​ ( m t ​ [ i ] ) \eta\cdot\textup{clip}(m_{t}[i]/\max\{\gamma\cdot h_{t}[i],\epsilon\},1)=\eta\cdot\textup{sign}(m_{t}[i]) , which is the same as stochastic momentum SignSGD. In other words, Sophia uses stochastic momentum SignSGD as a backup when the Hessian is negative (or mistakenly estimated to be negative or very small.) We also note that the clipping mechanism controls the worst-case size of the updates in all parameter dimensions to be at most ρ \rho , which also improves the stability (which could be a severe issue for second-order methods). Moreover, because for many parameter dimensions, the clipping is not activated and the update is automatically adjusted, our worst-case update size η ​ ρ \eta\rho can be chosen to be larger than the worst update size η \eta in stochastic momentum SignSGD.

Several previous works ( Becker & Le Cun, 1988 ; Chapelle et al., 2011 ; Schaul et al., 2013 ) , including the recent work AdaHessian ( Yao et al., 2021 ) , use diagonal Hessian as a pre-conditioner in optimizers for training neural networks. However, they use more frequent Hessian estimations, which leads to significant per-step computation overhead (more than two gradient computations), most likely because of the lack of the clipping mechanism that safeguards against inaccurate and changing Hessian. In general, to the best of our knowledge, there has not been previous reports that showed second-order optimizers achieve a speed-up on decoder-only large language models in wall-clock time or total compute (see more related work and discussions in Section 5 ).

### 2.3 Diagonal Hessian Estimators

We introduce two diagonal Hessian estimators, both of which have memory and run-time costs similar to computing a gradient (up to constant factors).

Option 1: Hutchinson’s unbiased estimator. For any loss function ℓ ⁡ ( θ ) \ell(\theta) on parameters θ ∈ ℝ d \theta\in\mathbb{R}^{d} , the Hutchinson’s estimator ( Hutchinson, 1989 ; Roosta-Khorasani & Ascher, 2015 ; Yao et al., 2021 ) first draws u ∈ ℝ d u\in\mathbb{R}^{d} from the spherical Gaussian distribution 𝒩 ⁡ ( 0 , I d ) \mathcal{N}(0,\mathrm{I}_{d}) , and then outputs h ^ = u ⊙ ( ∇ 2 ℓ ​ ( θ ) ​ u ) \hat{h}=u\odot(\nabla^{2}\ell(\theta)u) , where ⊙ \odot denotes the element-wise product, and ∇ 2 ℓ ​ ( θ ) ​ u \nabla^{2}\ell(\theta)u is the product of the Hessian with the vector u u . The Hutchinson’s estimator is an unbiased estimator for the diagonal of the Hessian, because 𝔼 ⁡ [ h ^ ] = diag ⁡ ( ∇ 2 ℓ ​ ( θ ) ) . \displaystyle\mathbb{E}[\hat{h}]=\mathrm{diag}(\nabla^{2}\ell(\theta))\,. (7) The estimator only requires a Hessian-vector product (i.e., ∇ 2 ℓ ​ ( θ ) ​ u \nabla^{2}\ell(\theta)u ), which have efficient implementations in PyTorch and JAX, instead of the full Hessian matrix.

Option 2: Gauss-Newton-Bartlett (GNB) estimator. We leverage the structure of the loss to design a biased stochastic estimator for the diagonal Hessian, following Schraudolph (2002) ; Martens (2020) ; Wei et al. (2020) . Suppose ℓ ⁡ ( θ , ( x , y ) ) \ell(\theta,(x,y)) is a loss function on an example ( x , y ) (x,y) of the form ℓ ⁡ ( θ , ( x , y ) ) = ℓ ce ​ ( f ⁡ ( θ , x ) , y ) \ell(\theta,(x,y))=\ell_{\textup{ce}}(f(\theta,x),y) where ℓ ce \ell_{\textup{ce}} is the cross-entropy loss and f ⁡ ( θ , x ) ∈ ℝ V f(\theta,x)\in\mathbb{R}^{V} is the logits, and V V is the number of items/classes in a multi-class classification problem (e.g., the vocabulary size in LLMs). First, the Hessian of ℓ ⁡ ( θ , ( x , y ) ) \ell(\theta,(x,y)) (w.r.t to variable θ \theta ) has the well-known Gauss-Newton decomposition ( Ortega & Rheinboldt, 2000 ; Schraudolph, 2002 ) (which is a simple consequence of the chain rule), ∇ θ 2 ​ ℓ ​ ( θ ) = J θ ​ f ​ ( θ , x ) ⋅ S ⋅ J θ ​ f ​ ( θ , x ) ⊤ + J θ ​ θ ​ f ​ ( θ , x ) ​ [ q ] \displaystyle\nabla_{\theta}^{2}~\ell(\theta)=J_{\theta}f(\theta,x)\cdot S\cdot J_{\theta}f(\theta,x)^{\top}+J_{\theta\theta}f(\theta,x)[q] (8) where J θ ​ f ​ ( θ , x ) J_{\theta}f(\theta,x) is the Jacobian of f f w.r.t to θ \theta viewed as a matrix in ℝ d × V \mathbb{R}^{d\times V} , S = ∂ 2 ℓ ce ​ ( t , y ) ∂ t 2 | t = f ⁡ ( θ , x ) ∈ ℝ V × V S=\frac{\partial^{2}\ell_{\textup{ce}}(t,y)}{\partial t^{2}}\Big|_{t=f(\theta,x)}\in\mathbb{R}^{V\times V} is the second-order derivatives of the loss w.r.t to the logits, q = ∂ ℓ ce ​ ( t , y ) ∂ t | t = f ⁡ ( θ , x ) ∈ ℝ V q=\frac{\partial\ell_{\textup{ce}}(t,y)}{\partial t}\Big|_{t=f(\theta,x)}\in\mathbb{R}^{V} is the first-order derivatives of the loss w.r.t to the logits, and J θ ​ θ ​ f ​ ( θ , x ) J_{\theta\theta}f(\theta,x) is the second-order derivatives of the multi-variate function f ⁡ ( θ , x ) f(\theta,x) w.r.t θ \theta , viewed as a linear map from ℝ V \mathbb{R}^{V} to ℝ d × d \mathbb{R}^{d\times d} , where d d is the dimension of the parameter θ \theta .

In the context of neural networks, past works have found that the second term J θ ​ θ ​ f ​ ( θ , x ) ​ [ q ] J_{\theta\theta}f(\theta,x)[q] in Equation 8 is often relative smaller than the first term J θ ​ f ​ ( θ , x ) ⋅ S ⋅ J θ ​ f ​ ( θ , x ) ⊤ J_{\theta}f(\theta,x)\cdot S\cdot J_{\theta}f(\theta,x)^{\top} ( Sankar et al., 2021 ) , which is often referred to as the Gauss-Newton matrix ( Dennis Jr & Schnabel, 1996 ; Ortega & Rheinboldt, 2000 ; Schraudolph, 2002 ; Chen, 2011 ) and used as pre-conditioners in second-order optimizers ( Botev et al., 2017 ; Martens, 2020 ; Gargiani et al., 2020 ) . Following this line of work, we build an unbiased estimator for the diagonal of the Gauss-Newton matrix, which is a biased estimator for the diagonal of the Hessian.

We first claim that S S only depends f ⁡ ( θ , x ) f(\theta,x) but not y y , even though the loss depends on y y . 2 2 2 Denote by p ⁡ ( θ , x ) = softmax ​ ( f ⁡ ( θ , x ) ) ∈ ℝ V p(\theta,x)=\textup{softmax}(f(\theta,x))\in\mathbb{R}^{V} the probability vector obtained by applying softmax on the logits. Indeed, a simple derivation shows that S = diagonal ​ ( p ⁡ ( θ , x ) ) − p ⁡ ( θ , x ) ​ p ​ ( θ , x ) ⊤ S=\textup{diagonal}(p(\theta,x))-p(\theta,x)p(\theta,x)^{\top} , where diagonal ​ ( p ​ ( θ , x ) ) \textup{diagonal}(p(\theta,x)) is the matrix with the vector p ⁡ ( θ , x ) p(\theta,x) residing on the diagonal. In fact, this is a general property of exponential families—the Hessian of the negative log-likelihood of any exponential family distribution only depends on the parameters of that exponential family, but not on the example on which the likelihood is evaluated. Thus, S = ∂ 2 ℓ ce ​ ( t , y ^ ) ∂ t 2 | t = f ⁡ ( θ , x ) S=\frac{\partial^{2}\ell_{\textup{ce}}(t,\hat{y})}{\partial t^{2}}\Big|_{t=f(\theta,x)} for any y ^ ∈ { 1 , … , V } \hat{y}\in\{1,\dots,V\} , which implies that S = 𝔼 y ^ ∼ p ⁡ ( θ , x ) ​ [ ∂ 2 ℓ ce ​ ( t , y ^ ) ∂ t 2 | t = f ⁡ ( θ , x ) ] . S=\mathbb{E}_{\hat{y}\sim p(\theta,x)}\left[\frac{\partial^{2}\ell_{\textup{ce}}(t,\hat{y})}{\partial t^{2}}\Big|_{t=f(\theta,x)}\right]. Because ℓ ce ​ ( t , y ) \ell_{\textup{ce}}(t,y) is the negative log-probability of the probabilistic model defined by the categorical distribution Cat ​ ( t ) \textup{Cat}(t) with parameter t t , by Bartlett’s second identity ( Bartlett, 1953 ) , we have that, S = 𝔼 y ^ ∼ Cat ​ ( t ) [ ∂ 2 ℓ ce ​ ( t , y ^ ) ∂ t 2 ] = 𝔼 y ^ ∼ Cat ​ ( t ) [ ∂ ℓ ce ​ ( t , y ^ ) ∂ t ​ ∂ ℓ ce ​ ( t , y ^ ) ∂ t ⊤ ] , \displaystyle S=\mathop{\mathbb{E}}_{\hat{y}\sim\textup{Cat}(t)}\left[\frac{\partial^{2}\ell_{\textup{ce}}(t,\hat{y})}{\partial t^{2}}\right]=\mathop{\mathbb{E}}_{\hat{y}\sim\textup{Cat}(t)}\left[\frac{\partial\ell_{\textup{ce}}(t,\hat{y})}{\partial t}\frac{\partial\ell_{\textup{ce}}(t,\hat{y})}{\partial t}^{\top}\right]\,, (9) where the first equality holds for t = f ⁡ ( θ , x ) t=f(\theta,x) and the second equality holds for all t t by Bartlett’s second identity. Therefore, the Gauss-Newton matrix satisfies

J θ ​ f ​ ( θ , x ) ⋅ S ⋅ J θ ​ f ​ ( θ , x ) ⊤ \displaystyle J_{\theta}f(\theta,x)\cdot S\cdot J_{\theta}f(\theta,x)^{\top} = 𝔼 y ^ ∼ Cat ​ ( t ) [ J θ ​ f ​ ( θ , x ) ​ ∂ ℓ ce ​ ( t , y ^ ) ∂ t ​ ∂ ℓ ce ​ ( t , y ^ ) ∂ t ⊤ ​ J θ ​ f ​ ( θ , x ) ⊤ ] \displaystyle=\mathop{\mathbb{E}}_{\hat{y}\sim\textup{Cat}(t)}\left[J_{\theta}f(\theta,x)\frac{\partial\ell_{\textup{ce}}(t,\hat{y})}{\partial t}\frac{\partial\ell_{\textup{ce}}(t,\hat{y})}{\partial t}^{\top}J_{\theta}f(\theta,x)^{\top}\right] = 𝔼 y ^ ∼ Cat ​ ( t ) [ ∇ θ ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) ​ ∇ θ ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) ⊤ ] , \displaystyle=\mathop{\mathbb{E}}_{\hat{y}\sim\textup{Cat}(t)}\left[\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x),\hat{y})\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x),\hat{y})^{\top}\right], (10) which implies that diag ​ ( J θ ​ f ​ ( θ , x ) ⋅ S ⋅ J θ ​ f ​ ( θ , x ) ⊤ ) = 𝔼 y ^ ∼ Cat ​ ( t ) [ ∇ θ ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) ⊙ ∇ θ ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) ] \textup{diag}(J_{\theta}f(\theta,x)\cdot S\cdot J_{\theta}f(\theta,x)^{\top})=\mathop{\mathbb{E}}_{\hat{y}\sim\textup{Cat}(t)}\left[\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x),\hat{y})\odot\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x),\hat{y})\right] . Hence, the quantity ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) ⊙ ∇ θ ℓ ce ​ ( f ⁡ ( θ , x ) , y ^ ) \ell_{\textup{ce}}(f(\theta,x),\hat{y})\odot\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x),\hat{y}) is an unbiased estimator of the Gauss-Newton matrix for the Hessian of a one-example loss ℓ ⁡ ( f ⁡ ( θ , x ) , y ) \ell(f(\theta,x),y) .