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

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

# Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training

###### Abstract

Given the massive cost of language model pre-training, a non-trivial improvement of the optimization algorithm would lead to a material reduction on the time and cost of training. Adam and its variants have been state-of-the-art for years, and more sophisticated second-order (Hessian-based) optimizers often incur too much per-step overhead. In this paper, we propose Sophia, S econd- o rder Cli p ped Stoc h astic Opt i miz a tion, a simple scalable second-order optimizer that uses a light-weight estimate of the diagonal Hessian as the pre-conditioner. The update is the moving average of the gradients divided by the moving average of the estimated Hessian, followed by element-wise clipping. The clipping controls the worst-case update size and tames the negative impact of non-convexity and rapid change of Hessian along the trajectory. Sophia only estimates the diagonal Hessian every handful of iterations, which has negligible average per-step time and memory overhead. On language modeling with GPT models of sizes ranging from 125M to 1.5B, Sophia achieves a 2x speed-up compared to Adam in the number of steps, total compute, and wall-clock time, achieving the same perplexity with 50 % \% fewer steps, less total compute, and reduced wall-clock time. Theoretically, we show that Sophia, in a much simplified setting, adapts to the heterogeneous curvatures in different parameter dimensions, and thus has a run-time bound that does not depend on the condition number of the loss.

## 1 Introduction

Language models (LLMs) have gained phenomenal capabilities as their scale grows ( Radford et al., 2019 ; Kaplan et al., 2020 ; Brown et al., 2020 ; Zhang et al., 2022b ; Touvron et al., 2023 ; OpenAI, 2023 ) . However, pre-training LLMs is incredibly time-consuming due to the massive datasets and model sizes—hundreds of thousands of updates to the model parameters are required. For example, PaLM was trained for two months on 6144 TPUs, which costed 10 million dollars ( Chowdhery et al., 2022 ) .

Pre-training efficiency is thus a major bottleneck in scaling up LLMs. This work aims to improve pre-training efficiency with a faster optimizer, which either reduces the time and cost to achieve the same pre-training loss, or alternatively achieves better pre-training loss with the same budget.

Adam ( Kingma & Ba, 2014 ) (or its variants ( Loshchilov & Hutter, 2017 ; Shazeer & Stern, 2018 ; You et al., 2019 ) ) is the dominantly used optimizer for training LLMs, such as GPT ( Radford et al., 2019 ; Brown et al., 2020 ) , OPT ( Zhang et al., 2022b ) , Gopher ( Rae et al., 2021 ) and LLAMA ( Touvron et al., 2023 ) . Designing faster optimizers for LLMs is challenging. First, the benefit of the first-order (gradient-based) pre-conditioner in Adam is not yet well understood ( Liu et al., 2020 ; Zhang et al., 2020 ; Kunstner et al., 2023 ) . Second, the choice of pre-conditioners is constrained because we can only afford light-weight options whose overhead can be offset by the speed-up in the number of iterations. For example, the block-diagonal Hessian pre-conditioner in K-FAC is expensive for LLMs ( Martens & Grosse, 2015 ; Grosse & Martens, 2016 ; Ba et al., 2017 ; Martens et al., 2018 ) . On the other hand, Chen et al. (2023) automatically search among the light-weight gradient-based pre-conditioners and identify Lion, which is substantially faster than Adam on vision Transformers and diffusion models but only achieves limited speed-up on LLMs ( Chen et al., 2023 ) .

This paper introduces Sophia, S econd- o rder Cli p ped Stoc h ast i c Optimiz a tion, a light-weight second-order optimizer that uses an inexpensive stochastic estimate of the diagonal of the Hessian as a pre-conditioner and a clipping mechanism to control the worst-case update size. On pre-training language models such as GPT-2, Sophia achieves the same validation pre-training loss with 50 % \% fewer number of steps than Adam. Because Sophia maintains almost the memory and average time per step, the speedup also translates to 50 % \% less total compute and 50 % \% less wall-clock time (See Figure 1 (a)&(b)). We also note that comparing the run-time to achieve the same loss is a correct way to compare the speed of optimizers for LLMs; see Section 3.2 for more details.

Moreover, the scaling law based on model size from 125M to 770M is in favor of Sophia over Adam—the gap between Sophia and Adam with 100K steps increases as the model size increases (Figure 1 (c)). In particular, Sophia on a 540M-parameter model with 100K steps gives the same validation loss as Adam on a 770M-parameter model with 100K steps. Note that the latter model needs 40% more training time and 40% more inference cost.

Concretely, Sophia estimates the diagonal entries of the Hessian of the loss using a mini-batch of examples every k k step (with k = 10 k=10 in our implementation). We consider two options for diagonal Hessian estimators: (a) an unbiased estimator that uses a Hessian-vector product with the same run-time as a mini-batch gradient up to a constant factor, and (b) a biased estimator that uses one mini-batch gradient calculated with resampled labels. Both the two estimators only introduce 5% overheads per step (in average). At every step, Sophia updates the parameter with an exponential moving average (EMA) of the gradient divided by the EMA of the diagonal Hessian estimate, subsequently clipped by a scalar. (All operations are element-wise.) See Algorithm 3 for the pseudo-code.

Additionally, Sophia can be seamlessly integrated into existing training pipelines, without any special requirements on the model architecture or computing infrastructure. With the either of the Hessian estimators, Sophia only require either standard mini-batch gradients, or Hessian-vector products which are supported in auto-differentiation frameworks such as PyTorch ( Paszke et al., 2019 ) and JAX ( Bradbury et al., 2018 ) .

Thanks to the Hessian-based pre-conditioner, Sophia adapts more efficiently, than Adam does, to the heterogeneous curvatures in different parameter dimensions, which can often occur in the landscape of LLMs losses and cause instability or slowdown. Sophia has a more aggressive pre-conditioner than Adam—Sophia applies a stronger penalization to updates in sharp dimensions (where the Hessian is large) than the flat dimensions (where the Hessian is small), ensuring a uniform loss decrease across all parameter dimensions. In contrast, Adam’s updates are mostly uniform across all parameter dimensions, leading to a slower loss decrease in flat dimensions. (See Section 2.1 for more discussions.) These make Sophia converge in fewer iterations. Thanks to the light-weight diagonal Hessian estimate, the speed-up in the number of steps translates to a speed-up in total compute and wall-clock time.

Sophia’s clipping mechanism controls the worst-case size of the updates in all directions, safeguarding against the negative impact of inaccurate Hessian estimates, rapid Hessian changes over time, and non-convex landscape (with which the vanilla Newton’s method may converge to local maxima or saddle points instead of local minima). The safeguard allows us to estimate Hessian infrequently (every k = 10 k=10 step) and stochastically. In contrast, prior second-order methods often update Hessian estimates every step ( Martens & Grosse, 2015 ; Grosse & Martens, 2016 ; Anil et al., 2020 ; Yao et al., 2021 ) .

We provide theoretical analyses of much simplified versions of Sophia on convex functions. The runtime bound 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.

## 2 Method

We first instantiate gradient descent (GD) and Adam on a simplified 2D problem and motivate the use of second-order information and per-coordinate clipping in Section 2.1 . Then, we present Sophia in detail in Section 2.2 , and the pseudo-code in Algorithm 3 . We introduce two choices of estimators of diagonal Hessian used in Sophia in Section 2.3 .

### 2.1 Motivations

Heterogeneous curvatures. The loss functions of modern deep learning problems often have different curvatures across different parameter dimensions ( Sagun et al., 2016 ; Ghorbani et al., 2019 ; Zhang et al., 2020 ; Yao et al., 2020 ) . E.g., on a 125M-parameter GPT-2 model, Figure 3 shows that the distribution of positive diagonal entries of the Hessian is dispersed.

We demonstrate the limitations of Adam and GD on heterogeneous landscapes by considering a two-dimensional loss function L ⁡ ( θ [ 1 ] , θ [ 2 ] ) = L 1 ​ ( θ [ 1 ] ) + L 2 ​ ( θ [ 2 ] ) L(\theta_{[1]},\theta_{[2]})=L_{1}(\theta_{[1]})+L_{2}(\theta_{[2]}) where L 1 L_{1} is much sharper than L 2 L_{2} . We plot the loss landscape of L ⁡ ( θ [ 1 ] , θ [ 2 ] ) L(\theta_{[1]},\theta_{[2]}) in Figure 2 . 1 1 1 Concretely, in Figure 2 , L 1 ​ ( θ [ 1 ] ) = 8 ​ ( θ [ 1 ] − 1 ) 2 ​ ( 1.3 ​ θ [ 1 ] 2 + 2 ​ θ [ 1 ] + 1 ) L_{1}(\theta_{[1]})=8(\theta_{[1]}-1)^{2}(1.3\theta_{[1]}^{2}+2\theta_{[1]}+1) and L 2 ​ ( θ [ 2 ] ) = 1 / 2 ​ ( θ [ 2 ] − 4 ) 2 L_{2}(\theta_{[2]})=1/2(\theta_{[2]}-4)^{2} . For simplicity, we discuss GD and deterministic versions of Adam. Recall that GD’s update in this setting is: θ [ 1 ] ← θ [ 1 ] − η ⋅ L 1 ′ ​ ( θ [ 1 ] ) and θ [ 2 ] ← θ [ 2 ] − η ⋅ L 2 ′ ​ ( θ [ 2 ] ) . \displaystyle\theta_{[1]}\leftarrow\theta_{[1]}-\eta\cdot L^{\prime}_{1}(\theta_{[1]})\ \ \textup{and}\ \ \theta_{[2]}\leftarrow\theta_{[2]}-\eta\cdot L^{\prime}_{2}(\theta_{[2]})\,. (1) A common simplification of Adam that is more amenable to analysis ( Balles & Hennig, 2018 ; Bernstein et al., 2018 ; Zhuang et al., 2020 ; Kunstner et al., 2023 ) is SignGD, which dates back to RProp ( Braun & Riedmiller, 1992 ) that motivated RMSProp ( Hinton et al., 2012 ) and Adam. Observe that without using the EMA (for both the gradient and second moments of the gradient), Adam’s update is simplified to η ⋅ ∇ L ​ ( θ ) / | ∇ L ​ ( θ ) | = η ⋅ sign ​ ( ∇ L ​ ( θ ) ) \eta\cdot\nabla L(\theta)/|\nabla L(\theta)|=\eta\cdot\textup{sign}(\nabla L(\theta)) (where all operations are entry-wise), which is called SignGD. Applying the update rule to our setting gives: θ [ 1 ] ← θ [ 1 ] − η ⋅ sign ​ ( L 1 ′ ​ ( θ [ 1 ] ) ) and θ [ 2 ] ← θ [ 2 ] − η ⋅ sign ​ ( L 2 ′ ​ ( θ [ 2 ] ) ) . \displaystyle\theta_{[1]}\leftarrow\theta_{[1]}-\eta\cdot\textup{sign}(L^{\prime}_{1}(\theta_{[1]}))\ \ \textup{and}\ \ \theta_{[2]}\leftarrow\theta_{[2]}-\eta\cdot\textup{sign}(L^{\prime}_{2}(\theta_{[2]}))\,. (2)

Limitations of GD and SignGD (Adam). It is well known that the optimal learning rate of GD should be proportional to the inverse of the curvature, that is, the Hessian/second derivative at the local minimum. More precisely, let h 1 h_{1} and h 2 h_{2} be the curvatures of L 1 L_{1} and L 2 L_{2} at the local minimum (and thus h 1 ≫ h 2 h_{1}\gg h_{2} ). The optimal learning rate for the update of θ [ 1 ] \theta_{[1]} in equation ( 1 ) is ≍ 1 / h 1 \asymp 1/h_{1} , which is much smaller than the optimal learning rate that the update of θ [ 2 ] \theta_{[2]} needs, which is ≍ 1 / h 2 \asymp 1/h_{2} . As a result, the largest shared learning rate can only be 1 / h 1 1/h_{1} ; consequently, the convergence in θ [ 2 ] \theta_{[2]} dimension is slow as demonstrated in the brown curve in Figure 2 .

The update size of SignGD is the learning rate η \eta in all dimensions. The same update size translates to less progress in decreasing the loss in the flat direction than in the sharp direction. As observed from the yellow curve in Figure 2 , the progress of SignGD in the flat dimension θ [ 2 ] \theta_{[2]} is slow because each step only decreases the loss L 2 ​ ( θ [ 2 ] ) L_{2}(\theta_{[2]}) slightly. On the other hand, along the direction θ [ 1 ] \theta_{[1]} , the iterate quickly travels to the valley in the first three steps and then starts to bounce. To fully converge in the sharp dimension, the learning rate η \eta needs to decay to 0, which will exacerbate the slow convergence in the flat dimension θ [ 2 ] \theta_{[2]} . The trajectory of Adam in this example is indeed similar to SignGD, which is also plotted as the red curve in Figure 2 .

The behavior of SignGD and Adam above indicates that a more aggressive pre-conditioning is needed—sharp dimensions should have relatively smaller updates than flat dimensions so that the decrease of loss is equalized in all dimensions. As suggested by well-established literature on second-order optimization ( Boyd & Vandenberghe, 2004 ) for convex functions, the optimal pre-conditioner should be the Hessian, which captures the curvature on each dimension; as in Newton’s method, the update is the gradient divided by the Hessian in each dimension: θ [ 1 ] ← θ [ 1 ] − η ⋅ L 1 ′ ​ ( θ [ 1 ] ) / h 1 and θ [ 2 ] ← θ [ 2 ] − η ⋅ L 2 ′ ​ ( θ [ 2 ] ) / h 2 . \displaystyle\theta_{[1]}\leftarrow\theta_{[1]}-\eta\cdot L^{\prime}_{1}(\theta_{[1]})/h_{1}\ \ \textup{and}\ \ \theta_{[2]}\leftarrow\theta_{[2]}-\eta\cdot L^{\prime}_{2}(\theta_{[2]})/h_{2}\,. (3)

Limitations of Newton’s method. Nevertheless, Newton’s method has known limitations as well. For non-convex functions, vanilla Newton’s method could converge to a global maximum when the local curvature is negative. In the blue curve of Figure 2 , Newton’s method quickly converges to a saddle point instead of a local minimum. The curvature might also change rapidly along the trajectory, making the second-order information unreliable. To address these limitations, we propose considering only pre-conditioners that capture positive curvature, and introduce a pre-coordinate clipping mechanism to mitigate the rapid change of Hessian (more detail in Section 2.2 ). Applying our algorithm on the toy case results in the following update: θ [ 1 ] ← θ [ 1 ] − η ⋅ clip ​ ( L 1 ′ ​ ( θ [ 1 ] ) / max ⁡ { h 1 , ϵ } , ρ ) ​ and ​ θ [ 2 ] ← θ [ 2 ] − η ⋅ clip ​ ( L 2 ′ ​ ( θ [ 2 ] ) / max ⁡ { h 2 , ϵ } , ρ ) , \displaystyle\theta_{[1]}\leftarrow\theta_{[1]}-\eta\cdot\textup{clip}(\nicefrac{{L^{\prime}_{1}(\theta_{[1]})}}{{\max\{h_{1},\epsilon\}}},\rho)\ \textup{and}\ \theta_{[2]}\leftarrow\theta_{[2]}-\eta\cdot\textup{clip}(\nicefrac{{L^{\prime}_{2}(\theta_{[2]})}}{{\max\{h_{2},\epsilon\}}},\rho)\,, (4) where ρ \rho is a constant to control the worst-case update size, ϵ \epsilon is a very small constant (e.g., 1e-12), which avoids dividing by 0. When the curvature of some dimension is rapidly changing or negative and thus the second-order information is misleading and possibly leads to a huge update before clipping, the clipping mechanism kicks in and the optimizer defaults to SignGD (even though this is sub-optimal for benign situations). Numerous prior methods such as trust region ( Conn et al., 2000 ) , backtracking line search ( Boyd & Vandenberghe, 2004 ) , and cubic regularization ( Nesterov & Polyak, 2006 ) also tackle the same issue of Newton’s method, but the clipping mechanism is much simpler and more efficient.

As shown in the black curve in Fig. 2 , the update in equation ( 4 ) starts off similarly to SignGD due to the clipping mechanism in the non-convex region, making descent opposed to converging to a local maximum. Then, in the convex valley, it converges to the global minimum with a few steps. Compared with SignGD and Adam, it makes much faster progress in the flat dimension θ [ 2 ] \theta_{[2]} (because the update is bigger in dimension θ [ 2 ] \theta_{[2]} ), while avoiding boucing in the sharp dimension θ [ 1 ] \theta_{[1]} (because the update is significantly shrunk in the sharp dimension θ [ 1 ] \theta_{[1]} ).

### 2.2 Sophia: S econd- o rder Cli p ped Stoc h astic Opt i miz a tion