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

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

The scaling law is in favor of Sophia over AdamW. In Figure 1 (d), we plot the validation loss on OpenWebText of models of different sizes pre-trained for 100K steps. The gap between Sophia and AdamW grows as we scale up the models. Moreover, the 540M model trained by Sophia-H has smaller loss than the 770M model trained by AdamW. The 355M model trained by Sophia-H has comparable loss as the 540M model trained by AdamW.

Few-shot Evaluation on Downstream Tasks (SuperGLUE). As shown in Figure 6 , as expected, the improvement in validation loss transfers to an improvement in downstream task accuracy. With the same number of steps in pre-training, GPT-2 medium, GPT-2 large and GPT NeoX 1.5B pre-trained with Sophia have better few-shot accuracy on most subtasks. Also, models pre-trained with Sophia have comparable few-shot accuracy as models pre-trained with AdamW for 2x number of steps.

### 3.4 Analysis

Comparison of wall-clock time and amount of compute. We compare the total compute (TFLOPs) per step and the wall-clock time on A100 GPUs in Table 1 . We report the average time per step (T(step)), the time spent in Hessian computation (T(Hessian)) and the total compute following Chowdhery et al. (2022) . Since we calculate the diagonal Hessian estimate with a reduced batch size every 10 steps, the computation of the Hessian accounts for 6 % \% of the total compute, and the overall wall-clock time overhead is less than 5 % \% compared with AdamW. In terms of memory usage, our optimizer has two states, m m and h h , which results in the same memory cost as AdamW.

Sensitivity to ρ \rho and β 2 \beta_{2} , and transferability of hyperparameters. On a 30M model, we perform a grid search to test the sensitivity of Sophia-H to hyperparamters (Figure 7 (c)). All combinations have a similar performance, while β 2 = 0.99 \beta_{2}=0.99 and γ = 0.01 \gamma=0.01 performs the best. Moreover, this hyperparameter choice is transferable across model sizes. For all the experiments on 125M, 355M and 770M, we use the hyperparameters searched on the 30M model, which is γ = 0.01 \gamma=0.01 , β 2 = 0.99 \beta_{2}=0.99 .

Training Stability. Sophia-H has better stability in pre-training compared to AdamW and Lion. Gradient clipping (by norm) is an important technique in language model pre-training as it avoids messing up the moment of gradients with one mini-batch gradient computed from rare data ( Zhang et al., 2020 ) . In practice, the frequency that gradients clipping is triggered is related to the training stability—if the gradient is frequently clipped, the iterate can be at a very instable state. We compare the proportion of steps where gradient clipping is triggered on GPT-2 small (125M) in Figure 7 (a). Although all methods use the same clipping threshold 1.0, Sophia-H seldomly triggers gradient clipping, while AdamW and Lion trigger gradient clipping in more than 10 % \% of the steps.

Another common trick of pre-training deep Transformers is scaling the product of keys and values by the inverse of the layer index as implemented by Mistral ( Karamcheti et al., 2021 ) and Huggingface ( Wolf et al., 2020 ) . This stabilizes training and increases the largest possible learning rate. Without this trick, the maximum learning rate of AdamW and Lion on GPT-2 medium (355M) can only be 1.5e-4, which is much smaller than 3e-4 with the trick (the loss will blow up with 3e-4 without the trick). Moreover, the loss decreases much slower without the trick as shown in Figure 7 (b). In all the experiments, Sophia-H does not require scaling the product of keys and values by the inverse of the layer index.

### 3.5 Ablation Study

Choices of Hessian update frequency k k . We study the effect of Hessian update frequency k k of Sophia-G on computational overhead and validation loss on a 30M GPT-2 model. We consider k = 1 , 10 , 100 k=1,10,100 and run each method for 100k, 200k, and 400k steps. All other hyperparameters are fixed, and we tune the peak learning rate with a grid search. We plot the amount of compute spent and the validation loss of each run in Figure 8 (a). While k = 1 k=1 has better validation loss with the same number of steps, the computational overhead is 50 % \% and the convergence speed with respect to amount of compute is worse than k = 10 k=10 . The choice of k = 100 k=100 still outperforms AdamW, but is not as good as k = 10 k=10 .

Diagonal Hessian pre-conditioners. We compare different diagonal Hessian pre-conditioners (with the same k = 10 k=10 and γ \gamma found by grid search): Empirical Fisher (E-F+clip), AdaHessian (AH+clip), Hutchinson (Sophia-H), and GNB (Sophia-G). Note that empirical Fisher is the EMA of squared gradients, which differs from GNB in label sampling. We run each method for 100k, 200k, and 400k steps and plot the results in Figure 8 (b). Results indicate that GNB is better than Empirical Fisher, which is consistent with Kunstner et al. (2019) . Sophia-H is also consistently better than AdaHessian. We hypothesize that the difference stems from that the EMA of the diagonal Hessian estimates (used in Sophia-H ) has more denoising effect than the EMA of the second moment of Hessian estimates (used in AdaHessian).

Element-wise Clipping. We compare the role of different update clipping strategy in Figure 8 (c). We include element-wise clipping without pre-conditioners (Clip), update normalization without pre-conditioners (Normalize), AdaHessian and Sophia-G without clipping (GNB). The learning rate is found by grid search. Note that clipping without pre-conditioner is essentially the same as sign momentum, or Lion with a single β \beta . Without element-wise clipping, we find that AdaHessian will diverge with k = 2 k=2 and GNB will diverge with k = 5 k=5 , thus we use k = 1 k=1 for AdaHessian and k = 2 k=2 for GNB. Results indicate that per-coordinate clipping itself is already better than AdamW. Further adding the GNB pre-conditioner makes Sophia-G much better than baselines.

## 4 Theoretical Analysis

This section provides runtime bounds for the deterministic version of Sophia that does not depend on the local condition number (the ratio between maximum and minimum curvature at the local minimum) and the worst-case curvature (that is, the smoothness parameter), demonstrating the advantage of Sophia in adapting to heterogeneous curvatures across parameter dimensions.

We start with standard assumptions on the differentiability and uniqueness of the minimizer.

###### Assumption 4.1 .

L : ℝ d → ℝ L:\mathbb{R}^{d}\to\mathbb{R} is a twice continuously differentiable, strictly convex function with θ ∗ \theta^{*} being its minimizer. For convenience, we denote λ min ​ ( ∇ 2 L ​ ( θ ∗ ) ) \lambda_{\min}(\nabla^{2}L(\theta^{*})) by μ \mu .

The following assumptions state that the Hessian has a certain form of continuity—within a neighborhood of size R R , the ratio between the Hessians, ∇ 2 L ​ ( θ ′ ) − 1 ​ ∇ 2 L ​ ( θ ) \nabla^{2}L(\theta^{\prime})^{-1}\nabla^{2}L(\theta) , is assumed to be bounded by a constant 2.

###### Assumption 4.2 .

There exists a constant R > 0 R>0 , such that ∀ θ , θ ′ ∈ ℝ d , ‖ θ − θ ′ ‖ 2 ≤ R ⟹ ‖ ∇ 2 L ​ ( θ ′ ) − 1 ​ ∇ 2 L ​ ( θ ) ‖ 2 ≤ 2 \displaystyle\forall\theta,\theta^{\prime}\in\mathbb{R}^{d},\left\|\theta-\theta^{\prime}\right\|_{2}\leq R\implies\left\|\nabla^{2}L(\theta^{\prime})^{-1}\nabla^{2}L(\theta)\right\|_{2}\leq 2 (15)

We analyze the convergence rate of the deterministic version of the Sophia on convex functions, θ t + 1 = θ t − η V t ⊤ clip ( V t ( ∇ 2 L ( θ t ) ) − 1 ∇ L ( θ t ) , ρ ) , \displaystyle\theta_{t+1}=\theta_{t}-\eta V_{t}^{\top}\textup{clip}(V_{t}(\nabla^{2}L(\theta_{t}))^{-1}\nabla L(\theta_{t}),\rho), (16) where ∇ 2 L ​ ( θ t ) = V t ⊤ ​ Σ t ​ V t \nabla^{2}L(\theta_{t})=V_{t}^{\top}\Sigma_{t}V_{t} is an eigendecomposition of ∇ 2 L ​ ( θ t ) \nabla^{2}L(\theta_{t}) . Here, we use the full Hessian as the pre-conditioner because the diagonal Hessian pre-conditioner cannot always work for general functions which may not have any alignment with the natural coordinate system. Moreover, the matrix V t V_{t} transforms ( ∇ 2 L ( θ t ) ) − 1 ∇ L ( θ t ) (\nabla^{2}L(\theta_{t}))^{-1}\nabla L(\theta_{t}) into eigenspace and thus the clipping can be done element-wise in the eigenspace. We do not need the max between Hessian and ϵ \epsilon in the original version of Sophia because the Hessian is always PSD for convex functions. Finally, the matrix V t ⊤ V_{t}^{\top} transforms the update back to the original coordinate system for the parameter update.

###### Theorem 4.3 .

Under Assumption 4.1 and Assumption 4.2 , let η = 1 / 2 , ρ = R 2 ​ d \eta=1/2,\rho=\frac{R}{2\sqrt{d}} , the update in Equation 16 reaches a loss at most ϵ \epsilon in T ≲ d ⋅ L ⁡ ( θ 0 ) − min ⁡ L μ ​ R 2 + ln ⁡ μ ​ R 2 32 ​ d ​ ϵ T\lesssim d\cdot\frac{L(\theta_{0})-\min L}{\mu R^{2}}+\ln\frac{\mu R^{2}}{32d\epsilon} steps.

The first term in the runtime bound is a burn-in time before reaching a local region, where the error decays exponentially fast so that the runtime bound is logarithmic in 1 / ϵ 1/\epsilon as the second term in the runtime bound shows. We remark that the bound does not depend on the condition number (the ratio between the maximum and minimum eigenvalue of Hessian), as opposed to the typical dependency on the maximum eigenvalue of the Hessian (or the smoothness parameter) in standard analysis of gradient descent in convex optimization ( Boyd & Vandenberghe, 2004 ) . Moreover, even on simple quadratic functions, the convergence rate of simplified Adam (SignGD) depends on the condition number (Appendix D.1 ). This demonstrates the advantage of Sophia in adapting to heterogeneous curvatures across parameter dimensions.

## 5 Related work

Stochastic Adaptive First-order Optimizers in Deep Learning. The idea of adaptive first-order optimizers dates back to RProp ( Braun & Riedmiller, 1992 ) . AdaGrad ( Duchi et al., 2011 ) adapted the learning rate of features by estimated geometry and assign larger learning rate to infrequent features. RMSProp ( Hinton et al., 2012 ) generalized RProp and is capable to work with smaller batch sizes. Adam ( Kingma & Ba, 2014 ) improved RMSProp by introducing a running average of gradients, and has so far become the dominant approach to solve optimization problems in deep learning, especially for training Transformers ( Vaswani et al., 2017 ) . Many follow-up works proposed variants of Adam ( Dozat, 2016 ; Shazeer & Stern, 2018 ; Reddi et al., 2019 ; Loshchilov & Hutter, 2017 ; Zhuang et al., 2020 ; You et al., 2019 ) . Chen et al. (2023) performed a search over adaptive first-order algorithms and discovered Lion, which is a improved version of sign momentum SGD.

Second-order Optimizers in Deep Learning. Second-order optimizers are believed to have the potential to outperform adaptive first-order optimizers. Classical second-order optimization algorithms pre-condition the gradient with curvature information ( Broyden, 1970 ; Nesterov & Polyak, 2006 ; Conn et al., 2000 ) . Over the years, people have developed numerous ways to adapt these methods to deep learning. To the best of our knowledge, Becker & Le Cun (1988) was the first to use diagonal Hessian as the pre-conditioner. Martens et al. (2010) approximated the Hessian with conjugate gradient. Schaul et al. (2013) automatically tuned learning rate of SGD by considering diagonal Hessian. Pascanu & Bengio (2013) considered Gaussian Newton’s approximation of Hessian and Fisher information matrix. Martens & Grosse (2015) and follow-up works ( Ba et al., 2017 ; George et al., 2018 ; Martens et al., 2018 ; Zhang et al., 2022a ) proposed to approximate the Hessian based on the structure of neural networks. Yao et al. (2021) ; Jahani et al. (2021) proposed to use the EMA of diagonal Hessian estimator as the pre-conditioner.

Despite these progress on deep learning applications, for decoder-only large language models, Adam still appears to the most popular optimizer. The authors of this paper suspect that many previous second-order optimizers face the challenge that the computational / memory overhead due to frequent Hessian computation hinders improvements in wall-clock time ( Martens & Grosse, 2015 ; Gupta et al., 2018 ) . Some of them also depend on specific model architecture or hardware structures, e.g., Anil et al. (2020) offloads hessian computation to CPUs, and George et al. (2018) needs ResNets and very large batch size to approximate the Fisher information matrix. To the best of our knowledge, there was no previous report that second-order optimizers can achieve a speed-up on large language models in total compute.

Gradient Clipping. Global gradient clipping has been a standard practice in pre-training language models ( Merity et al., 2017 ; Radford et al., 2019 ; Izsak et al., 2021 ; Zhang et al., 2022b ) . It helps stabilizes training and avoids the effect of rare examples and large gradient noise. Zhang et al. (2019) ; Mai & Johansson (2021) showed that global gradient clipping is faster than standard SGD when global smoothness does not hold. Zhang et al. (2020) ; Crawshaw et al. (2022) found out per-coordinate gradient clipping can function as adaptivity. In addition to gradient clipping, Sophia is the first to clip the update (coordinate-wise) in second-order methods to avoid the effect of Hessian’s changing along the trajectory and the inaccuracy of Hessian approximation.

Optimization Algorithms in LM Pre-training. Adam ( Kingma & Ba, 2014 ) (with decoupled weight decay ( Loshchilov & Hutter, 2017 ) ) has become the dominant approach for language model pre-training ( Vaswani et al., 2017 ; Devlin et al., 2018 ; Radford et al., 2019 ; Brown et al., 2020 ; Zhang et al., 2022b ; Touvron et al., 2023 ) . Different from vision tasks with CNNs ( He et al., 2016 ) where models trained with SGD generalize better than models trained with Adam, Adam outperforms SGD by a huge margin on language modeling tasks with Transformers ( Anil et al., 2019 ; Liu et al., 2020 ; Kunstner et al., 2023 ) . Raffel et al. (2020) ; Chowdhery et al. (2022) trained Transformers with AdaFactor ( Shazeer & Stern, 2018 ) , which is a low rank version of Adam. You et al. (2019) proposed to make the update of Adam proportional to per-layer paramter norm to stably train LLMs.

## 6 Conclusion

We introduced Sophia, a scalable second-order optimizer for language model pre-training. Sophia converges in fewer steps than first-order adaptive methods, while maintaining almost the same per-step cost. On language modeling with GPT models, Sophia achieves a 2x speed-up compared with AdamW in the number of steps, total compute, and wall-clock time.

### Acknowledgements

We thank Jeff Z. HaoChen, Neil Band, Garrett Thomas for valuable feedbacks. HL is supported by Stanford Graduate Fellowship. The authors would like to thank the support from NSF IIS 2211780.

## Appendix A Additional Experiment Results

Dynamics of Sophia in training. We measure the ℓ 2 \ell_{2} norm of the EMA of the diagonal Hessian h t h_{t} , and the proportion of parameters where clipping happens (that is, | m t ​ [ i ] / max ⁡ { γ ​ h t ​ [ i ] , ϵ } | |m_{t}[i]/\max\{\gamma h_{t}[i],\epsilon\}| is larger than 1) during pre-training in Figure 9 . After the initial stage, the norm of the Hessian steadily grows. The proportion of parameters where clipping happens approaches 60 % \% , which corroborates the importance of per-coordinate clipping in the algorithm. We also note that the clipping proportion generally should be between 50% and 90% for Sophia to be effective, and the users are recommended to tune γ \gamma to achieve some a clipping proportion (See Section 3.1 for details).