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

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

Results with different number of steps. Runs with different number of steps and their comparison are provided in Figure 10 . Across different choices of the total number of steps, Sophia outperforms AdamW and Lion with a large margin as the main experiments we presented in Section 3.3 .

## Appendix B Additional Experiment Details

### B.1 Hyperparamter Tuning

The hyperparameters we consider for baselines are as follows: peak learning rate, β 1 \beta_{1} , β 2 \beta_{2} and weight decay. All hyperparameters except the peak learning rate are tuned with grid search on a 30M GPT-2 trained for 50K steps. The peak learning rate is tuned on models of different sizes with grid search separately. We search β 1 \beta_{1} in [ 0.8 , 0.9 , 0.95 , 0.96 , 0.99 ] [0.8,0.9,0.95,0.96,0.99] and β 2 \beta_{2} in [ 0.9 , 0.95 , 0.98 , 0.99 , 0.995 ] [0.9,0.95,0.98,0.99,0.995] . Weight decay is chosen from 0.05 , 0.1 , 0.2 , 0.5 0.05,0.1,0.2,0.5 . For Lion, we also include β 1 = 0.95 , β 2 = 0.98 \beta_{1}=0.95,\beta_{2}=0.98 as suggested by Chen et al. (2023) . On 30M models, we found AdamW is sensitive to the choice of β 1 \beta_{1} but not β 2 \beta_{2} . β 1 = 0.9 \beta_{1}=0.9 works the best for AdamW while β 1 = 0.95 \beta_{1}=0.95 works the best for Lion. We use β 2 = 0.95 \beta_{2}=0.95 for AdamW since this is the dominantly used configuration in the LLM pre-training literature, and β 2 = 0.98 \beta_{2}=0.98 for Lion as it is recommended by Chen et al. (2023) . We found that weight decay 0.1 works the best for AdamW, while 0.2 works the best for Lion and Sophia.

For peak learning rate on 125M and 355M, we perform a grid search. For larger models, we search for the largest possible learning rate with which training does not blow up for each model in the following list: [6e-4, 4e-4, 3e-4, 2e-4, 1.5e-4, 1.2e-4, 1e-4, 8e-5, 6e-5, 4e-5]. For example, a 6.6B GPT NeoX model work with 1.2e-4 peak learning rate, but the loss will blow up if we increase the learning rate to 1.5e-4 as shown in Figure 12 . The result of grid search of peak learning rate is provided in Table 2 . We provide additional fine-grained tuning of peak learning rate of the 355M GPT-2 and 1.5B GPT Neo-X in Figure 12 . Concretely, we tried 3.6e-4, 4.2e-4, 4.5e-4 and 4.8e-4 for the 355M model, 1.8e-4 for the 1.5B GPT Neo-X. Results indicate that the learning rate choices in Table 2 is indeed the optimal for all the values we have tried.

We use β 1 = 0.96 \beta_{1}=0.96 , β 2 = 0.99 \beta_{2}=0.99 , ϵ = \epsilon= 1e-12 and k = 10 k=10 for Sophia. We adopt the following procedure to obtain these hyper parameters. We first fix γ = 0.01 \gamma=0.01 , k = 10 k=10 , and tune β 1 \beta_{1} and β 2 \beta_{2} with grid search on a 30M model, and directly use β 1 \beta_{1} and β 2 \beta_{2} from the 30M model on models of larger sizes. Similar to AdamW, we find that Sophia is not sensitive to β 2 \beta_{2} . We then fix β 1 = 0.96 \beta_{1}=0.96 , β 2 = 0.99 \beta_{2}=0.99 and tuning k = 10 k=10 . As shown in in Figure 8 (a), k = 10 k=10 is better than k = 1 k=1 or k = 100 k=100 in terms of the balance between convergence speed and the computation overhead.

After finding out β 1 = 0.96 \beta_{1}=0.96 , β 2 = 0.99 \beta_{2}=0.99 , ϵ = 1 ​ e − 12 \epsilon=1e-12 and k = 10 k=10 with the method above, we tune γ \gamma and peak learning rate jointly. We first 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\% . We search for γ \gamma in the list of [0.005,0.01,0.02,0.05,0.1,0.2]. As a result we find out γ = 0.01 \gamma=0.01 works the best for Sophia-H while γ = 0.05 \gamma=0.05 works the best for Sophia-G. We then fix β 1 , β 2 = 0.99 , ϵ , γ , k \beta_{1},\beta_{2}=0.99,\epsilon,\gamma,k for all larger models.

To tune the peak learning rate, we adopt the same procedure as we use for baseline methods. The result of grid search of peak learning rate is also provided in Table 2 .

### B.2 Model and Implementation Details

We consider three sizes of GPT-2 corresponding to small, medium, and large in Radford et al. (2019) . We also introduce a 30M model for efficient hyperparameter grid search and a 540M model for scaling law visualization. We provide the model specifications in Table 2 . We use the nanoGPT ( https://github.com/karpathy/nanoGPT/ ) code base. Following nanoGPT, we use GELU activations and disable bias and Dropout Srivastava et al. (2014) during pre-training.

GPT-2 models are trained on OpenWebText ( Gokaslan & Cohen, 2019 ) . The text is tokenized with the GPT-2 tokenizer ( Radford et al., 2019 ) . We use the train and validation split from nanoGPT. The training set contains 9B tokens, and the validation set contains 4.4M tokens.

We use distributed data parallel with gradient accumulation to enable a batch size of 480. All models are trained with bfloat16. The 125M and 355M models are trained on machines with 10 A5000 GPUs, while the 770M models are trained on an AWS p4d.24xlarge instance with 8 A100 GPUs.

We consider 1.5B and 6.6B GPT NeoX ( Black et al., 2022 ) models trained on the Pile ( Gao et al., 2020 ) . The models use GPT NeoX tokenizers. We use levanter ( https://github.com/stanford-crfm/levanter/tree/main ) for GPT NeoX. We use fully sharded data parallel with gradient accumulation to enable a batch size of 512 for the 1.5B model and 1024 for the 6.6B model. These models are trained on a TPU v3-128 slice.

We observed AdamW and Lion does not perform well on standard transformers which are larger than 355M. The iterates become unstable when the learning rate is close to the choice of Radford et al. (2019) . We introduce scaling attention by the inverse of layer index to address this issue following Karamcheti et al. (2021) ; Wolf et al. (2020) . Note that Sophia does not need this trick as mentioned in Section 3.4 .

### B.3 Downtream Evaluation

We perform few-shot evaluation of the models on 4 subtasks of SuperGLUE. We use 2-shot prompting and greedy decoding. The prompt consists of an instruction followed by two examples. The examples are sampled from the train split while we report the accuracy on validation split averaged over 5 selection of exemplars. Prompts for each subtask are illustrated in Figure 11 .

## Appendix C Limitations

Scaling up to larger models and datasets. Although Sophia demonstrates scalability up to 6.6B-parameter models, and there is no essential constraints from further scaling up, we do not compare with AdamW and Lion on larger models due to limited resources. We believe Sophia is faster than AdamW and Lion on larger models given the improvement in scaling laws and better pre-training stability.

Holistic downstream evaluation. We evaluate pre-trained checkpoints on the pre-training validation losses and only four SuperGLUE subtasks. The limited evaluation in downstream evaluation is partly due to the limited model size, because language models at this scale do not have enough capabilities such as in-context learning, and mathematical reasoning.

Evaluation on other domains. While this paper focuses on optimizers for large language modeling, a more general optimizer should also be evaluated in other domains such as computer vision, reinforcement learning, and multimodal tasks. Due to the limitation of computation resources, we leave the application to other domains and models to future works.

## Appendix D Theoretical Analyses: Details of Section 4

Theorem 4.3 is a direct combination of the Lemma D.10 (Descent Lemma), Lemma D.9 and Lemma D.11 . In the analysis, there will be two phases. In the first phase decrease loss to μ ​ ρ 2 8 \frac{\mu\rho^{2}}{8} in 8 ​ L ⁡ ( θ ⁡ ( 0 ) ) − min ⁡ L η ​ μ ​ ρ 2 8\frac{L(\theta(0))-\min L}{\eta\mu\rho^{2}} steps. In the second phase, there will be an exponential decay of error.

###### Lemma D.1 .

Under Assumption 4.1 , we have that L ⁡ ( θ ) → ∞ L(\theta)\to\infty whenever ‖ θ ‖ 2 → ∞ \|\theta\|_{2}\to\infty .

###### Proof of Lemma D.1 .

By convexity of L L , we have ∀ θ ∈ ℝ d \forall\theta\in\mathbb{R}^{d} with ‖ θ − θ ∗ ‖ 2 ≥ 1 \left\|\theta-\theta^{*}\right\|_{2}\geq 1 , 1 ‖ θ − θ ∗ ‖ 2 ​ L ​ ( θ ) + ‖ θ − θ ∗ ‖ 2 − 1 ‖ θ − θ ∗ ‖ 2 ​ L ​ ( θ ∗ ) ≥ L ⁡ ( θ ∗ + θ − θ ∗ ‖ θ − θ ∗ ‖ 2 ) ≥ min ‖ θ ¯ ‖ 2 = 1 ⁡ L ⁡ ( θ ∗ + θ ¯ ) . \displaystyle\frac{1}{\left\|\theta-\theta^{*}\right\|_{2}}L(\theta)+\frac{\left\|\theta-\theta^{*}\right\|_{2}-1}{\left\|\theta-\theta^{*}\right\|_{2}}L(\theta^{*})\geq L(\theta^{*}+\frac{\theta-\theta^{*}}{\left\|\theta-\theta^{*}\right\|_{2}})\geq\min_{\left\|\bar{\theta}\right\|_{2}=1}L(\theta^{*}+\bar{\theta}). (17)

Since L L is strictly convex, Δ ≜ min ‖ θ ¯ ‖ 2 = 1 ⁡ L ⁡ ( θ ∗ + θ ¯ ) − L ⁡ ( θ ∗ ) > 0 \Delta\triangleq\min_{\left\|\bar{\theta}\right\|_{2}=1}L(\theta^{*}+\bar{\theta})-L(\theta^{*})>0 . Thus we conclude that L ⁡ ( θ ) ≥ ‖ θ − θ ∗ ‖ 2 ​ Δ + L ⁡ ( θ ∗ ) ≥ ( ‖ θ ‖ 2 − ‖ θ ∗ ‖ 2 ) ​ Δ + L ⁡ ( θ ∗ ) . \displaystyle L(\theta)\geq\left\|\theta-\theta^{*}\right\|_{2}\Delta+L(\theta^{*})\geq(\left\|\theta\right\|_{2}-\left\|\theta^{*}\right\|_{2})\Delta+L(\theta^{*}). (18) Therefore when ‖ θ ‖ 2 → ∞ \left\|\theta\right\|_{2}\to\infty , L ⁡ ( θ ) → ∞ L(\theta)\to\infty as well. ∎

Note that we don’t assume the Hessian of loss is Lipschitz. Assumption 4.2 only assumes the Hessian in a neighborhood of constant radius only differs by a constant in the multiplicative sense.

###### Lemma D.2 .

For any θ ∈ ℝ d \theta\in\mathbb{R}^{d} satisfying L ⁡ ( θ ) − min ⁡ L ≤ μ ​ R 2 4 L(\theta)-\min L\leq\frac{\mu R^{2}}{4} , it holds that ‖ θ − θ ∗ ‖ 2 ≤ 2 ​ L ⁡ ( θ ) − min ⁡ L μ ≤ R \left\|\theta-\theta^{*}\right\|_{2}\leq 2\sqrt{\frac{L(\theta)-\min L}{\mu}}\leq R .

###### Proof of Lemma D.2 .

We will prove by contradiction. Suppose there exists such θ \theta with L ⁡ ( θ ) ≤ μ ​ R 2 4 L(\theta)\leq\frac{\mu R^{2}}{4} but ‖ θ − θ ∗ ‖ 2 > 2 ​ L ⁡ ( θ ) − min ⁡ L μ \left\|\theta-\theta^{*}\right\|_{2}>2\sqrt{\frac{L(\theta)-\min L}{\mu}} . We consider θ ′ ≜ θ ∗ + 2 ​ L ​ ( θ ) μ ⋅ θ − θ ∗ ‖ θ − θ ∗ ‖ 2 \theta^{\prime}\triangleq\theta^{*}+\sqrt{\frac{2L(\theta)}{\mu}}\cdot\frac{\theta-\theta^{*}}{\left\|\theta-\theta^{*}\right\|_{2}} . Since θ ′ \theta^{\prime} is between θ \theta and θ ∗ \theta^{*} and that L L is strictly convex, we know that L ⁡ ( θ ′ ) < L ⁡ ( θ ) L(\theta^{\prime})<L(\theta) . However, by Taylor expansion on function f ⁡ ( t ) ≜ L ⁡ ( θ ∗ + t ⁡ ( θ ′ − θ ∗ ) ) f(t)\triangleq L(\theta^{*}+t(\theta^{\prime}-\theta^{*})) , we have that f ⁡ ( 1 ) = f ⁡ ( 0 ) + f ′ ​ ( 0 ) + f ′′ ​ ( t ) 2 , for some t ∈ [ 0 , 1 ] . \displaystyle f(1)=f(0)+f^{\prime}(0)+\frac{f^{\prime\prime}(t)}{2},\quad\text{for some $t\in[0,1]$.} (19) Note that ‖ θ ′ − θ ∗ ‖ 2 ≤ ‖ θ − θ ∗ ‖ 2 ≤ R \left\|\theta^{\prime}-\theta^{*}\right\|_{2}\leq\left\|\theta-\theta^{*}\right\|_{2}\leq R , by Assumption 4.2 and Assumption 4.1 , we have OPEN f ′′ ​ ( t ) = ( θ ′ − θ ∗ ) ⊤ ​ ∇ 2 L ​ ( t ​ θ ′ + ( 1 − t ) ​ θ ∗ ) ​ ( θ ′ − θ ∗ ) ≥ 1 2 ​ ( θ ′ − θ ∗ ) ⊤ ​ ∇ 2 L ​ ( θ ∗ ) ​ ( θ ′ − θ ∗ ) ≥ μ 2 ​ ‖ θ ′ − θ ∗ ‖ 2 2 = 2 ​ ( L ⁡ ( θ ) − min ⁡ L ) ) f^{\prime\prime}(t)=(\theta^{\prime}-\theta^{*})^{\top}\nabla^{2}L(t\theta^{\prime}+(1-t)\theta^{*})(\theta^{\prime}-\theta^{*})\geq\frac{1}{2}(\theta^{\prime}-\theta^{*})^{\top}\nabla^{2}L(\theta^{*})(\theta^{\prime}-\theta^{*})\geq\frac{\mu}{2}\left\|\theta^{\prime}-\theta^{*}\right\|_{2}^{2}=2(L(\theta)-\min L)) . Also note that f ⁡ ( 1 ) = L ⁡ ( θ ′ ) , f ⁡ ( 0 ) = L ⁡ ( θ ∗ ) f(1)=L(\theta^{\prime}),f(0)=L(\theta^{*}) and f ′ ​ ( 0 ) = 0 f^{\prime}(0)=0 , we conclude that L ⁡ ( θ ′ ) − L ⁡ ( θ ∗ ) ≥ L ⁡ ( θ ) − L ⁡ ( θ ∗ ) L(\theta^{\prime})-L(\theta^{*})\geq L(\theta)-L(\theta^{*}) , namely ( θ ′ ) ≥ L ⁡ ( θ ) (\theta^{\prime})\geq L(\theta) . Contradiction! ∎

###### Lemma D.3 .

For any θ ∈ ℝ d \theta\in\mathbb{R}^{d} satisfying that ‖ ∇ L ​ ( θ ) ‖ 2 ≤ R ​ μ 2 \left\|\nabla L(\theta)\right\|_{2}\leq\frac{R\mu}{2} , it holds that ‖ θ − θ ∗ ‖ 2 ≤ 2 ​ ‖ ∇ L ​ ( θ ) ‖ μ ≤ R \left\|\theta-\theta^{*}\right\|_{2}\leq\frac{2\left\|\nabla L(\theta)\right\|}{\mu}\leq R .

###### Proof of Lemma D.3 .

We will prove by contradiction. We consider function f ⁡ ( t ) ≜ ⟨ θ − θ ∗ ‖ θ − θ ∗ ‖ 2 , ∇ L ​ ( θ ∗ + t ⋅ θ − θ ∗ ‖ θ − θ ∗ ‖ 2 ) ⟩ f(t)\triangleq\left\langle\frac{\theta-\theta^{*}}{\left\|\theta-\theta^{*}\right\|_{2}},\nabla L(\theta^{*}+t\cdot\frac{\theta-\theta^{*}}{\left\|\theta-\theta^{*}\right\|_{2}})\right\rangle . Because of the strict convexity of L L , f f is a strict monotone increasing function. If ‖ θ − θ ∗ ‖ > 2 ​ ‖ ∇ L ​ ( θ ) ‖ μ \left\|\theta-\theta^{*}\right\|>\frac{2\left\|\nabla L(\theta)\right\|}{\mu} but ‖ ∇ L ​ ( θ ) ‖ 2 ≤ R ​ μ 2 \left\|\nabla L(\theta)\right\|_{2}\leq\frac{R\mu}{2} , then we have f ⁡ ( R ) < f ⁡ ( ‖ θ − θ ∗ ‖ 2 ) ≤ ‖ ∇ L ​ ( θ ) ‖ 2 f(R)<f(\left\|\theta-\theta^{*}\right\|_{2})\leq\left\|\nabla L(\theta)\right\|_{2} . On the other hand, by Assumption 4.2 and Assumption 4.1 , f ′ ​ ( t ) ≥ μ 2 f^{\prime}(t)\geq\frac{\mu}{2} for t ∈ [ 0 , R ] t\in[0,R] . Thus f ⁡ ( R ) ≥ f ⁡ ( 0 ) + ∫ t = 0 2 ​ ‖ ∇ L ​ ( θ ) ‖ μ f ′ ​ ( t ) ​ 𝑑 t = ‖ ∇ L ​ ( θ ) ‖ f(R)\geq f(0)+\int_{t=0}^{\frac{2\left\|\nabla L(\theta)\right\|}{\mu}}f^{\prime}(t)\mathrm{d}t=\left\|\nabla L(\theta)\right\| . Contradiction! ∎

###### Lemma D.4 .

For any θ ∈ ℝ d \theta\in\mathbb{R}^{d} , the following differential equation has at least one solution on interval [ 0 , 1 ] [0,1] : d ​ θ ​ ( t ) d ​ t = − ( ∇ 2 L ( θ ( t ) ) ) − 1 ∇ L ( θ ) , θ ( 0 ) = θ , \displaystyle\frac{\mathrm{d}\theta(t)}{\mathrm{d}t}=-(\nabla^{2}L(\theta(t)))^{-1}\nabla L(\theta),\quad\theta(0)=\theta, (20) and the solution satisfies that ∇ L ( θ ( t ) ) = ( 1 − t ) ∇ L ( θ ) \nabla L(\theta(t))=(1-t)\nabla L(\theta) for all t ∈ [ 0 , 1 ] t\in[0,1] and θ ⁡ ( 0 ) = θ ∗ \theta(0)=\theta^{*} .

###### Proof of Lemma D.4 .

Since ∇ 2 L \nabla^{2}L is continuous and positive definite by Assumption 4.1 , ( ∇ 2 L ) − 1 (\nabla^{2}L)^{-1} is continuous and thus the above ODE ( 45 ) has a solution over interval [ 0 , T ) [0,T) for some positive T T and we let T max T_{\max} be the largest positive number such that the solution exists (or T max = ∞ T_{\max}=\infty ). Now we claim T max ≥ 1 T_{\max}\geq 1 , otherwise ‖ θ ⁡ ( t ) − θ ∗ ‖ 2 \left\|\theta(t)-\theta^{*}\right\|_{2} must diverge to infinity when t → T max t\to T_{\max} . However, for any t ≤ 1 t\leq 1 , we have d ∇ L ( θ ( t ) ) d ​ t = − ∇ L ​ ( θ ) , \displaystyle\frac{\mathrm{d}\nabla L(\theta(t))}{\mathrm{d}t}=-\nabla L(\theta), (21) which implies that ∇ L ( θ ( t ) ) = ( 1 − t ) ∇ L ( θ ) \nabla L(\theta(t))=(1-t)\nabla L(\theta) for all t ∈ [ 0 , 1 ] t\in[0,1] . Therefore, d ​ L ​ ( θ ⁡ ( t ) ) d ​ t = − ( ∇ L ( θ ( t ) ) ) ⊤ ( ∇ 2 L ( θ ( t ) ) ) − 1 ∇ L ( θ ) = ( 1 − t ) ( ∇ L ( θ ) ) ⊤ ( ∇ 2 L ( θ ( t ) ) ) − 1 ∇ L ( θ ) ≤ 0 . \displaystyle\frac{\mathrm{d}L(\theta(t))}{\mathrm{d}t}=-(\nabla L(\theta(t)))^{\top}(\nabla^{2}L(\theta(t)))^{-1}\nabla L(\theta)=(1-t)(\nabla L(\theta))^{\top}(\nabla^{2}L(\theta(t)))^{-1}\nabla L(\theta)\leq 0. (22) Thus L ⁡ ( θ ⁡ ( t ) ) ≤ L ⁡ ( θ ⁡ ( 0 ) ) L(\theta(t))\leq L(\theta(0)) . By Lemma D.1 , we know that ‖ θ ⁡ ( t ) ‖ \|\theta(t)\| remains bounded for all t ∈ [ 0 , T max ] t\in[0,T_{\max}] , thus T max ≥ 1 T_{\max}\geq 1 . Note that θ ⁡ ( 1 ) \theta(1) has zero gradient, θ ⁡ ( 1 ) \theta(1) must be θ ∗ \theta^{*} . This completes the proof. ∎

###### Lemma D.5 .

For any θ ∈ ℝ d \theta\in\mathbb{R}^{d} satisfying (1) L ⁡ ( θ ) − min ⁡ L ≤ μ ​ R 2 16 L(\theta)-\min L\leq\frac{\mu R^{2}}{16} or (2) ‖ ∇ L ​ ( θ ) ‖ 2 ≤ R ​ μ 4 \left\|\nabla L(\theta)\right\|_{2}\leq\frac{R\mu}{4} , it holds that L ( θ ) − min L ≤ ∇ L ( θ ) ⊤ ( ∇ 2 L ( θ ) ) − 1 ∇ L ( θ ) ≤ 4 ( L ( θ ) − min L ) . \displaystyle L(\theta)-\min L\leq\nabla L(\theta)^{\top}(\nabla^{2}L(\theta))^{-1}\nabla L(\theta)\leq 4(L(\theta)-\min L). (23)

###### Proof of Lemma D.5 .