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

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

For the second term ( Equation 50 ), by Assumption 4.2 , we have that ‖ ∇ 2 L ​ ( θ + V I θ ​ w ​ ( t ) ) V I θ V I θ ⊤ ( ∇ 2 L ( θ + V I θ w ( t ) ) ) − 1 ∇ L ( θ ) ‖ 2 2 \displaystyle\left\|\sqrt{\nabla^{2}L(\theta+V_{I_{\theta}}w(t))}V_{I_{\theta}}V_{I_{\theta}}^{\top}(\nabla^{2}L(\theta+V_{I_{\theta}}w(t)))^{-1}\nabla L(\theta)\right\|_{2}^{2} (52) ≤ \displaystyle\leq 8 ‖ ∇ 2 L ​ ( θ ) V I θ V I θ ⊤ ( ∇ 2 L ( θ ) ) − 1 ∇ L ( θ ) ‖ 2 2 \displaystyle 8\left\|\sqrt{\nabla^{2}L(\theta)}V_{I_{\theta}}V_{I_{\theta}}^{\top}(\nabla^{2}L(\theta))^{-1}\nabla L(\theta)\right\|_{2}^{2} (53) = \displaystyle= 8 ∇ L ( θ ) ⊤ V I θ V I θ ⊤ ( ∇ 2 L ( θ ) ) − 1 V I θ V I θ ⊤ ∇ L ( θ ) \displaystyle 8\nabla L(\theta)^{\top}V_{I_{\theta}}V_{I_{\theta}}^{\top}(\nabla^{2}L(\theta))^{-1}V_{I_{\theta}}V_{I_{\theta}}^{\top}\nabla L(\theta) (54) = \displaystyle= 8 ∑ i ∈ I θ σ i − 1 | v i ⊤ ∇ L ( θ ) | 2 ≤ 8 Δ . \displaystyle 8\sum_{i\in I_{\theta}}\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}\leq 8\Delta. (55) Thus we conclude that ‖ V I θ c ​ d ∇ L ( θ + V I θ w ( t ) ) d ​ t ‖ 2 ≤ 4 ​ Δ ρ \left\|V_{I_{\theta}^{c}}\frac{\mathrm{d}\nabla L(\theta+V_{I_{\theta}}w(t))}{\mathrm{d}t}\right\|_{2}\leq\frac{4\Delta}{\rho} , which implies that ‖ ∇ L ( θ + V I θ ⊤ w ∗ ) ‖ 2 = ‖ V I θ c ∇ L ( θ + V I θ ⊤ w ∗ ) ‖ 2 \displaystyle\left\|\nabla L(\theta+V_{I_{\theta}}^{\top}w^{*})\right\|_{2}=\left\|V_{I_{\theta}^{c}}\nabla L(\theta+V_{I_{\theta}}^{\top}w^{*})\right\|_{2} (56) = \displaystyle= ‖ V I θ c ∇ L ( θ ) + ∫ t = 0 1 V I θ c d ∇ L ( θ + V I θ w ( t ) ) d ​ t d t ‖ 2 \displaystyle\left\|V_{I_{\theta}^{c}}\nabla L(\theta)+\int_{t=0}^{1}V_{I_{\theta}^{c}}\frac{\mathrm{d}\nabla L(\theta+V_{I_{\theta}}w(t))}{\mathrm{d}t}\mathrm{d}t\right\|_{2} (57) ≤ \displaystyle\leq ‖ V I θ c ∇ L ( θ ) ‖ 2 + ∫ t = 0 1 ‖ V I θ c d ∇ L ( θ + V I θ w ( t ) ) d ​ t ‖ 2 d t \displaystyle\left\|V_{I_{\theta}^{c}}\nabla L(\theta)\right\|_{2}+\int_{t=0}^{1}\left\|V_{I_{\theta}^{c}}\frac{\mathrm{d}\nabla L(\theta+V_{I_{\theta}}w(t))}{\mathrm{d}t}\right\|_{2}\mathrm{d}t (58) ≤ \displaystyle\leq Δ ρ + 4 ​ Δ ρ = 5 ​ Δ ρ . \displaystyle\frac{\Delta}{\rho}+\frac{4\Delta}{\rho}=\frac{5\Delta}{\rho}. (59)

Applying Lemma D.6 , we have that L ⁡ ( θ + V I θ ⊤ ​ w ∗ ) − min ⁡ L ≤ μ − 1 ​ ‖ ∇ L ​ ( θ + V I θ ⊤ ​ w ∗ ) ‖ 2 2 = 25 ​ Δ 2 ρ 2 ​ μ . \displaystyle L(\theta+V_{I_{\theta}}^{\top}w^{*})-\min L\leq\mu^{-1}\left\|\nabla L(\theta+V_{I_{\theta}}^{\top}w^{*})\right\|_{2}^{2}=\frac{25\Delta^{2}}{\rho^{2}\mu}. (60) This completes the proof. ∎

###### Lemma D.10 (Descent Lemma) .

For any η , ρ > 0 \eta,\rho>0 with η ​ ρ ≤ R / d \eta\rho\leq R/\sqrt{d} , θ ∈ ℝ d \theta\in\mathbb{R}^{d} and any eigendecomposition of ∇ 2 L ​ ( θ ) \nabla^{2}L(\theta) , where V t ​ V t ⊤ = I d V_{t}V_{t}^{\top}=I_{d} , σ t \sigma_{t} is diagonal ∇ 2 L ​ ( θ ) = V ⊤ ​ Σ ​ V \nabla^{2}L(\theta)=V^{\top}\Sigma V , define θ + ≜ θ − η V ⊤ clip ( V ( ∇ 2 L ( θ ) ) − 1 ∇ L ( θ ) , ρ ) , \displaystyle\theta_{+}\triangleq\theta-\eta V^{\top}\textup{clip}(V(\nabla^{2}L(\theta))^{-1}\nabla L(\theta),\rho), (61) it holds that L ( θ + ) − L ( θ ) ≤ − ( η − η 2 ) ∑ i = 1 d min { ρ | v i ⊤ ∇ L ( θ ) | , σ i − 1 | v i ⊤ ∇ L ( θ ) | 2 } , \displaystyle L(\theta_{+})-L(\theta)\leq-(\eta-\eta^{2})\sum_{i=1}^{d}\min\{\rho\left|v_{i}^{\top}\nabla L(\theta)\right|,\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}\}, (62) where v i v_{i} is the i i th row of matrix V V .

###### Proof of Lemma D.10 .

Let u ≜ clip ( V ( ∇ 2 L ( θ ) ) − 1 ∇ L ( θ ) , ρ ) u\triangleq\textup{clip}(V(\nabla^{2}L(\theta))^{-1}\nabla L(\theta),\rho) . By the definition of clip operation, we know that ‖ V ⊤ ​ u ‖ 2 = ‖ u ‖ 2 ≤ d ​ ρ \left\|V^{\top}u\right\|_{2}=\left\|u\right\|_{2}\leq\sqrt{d}\rho . Thus we have ‖ θ + − θ ‖ = η ​ ‖ V ⊤ ​ u ‖ 2 ≤ η ​ ρ ​ d \left\|\theta_{+}-\theta\right\|=\eta\left\|V^{\top}u\right\|_{2}\leq\eta\rho\sqrt{d} . Define f ⁡ ( t ) = L ⁡ ( t ​ θ + + ( 1 − t ) ​ θ ) f(t)=L(t\theta_{+}+(1-t)\theta) . By Assumption 4.2 , we know that f ′′ ​ ( t ) ≤ 2 ​ f ′′ ​ ( 0 ) f^{\prime\prime}(t)\leq 2f^{\prime\prime}(0) for all t ∈ [ 0 , 1 ] t\in[0,1] and thus f ⁡ ( 1 ) = f ⁡ ( 0 ) + f ′ ​ ( 0 ) + ∫ s = 0 1 ∫ t = 0 s f ′′ ​ ( s ) ​ 𝑑 s ​ 𝑑 t ≤ f ⁡ ( 0 ) + f ′ ​ ( 0 ) + f ′′ ​ ( 0 ) . \displaystyle f(1)=f(0)+f^{\prime}(0)+\int_{s=0}^{1}\int_{t=0}^{s}f^{\prime\prime}(s)\mathrm{d}s\mathrm{d}t\leq f(0)+f^{\prime}(0)+f^{\prime\prime}(0). (63) It remains to show that 1. f ′ ( 0 ) = − η ∑ i = 1 d min { ρ | v i ⊤ ∇ L ( θ ) | , σ i − 1 | v i ⊤ ∇ L ( θ ) | 2 } f^{\prime}(0)=-\eta\sum_{i=1}^{d}\min\{\rho\left|v_{i}^{\top}\nabla L(\theta)\right|,\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}\} ;

2. f ′′ ( 0 ) ≤ η 2 ∑ i = 1 d min { ρ | v i ⊤ ∇ L ( θ ) | , σ i − 1 | v i ⊤ ∇ L ( θ ) | 2 } f^{\prime\prime}(0)\leq\eta^{2}\sum_{i=1}^{d}\min\{\rho\left|v_{i}^{\top}\nabla L(\theta)\right|,\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}\} ;

First, by chain rule, we have f ′ ( 0 ) = ⟨ ∇ L ( θ ) , − η V ⊤ u ⟩ = ⟨ V ∇ L ( θ ) , − η u ⟩ = − η ⟨ V ∇ L ( θ ) , clip ( Σ − 1 V ∇ L ( θ ) , ρ ) ⟩ = − η ∑ i = 1 d min { ρ | v i ⊤ ∇ L ( θ ) | , σ i − 1 | v i ⊤ ∇ L ( θ ) | 2 } f^{\prime}(0)=\left\langle\nabla L(\theta),-\eta V^{\top}u\right\rangle=\left\langle V\nabla L(\theta),-\eta u\right\rangle=-\eta\left\langle V\nabla L(\theta),\textup{clip}(\Sigma^{-1}V\nabla L(\theta),\rho)\right\rangle=-\eta\sum_{i=1}^{d}\min\{\rho\left|v_{i}^{\top}\nabla L(\theta)\right|,\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}\} . Second, again by chain rule, we have f ′′ ​ ( 0 ) = η 2 ​ ⟨ V ⊤ ​ u , ∇ 2 L ​ ( θ ) ​ V ⊤ ​ u ⟩ = η 2 ​ ⟨ u , Σ ​ u ⟩ = ∑ i = 1 d | u i | 2 ​ σ i f^{\prime\prime}(0)=\eta^{2}\left\langle V^{\top}u,\nabla^{2}L(\theta)V^{\top}u\right\rangle=\eta^{2}\left\langle u,\Sigma u\right\rangle=\sum_{i=1}^{d}\left|u_{i}\right|^{2}\sigma_{i} . Note that by definition | u i | = min { | v i ⊤ ∇ L ( θ ) | / σ i , ρ } \left|u_{i}\right|=\min\{\left|v_{i}^{\top}\nabla L(\theta)\right|/\sigma_{i},\rho\} , we have | u i | 2 σ i ≤ min { | v i ⊤ ∇ L ( θ ) | / σ i , ρ } ⋅ | v i ⊤ ∇ L ( θ ) | / σ i ⋅ σ i = min { | v i ⊤ ∇ L ( θ ) | 2 / σ i , ρ | v i ⊤ ∇ L ( θ ) | } \left|u_{i}\right|^{2}\sigma_{i}\leq\min\{\left|v_{i}^{\top}\nabla L(\theta)\right|/\sigma_{i},\rho\}\cdot\left|v_{i}^{\top}\nabla L(\theta)\right|/\sigma_{i}\cdot\sigma_{i}=\min\{\left|v_{i}^{\top}\nabla L(\theta)\right|^{2}/\sigma_{i},\rho\left|v_{i}^{\top}\nabla L(\theta)\right|\} , which completes the proof. ∎

###### Lemma D.11 .

If η ​ ρ ≤ R / d \eta\rho\leq R/\sqrt{d} and for some T ∈ ℕ T\in\mathbb{N} , L ⁡ ( θ T ) − min ⁡ L ≤ μ ​ ρ 2 8 L(\theta_{T})-\min L\leq\frac{\mu\rho^{2}}{8} , then if holds that for all t ≥ T t\geq T , 1. θ t + 1 = θ t − η ( ∇ 2 L ( θ t ) ) − 1 ∇ L ( θ t ) \theta_{t+1}=\theta_{t}-\eta(\nabla^{2}L(\theta_{t}))^{-1}\nabla L(\theta_{t}) ;

2. L ⁡ ( θ t ) − min ⁡ L ≤ ( 1 − η ⁡ ( 1 − η ) ) t − T ​ ( L ⁡ ( θ T ) − min ⁡ L ) L(\theta_{t})-\min L\leq(1-\eta(1-\eta))^{t-T}(L(\theta_{T})-\min L) .

###### Proof of Lemma D.11 .

First by Lemma D.10 , we have for all t ≥ T t\geq T , Ł ⁡ ( θ t ) − min ⁡ L ≤ L ⁡ ( θ T ) − min ⁡ L ≤ μ ​ ρ 2 8 \L(\theta_{t})-\min L\leq L(\theta_{T})-\min L\leq\frac{\mu\rho^{2}}{8} , therefore by Lemma D.7 , we have ‖ ( ∇ 2 L ( θ t ) ) − 1 ∇ L ( θ t ) ‖ 2 ≤ ρ \left\|(\nabla^{2}L(\theta_{t}))^{-1}\nabla L(\theta_{t})\right\|_{2}\leq\rho for all t ≥ T t\geq T , which implies clipping will not happen. This completes the proof of the first claim.

For the second claim, by Lemmas D.10 and D.5 , we have that L ⁡ ( θ t + 1 ) − L ⁡ ( θ t ) ≤ \displaystyle L(\theta_{t+1})-L(\theta_{t})\leq − ( η − η 2 ) ∑ i = 1 d σ i − 1 | v i ⊤ ∇ L ( θ t ) | 2 \displaystyle-(\eta-\eta^{2})\sum_{i=1}^{d}\sigma_{i}^{-1}\left|v_{i}^{\top}\nabla L(\theta_{t})\right|^{2} (64) = \displaystyle= − ( η − η 2 ) ∇ L ( θ t ) ( ∇ 2 L ( θ t ) ) − 1 ∇ L ( θ t ) \displaystyle-(\eta-\eta^{2})\nabla L(\theta_{t})(\nabla^{2}L(\theta_{t}))^{-1}\nabla L(\theta_{t}) (65) ≤ \displaystyle\leq − η ⁡ ( 1 − η ) ​ ( L ⁡ ( θ t ) − min ⁡ L ) , \displaystyle-\eta(1-\eta)(L(\theta_{t})-\min L), (66) which completes the proof. ∎

### D.1 Lower bound for SignGD on 2-dimensional quadratic loss

Define L μ , β : ℝ 2 → ℝ L_{\mu,\beta}:\mathbb{R}^{2}\to\mathbb{R} as a quadratic function with parameter μ , β \mu,\beta as L μ , β ​ ( θ ) ≜ μ 2 ​ θ [ 1 ] 2 + β 2 ​ θ [ 2 ] 2 L_{\mu,\beta}(\theta)\triangleq\frac{\mu}{2}\theta_{[1]}^{2}+\frac{\beta}{2}\theta_{[2]}^{2} . We have the following lower bound, which shows signGD’s convergence rate has to depend on the condition number β / μ \beta/\mu .

###### Theorem D.12 .

For any μ , β , Δ , ϵ > 0 \mu,\beta,\Delta,\epsilon>0 , suppose there exist a learning rate η \eta and a time T T such that for all θ 0 \theta_{0} satisfying that L μ , β ​ ( θ 0 ) ≤ Δ L_{\mu,\beta}(\theta_{0})\leq\Delta , signGD reaches loss at most ϵ \epsilon at step T − 1 T-1 and T T (in the sense that L μ , β ​ ( θ T ) ≤ ϵ L_{\mu,\beta}(\theta_{T})\leq\epsilon and L μ , β ​ ( θ T − 1 ) ≤ ϵ L_{\mu,\beta}(\theta_{T-1})\leq\epsilon ). Then, T T must satisfy T ≥ 1 2 ​ ( Δ ϵ − 2 ) ​ β μ T\geq\frac{1}{2}(\sqrt{\frac{\Delta}{\epsilon}}-\sqrt{2})\sqrt{\frac{\beta}{\mu}} .

###### Proof of Theorem D.12 .

We consider two initialization: θ 0 = ( 0 , 2 ​ Δ β ) \theta_{0}=(0,\sqrt{\frac{2\Delta}{\beta}}) and θ 0 ′ = ( 2 ​ Δ μ , 0 ) \theta^{\prime}_{0}=(\sqrt{\frac{2\Delta}{\mu}},0) , and let θ t \theta_{t} and θ t ′ \theta_{t}^{\prime} be the iterates under the two initializations. For each coordinate i ∈ { 1 , 2 } i\in\{1,2\} , because | ( θ t ) [ i ] − ( θ t + 1 ) [ i ] | = η |(\theta_{t})_{[i]}-(\theta_{t+1})_{[i]}|=\eta , we have that | ( θ t ) [ i ] | + | ( θ t + 1 ) [ i ] | ≥ η |(\theta_{t})_{[i]}|+|(\theta_{t+1})_{[i]}|\geq\eta . Thus 2 ​ ϵ ≥ L μ , β ​ ( θ T ) + L μ , β ​ ( θ T − 1 ) ≥ β 2 ​ ( ( θ T ) [ 2 ] 2 + ( θ T − 1 ) [ 2 ] 2 ) ≥ β ​ η 2 4 2\epsilon\geq L_{\mu,\beta}(\theta_{T})+L_{\mu,\beta}(\theta_{T-1})\geq\frac{\beta}{2}((\theta_{T})_{[2]}^{2}+(\theta_{T-1})_{[2]}^{2})\geq\frac{\beta\eta^{2}}{4} , which implies η ≤ 8 ​ ϵ β \eta\leq\sqrt{\frac{8\epsilon}{\beta}} .

The fact that L μ , β ​ ( θ T ′ ) + L μ , β ​ ( θ T − 1 ′ ) ≤ 2 ​ ϵ L_{\mu,\beta}(\theta_{T}^{\prime})+L_{\mu,\beta}(\theta_{T-1}^{\prime})\leq 2\epsilon implies ( θ T ′ ) [ 1 ] ≤ 4 ​ ϵ μ (\theta_{T}^{\prime})_{[1]}\leq\sqrt{\frac{4\epsilon}{\mu}} . Because SignGD can only move each coordinate by η \eta at most, we have ( T − 1 ) ​ η ≥ 2 ​ Δ / μ − 4 ​ ϵ μ (T-1)\eta\geq\sqrt{2\Delta/\mu}-\sqrt{\frac{4\epsilon}{\mu}} . Using the fact that η ≤ 8 ​ ϵ β \eta\leq\sqrt{\frac{8\epsilon}{\beta}} , we have that 2 ​ ( T − 1 ) ≥ ( Δ ϵ − 2 ) ​ β μ 2(T-1)\geq(\sqrt{\frac{\Delta}{\epsilon}}-\sqrt{2})\sqrt{\frac{\beta}{\mu}} , which completes the proof. ∎

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .