Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training — part 3/7
Cleaned-text code points [30020, 45778). Bibliography omitted here; source.md includes it.
Mini-batch version. Given a mini-batch of inputs { ( x b , y b ) } b = 1 B \{(x_{b},y_{b})\}_{b=1}^{B} . The most natural way to build an estimator for the diagonal of the Gauss-Newton matrix for the Hessian of the mini-batch loss is using 1 B ∑ b = 1 B ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) ⊙ ∇ θ ℓ ce ( f ( θ , x b ) , y ^ b ) , \displaystyle\frac{1}{B}\sum_{b=1}^{B}\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\odot\nabla_{\theta}\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\,, (11) where y ^ b \hat{y}_{b} ’s are labels sampled from the model on inputs x b x_{b} ’s respectively. However, as noted by Grosse (2022) , implementing this estimator is inconvenient under the current auto-differentiation frameworks, where the users only have access to the average gradient over a mini-batch (as opposed to the individual ones). Fortunately, by the Bartlett’s first identity ( Bartlett, 1953 ) (which generally holds for the negative log-likelihood loss of any probabilistic model), we have: ∀ b , 𝔼 y ^ b ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) = 0 . \displaystyle\forall b,~~\mathbb{E}_{\hat{y}_{b}}\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})=0\,. (12) Let L ^ ( θ ) = 1 B ∑ b = 1 B ℓ ce ( f ( θ , x b ) , y ^ b ) \widehat{L}(\theta)=\frac{1}{B}\sum_{b=1}^{B}\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b}) be the mini-batch loss on the sampled labels (as opposed to the original labels). Observing that y ^ b \hat{y}_{b} ’s are independent with each other, we have 𝔼 y ^ b ′ s [ B ⋅ ∇ θ L ^ ( θ ) ⊙ ∇ θ L ^ ( θ ) ] \displaystyle\mathbb{E}_{\hat{y}_{b}^{\prime}s}\left[B\cdot\nabla_{\theta}\widehat{L}(\theta)\odot\nabla_{\theta}\widehat{L}(\theta)\right] = 𝔼 y ^ b ′ s [ 1 B ∑ b = 1 B ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) ⊙ ∑ b = 1 B ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) ] \displaystyle=\mathbb{E}_{\hat{y}_{b}^{\prime}s}\left[\frac{1}{B}\sum_{b=1}^{B}\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\odot\sum_{b=1}^{B}\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\right] = 𝔼 y ^ b ′ s [ 1 B ∑ b = 1 B ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) ⊙ ∇ ℓ ce ( f ( θ , x b ) , y ^ b ) ] . \displaystyle=\mathbb{E}_{\hat{y}_{b}^{\prime}s}\left[\frac{1}{B}\sum_{b=1}^{B}\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\odot\nabla\ell_{\textup{ce}}(f(\theta,x_{b}),\hat{y}_{b})\right]\,. (13) Note that the RHS of Equation 13 is the same as the expectation of Equation 11 , which, by Equation 10 , also equals to the diagonal of the Gauss-Newton matrix for the mini-batch loss. Hence, we use B ⋅ ∇ θ L ^ ( θ ) ⊙ ∇ θ L ^ ( θ ) B\cdot\nabla_{\theta}\widehat{L}(\theta)\odot\nabla_{\theta}\widehat{L}(\theta) as the estimator.
GNB estimator for exponential family. If y y is drawn from an exponential family p ( y , η ) p(y;\eta) where the natural parameter η \eta is set to be f ( θ , x ) f(\theta,x) and the loss function ℓ ( f ( θ , x ) , y ) \ell(f(\theta,x),y) is the negative log-likelihood loss for the corresponding probabilistic distribution, then all the derivations above still follow because (1) S S still only depends on f ( θ , x ) f(\theta,x) but not y y , and (2) both the first and second Bartlett’s identities still hold.
GNB estimator for squared loss. When y , f ( θ , x ) ∈ ℝ y,f(\theta,x)\in\mathbb{R} and ℓ ( f ( θ , x ) , y ) = 1 2 ( f ( θ , x ) − y ) 2 \ell(f(\theta,x),y)=\frac{1}{2}(f(\theta,x)-y)^{2} , the S S matrix is identity, and thus one can simply use J θ f ( θ , x ) J θ f ( θ , x ) ⊤ J_{\theta}f(\theta,x)J_{\theta}f(\theta,x)^{\top} as the estimator. 3 3 3 One can verify that the GNB estimator gives the same quantity in expectation if we use a probabilistic model y ∼ 𝒩 ( f ( θ , x ) , σ 2 ) y\sim\mathcal{N}(f(\theta,x),\sigma^{2}) and go through the derivation of the GNB estimator.
To the best of our knowledge, Wei et al. (2020) is the first paper that uses this estimator of Gauss-Newton matrix. Given the use Bartlett’s first and second identities that are central to the estimator, we call it Gauss-Newton-Bartlett (GNB) estimator.
Comparisons of Hessian estimators. The Hutchinson’s estimator does not assume any structure of the loss, but requires a Hessian-vector product. The GNB estimator only estimates the Gauss-Newton term but always gives a positive semi-definite (non-negative) diagonal Hessian estimate. The PSDness ensures that the pre-conditioned update is always a descent direction ( Dennis Jr & Schnabel, 1996 ) . The GNB estimator can also be easily extended to the negative log-likelihood loss of any exponential family distribution, and be adapted to estimating the trace of the Gauss-Newton matrix as in Wei et al. (2020) or efficiently implementing the product of Gauss-Newton matrix with a vector. The authors suspect the GNB estimator has a smaller variance than the Hutchinson’s estimator, but more empirical and theoretical investigation are needed to support the hypothesis.
## 3 Experiments
We name the algorithm using the Hutchinson’s estimator and the GNB estimator Sophia-H and Sophia-G, respectively. We evaluate Sophia on auto-regressive language modeling with GPT-2 ( Radford et al., 2019 ) of model sizes ranging from 125M to 770M, and GPT NeoX ( Black et al., 2022 ) of sizes 1.5B and 6.6B. Results indicate that Sophia is 2x faster than AdamW ( Loshchilov & Hutter, 2017 ) in number of steps, total compute, and wall-clock time across all model sizes. Moreover, the scaling law is in favor of Sophia over AdamW.
### 3.1 Experimental Setup
Language modeling. We train autoregressive models on OpenWebText ( Gokaslan & Cohen, 2019 ) and the Pile ( Gao et al., 2020 ) . Following standard protocol, we set the context length of GPT-2 to 1024, and the context length of GPT-2 NeoX ( Black et al., 2022 ) to 2048. We consider GPT-2 with 125M (small), 355M (medium), and 770M (large) parameters, and GPT NeoX with 1.5B and 6.6B parameters, respectively. Detailed model configurations are deferred to Section B.2 .
Baselines. We mainly compare Sophia and Adam with decoupled weight decay (AdamW) ( Loshchilov & Hutter, 2017 ) which is the dominantly used optimizer on language modeling tasks, AdaHessian ( Yao et al., 2021 ) which uses the EMA of the square of the diagonal Hessian estimate in its denominator, and Lion ( Chen et al., 2023 ) , which is an first-order adaptive optimizer discovered by symbolic search. For the 30M model, all hyperparameters are tuned with grid search. For other models, all hyperparmeters but the peak learning rate are configured as identical to those found on the 30M model. For models with size 125M and 355M, the peak learning rates are obtained through grid search. For larger models, we gradually increase the peak learning rate to search for the largest possible peak learning rate such that the training does not blow up, and ensure that the chosen learning rate is approximately the largest in the sense that 1.25 times the chosen learning rate will lead to a blow-up. For AdamW we found the well-established practice ( β 1 = 0.9 \beta_{1}=0.9 and β 2 = 0.95 \beta_{2}=0.95 ) works consistently better than other choices ( Radford et al., 2019 ; Karamcheti et al., 2021 ) . For Lion, we use β 1 = 0.95 \beta_{1}=0.95 and β 2 = 0.98 \beta_{2}=0.98 following Chen et al. (2023) . Although Chen et al. (2023) suggests using 0.1 0.1 times the learning rate (LR) of AdamW for vision tasks, we find out the LR should be larger on LMs from the grid search. For AdaHessian, we found β 1 = 0.92 \beta_{1}=0.92 and β 2 = 0.99 \beta_{2}=0.99 works the best in the grid search. Details on hyperparameter tuning are deferred to Section B.1 .
Implementation. We set batch size to 480 for GPT-2 and 2048 for GPT NeoX. We use cosine LR schedule with the final LR equal to 0.05 times the peak LR, following Rae et al. (2021) . We use the standard gradient clipping (by norm) threshold 1.0. We adopt a fixed 2k steps of LR warm-up. For Sophia, we use β 1 = 0.96 \beta_{1}=0.96 , β 2 = 0.99 \beta_{2}=0.99 , ϵ = \epsilon= 1e-12 and update diagonal Hessian every 10 steps. For Sophia-H, we use only a subset of 32 examples from the mini-batch to calculate the diagonal Hessian to further reduce overhead. For Sophia-G, we use a subset of 240 examples from the mini-batch to calculate the diagonal Gauss-Newton. We implement the algorithms in PyTorch ( Paszke et al., 2019 ) and JAX ( Bradbury et al., 2018 ) and train all the models in bfloat16. The 125M and 355M models are trained on A5000 GPUs, while the 770M models are trained on A100 GPUs. We use a TPU v3-128 slice to train the 1.5B and 6.6B GPT NeoX.
#### Hyperparamter tuning strategy.
We refer to Section B.1 for the details on hyperparameters and only discuss two key hyperparameters, γ \gamma and the peak learning rate η \eta in the main text. Similar to the protocol of baselines, all other hyperparameters are tuned on a 30M model and remain fixed for all the model sizes. For the peak learning rate and γ \gamma , we found the following strategy general works well, and delivers almost the same performance as those found by grid search. • On a small model, tune γ \gamma to make the proportion of coordinates where the update is not clipped (i.e., | m t / max { γ ⋅ h t , ϵ } | < 1 |m_{t}/\max\{\gamma\cdot h_{t},\epsilon\}|<1 ) in the range of 10 % − 50 % 10\%-50\% . If the proportion is too large (or too small), multiply γ \gamma by 0.5 0.5 (or 2 2 ) and restart. The same γ \gamma likely can be transferred to models with the same architecture and data but different number of parameters. We use γ = 0.01 \gamma=0.01 for Sophia-H and γ = 0.05 \gamma=0.05 for Sophia-G in this paper.
• Suppose we already find a suitable γ \gamma following the above procedure. We can then set the learning rate of Sophia to be either 3-5 times the learning rate that one would have used for Lion, or 0.8 times the learning rate that one would have used for AdamW.
### 3.2 Evaluation
#### Methodology for comparing the optimizers for LLMs.
This paper argues that one correct (and arguably preferred) way of claiming optimizer O 2 O_{2} is 2x faster than optimizer O 1 O_{1} is comparing the following two experiments (for a reasonable variety of T T ): 1. running optimizer O 1 O_{1} (e.g. Adam) with T T steps, with the optimal learning rate and learning rate schedule (tuned for running for T T steps),
2. running optimizer O 2 O_{2} (e.g., Sophia) with T / 2 T/2 steps, with any learning rate schedule.
If Experiment 2 achieves a loss that is smaller than or equal to the loss of Experiment 1 (for a reasonable sets of choices of T T ), then we say optimizer O 2 O_{2} is 2x faster than optimizer O 1 O_{1} .
In more formal language, suppose the loss (or any minimization metric) of an optimizer O O with runtime budget T T and hypeparameter H H (excluding the runtime T T ) is denoted by Eval ( O , T , H ) \textup{Eval}(O,T,H) . Then, we say optimizer O 2 O_{2} is k k -times faster than O 1 O_{1} if ∃ H 2 , min H 1 Eval ( O 1 , T , H 1 ) ≥ Eval ( O 2 , T / k , H 2 ) . \displaystyle\exists H_{2},~~\min_{H_{1}}~\textup{Eval}(O_{1},T,H_{1})\geq\textup{Eval}(O_{2},T/k,H_{2})\,. (14) Note that the equation above implies min H 1 Eval ( O 1 , T , H 1 ) ≥ min H 2 Eval ( O 2 , T / k , H 2 ) \min_{H_{1}}\textup{Eval}(O_{1},T,H_{1})\geq\min_{H_{2}}\textup{Eval}(O_{2},T/k,H_{2}) , but is simpler to verify than the latter. We note that many modern learning rate schedulers such as cosine learning rate ( Loshchilov & Hutter, 2016 ) are highly sensitive to a pre-specific total number of steps. The optimal hyperparameters are also sensitive to the total budget T T . Thus, tuning the hyperparameters of the baseline O 1 O_{1} with a pre-specified budget T T (aka, minimizing H 1 H_{1} given T T ) is critical to ensure fair comparison. (On the other hand, the tuning strategy for H 2 H_{2} is less important because we just need to show the existence of a good H 2 H_{2} .)
Moreover, we note that, in equation ( 14 ), we need to insist ≥ \geq without any approximations. The difference in losses of various experiments might be seemingly small— Eval ( Adam , T / 2 , H ) \textup{Eval}(\textup{Adam},T/2,H) is likely to be only slightly higher than Eval ( Adam , T , H ) \textup{Eval}(\textup{Adam},T,H) (assuming the peak learning rate is a hyperparameter), and thus allowing any approximation in the criterion ( 14 ) may lead to the fallacy statement that “Adam is 2x faster than Adam”. In fact, Figure 4 (a) shows that even with the same peak learning rate, the learning rate in a run with T / 2 T/2 steps decays faster than the learning rate in a run with T T steps. Moreover, the T / 2 T/2 -steps run tends to have a initial faster decay of loss than the T T -steps run but stop at a higher final loss, and the latter is not a continuation of the former.
Finally, a tempting comparison would be comparing the following experiment with Experiment 1: 2’. running optimizer O 2 O_{2} (e.g., Sophia) with T T steps, with any learning rate schedule, and recording the performance of the T / 2 T/2 -th checkpoint.
We note that if one were able to find an optimizer O 2 O_{2} that is 2x faster than O 1 O_{1} under the proposed criterion (comparing 1&2), then one can also design another optimizer O 2 ′ O_{2}^{\prime} (which run for T T steps) that is 2x faster than O 1 O_{1} under the second criterion (comparing 1&2’): first running optimizer O 2 O_{2} for T / 2 T/2 steps and do nothing for the rest of the T / 2 T/2 steps. Therefore, even though we will also provide comparison between Experiment 1&2’ in Figure 5 as an alternative, we use the comparison between Experiment 1&2 quantitatively in most of the paper.
#### Technical details.
Following the methodology above, we train baselines and Sophia for 100K, 200K, or 400K, and mainly compare 400k-steps baseline vs 200K-steps Sophia, and 200k-steps baseline vs 100k-steps Sophia. We primarily evaluate the models with their log perplexity and plot the loss curves. We also report in-context learning results (with 2-shot exemplars and greedy decoding) on SuperGLUE ( Wang et al., 2019 ) . We average the results of 5 prompts (Section B.3 ).
### 3.3 Results
Figure 5 illustrates the validation loss curve (token-level log perplexity) on OpenWebText or the Pile with the same number of steps. Sophia-H and Sophia-G consistently achieve better validation loss than AdamW, Lion, and AdaHessian, while Sophia-G is better than Sophia-H. As the model size grows, the gap between Sophia and baselines also becomes larger. Note that the cost of each step AdaHessian is more than 2x that of AdamW, while Sophia only introduces a per step compute overhead of less than 5 % \% . Sophia-H and Sophia-G both achieve a 0.04 smaller validation loss on the 355M model (Figure 5 (b)). Sophia-G and Sophia-H achieve a 0.05 smaller validation loss on the 770M model (Figure 5 (c)), with the same 100K steps. This is a significant improvement since according to scaling laws in this regime ( Kaplan et al., 2020 ; Hoffmann et al., 2022 ) and results in Figure 4 , an improvement in loss of 0.05 is equivalent to 2x improvement in terms of number of steps or total compute to achieve the same validation loss.
Sophia is 2x faster in terms of number of steps, total compute and wall-clock time. The improvement in validation loss brought by Sophia can be translated into reduction of number of steps or total compute. In Figure 1 (a)(b)(c) and Figure 4 , we evaluate the optimizers by comparing the number of steps or total compute needed to achieve the same validation loss level . As can be observed in Figure 1 (a)(b)(c), Sophia-H and Sophia-G achieve a 2x speedup compared with AdamW and Lion across different model sizes.