Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training
Authors: Hong Liu
Zhiyuan Li
David Leo Wright Hall
Percy Liang
Tengyu Ma
Venue: ICLR 2024 poster (2024)
Keywords: large language models
pretraining
optimization in deep learning
Primary area: optimization
TL;DR: We design a stochastic second-order optimizer for language model pre-training which is significantly faster than Adam.

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, 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.

Published as a conference paper at ICLR 2024
SOPHIA: A SCALABLE STOCHASTIC SECOND-ORDER OPTIMIZER FOR LANGUAGE MODEL PRE-TRAINING
Hong Liu
Zhiyuan Li
David Hall
Perey Liang
Tengyu Ma
Stanford University
{hliu99, zhiyuanli, dlwh, pliang, tengyuma (cs. stanford,
ABSTRACT
Given the massive cost of language model pre-training, a non-trivial improve- ment of the optimization algorithm would lead to 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 in- cur too much per-step overhead. In this paper, propose Sophia, Second-order Clipped Stochastic Optimization, a simple scalable second-order optimizer that uses fight-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-wist 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 Hes- sian 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 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.
INTRODUCTION
Language models (LLMs) have gained phenomenal capabilities their scale grows (Radford al., 2019; Kaplan et al., 2020; Brown et al., 2020; Zhang et al., 2022b; Touvron et al., 2023; OpenAT, 2023). However, pre-training LLMs is incredibly time-consuming due to the massive datasets and model sizes-hundreds of thousands of updates to parameters are required. For example, PaLM was trained for two months on 6144 TPUs, which costed 10 million dollars (Chowdhery et al., 2022).
Tre-training efficiency is thus 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 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 sradieent-bseed 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 weihht options whose overhead can be offset by the speed-up in the number of iterations. On the other hand, Chen al. (2023) automatically search among the liight 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, Second-order Clipped Stochastic Optimization, a light-weight second- order optimizer that uses an inexpensive stochastic estimate of the diagonal of the Hessian as 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 (a)&(b)). We also note that comparing the run-time to achieve the same loss is the correct way to compare the speed of optimizers for LLMs; see Section 3.2 for more details.

---

Published as a conference paper at ICLR 2024
(N) GPT2 Medium (25SH)
3ct GGF12 Lamg (720M)
(c) CPT Neox, 1.38
Sca
/daml)
Adben!
Sapfea
S
Scphlaii 150K
Sophia-d
52.
Sapria 100x
244
2.7
2.4
2.5
Spend.p
56.75
1166 120.35
2
10Mc 216K
<04K
310
4:00 G20
Compuse aeattots
"AATIC Stggs
Namber Sluse
Model Ste
Figure 1: Sophia achieves significant speedup over AdamW in GPT-2 pre-trained on Open WeenWe and GPT NeoX pre-trained on the Pile. (a) (b) (c) Comparison the number of steps needed to achieve the same validation loss on (a) GPT-2-medium (355M), (b) G-ree-large (770M) and (c) GPT NeoX 1.5B. Across all model sizes, Sophia needs 50% less time to reach the same validation loss as AdamW. (d) Validation losses of models with different sizes pre-trained. The gap between Sophia and AdamW gets larger as models size grows. Notably, a 5400-parameee model pre-trained with Sophia for 100K steps has the same loss as a 770M- parameter model pre-trained with AdamW for 100K steps. See Section 3 for details and more results.
Algorithm 1 Hutchinson( (0)
Algorithm 2 -eeess-eettto---rreee (GNB)(0)
1: Input: parameter 6.
1: Input: parameter 0.
2: Compute mini batch loss L(0).
2: Draw mini-batch of input {2
3: Draw u from (0 Ia).
3: Compute logits on batch: (f() [b)
4: return Φ (((VL(0), u))-
4: Sample in. ~ softmax(f(C [(0 xn)), 4b E [B]
5: Calculate 9 VLIBB ((((0,)))) 96)).
6: return B g C 9.
SinnGD
Adam
Algorithm 3 Sophia
Newton
Sophia
1: Input: 01, learning rate {n: hyperpa- rameters X, B,, .2, C, and estimator choice
Estimator E (Hutchinson,GNI)
2: Set mo = 0, vo = 0, h =0
3: for f = 1 to T do
e>
4: Compute minibach loss (0.).
5: Compute ga = LL (0.).
The motivating toy example. Ouj the sharp dimension and fp]] is the flat di- mension. GD's learning rate is limited by the sharpness in 0,. It makes slow progress along 0/-)- Adam and SignGD bounce along Onj while making slow progress along 0---- Newton's 's method converges to saddle point. Sophia makes fast progress in [[] and d/lj and converges to the minimum with few steps.
Figure 2:
6: my = - + (1 3))
7:
if t mod k then
8:
Compute h, Estimator(O.).
9:
ht B2h4 k + (1 - 82)
10:
else
11:
ht = h4-1
12:
O₂ = 02 11 (weight decay)
13:
0c+1 = B. - m clip(m:/ max <), 1)
Moreover, the scaling law based on model size from 125M to 770M 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 540a-marrmeter model with 100K steps gives the same validation loss as Adam
model with 100K steps.
Sophia estimates the diagonal entries of the Hessian of the loss using a mini-batch of examples ev- ery l step (with k= = 10 in practice). We consider two options for diagonal Hessian estimators: (a) an unbiased estimator that uses a Hessian vector product, 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, subse- quently clipped by a scalar. (All operations are clement-wise.) See Algorithm 3 for the pseudo-code.
Sophia has a more aggressive pre-conditioner-in it applies stronger penalization to updates in sharp dimensions (where the Hessian is large) than the flat dimensions (where the Hessian is small), ensur- ing a uniform loss decrease across all parameter dimensions. contrast, Adam's 's updates are mostly uniform across all parameter dimensions, leading to 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 or wall-clock time.
Sophia's is clipping mechanism controls the worst-case size of the updates in all directions, safeguard- ing against the negative impact of inaccurate Hessian estimates, rapid Hessian changes, and non-
2

---

Published as a conference paper at ICLR 2024
convex landscape (with which 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 steps). In contrast, prior second-order methods often update Hessian estimates every step.
METHOD
We instantiate gradient descent and Adam on a simplified 2D problem to motivate the use of second- order information and per-coordinate clipping in Section 2.1. We present Sophia in Section 2.2, with pseudo-code in Algorithm 3. We introduce two estimators of diagonal Hessian in Section 2.3.
2.1 MOTIVATIONS
Heterogeneous curvatures. The loss functions in deep learning often have different curvatures across different parameter dimensions (Sagun et al., 2016; Ghorbani et al., 2019; Zhang et al., 2020; Yao et al., 2020). We demonstrate the limitations of Adam and GD on heterogeneous landscapes in a two-dimensional loss function [(([u): 012)) = I. (9n)) + L (0())) where In is much sharper
than L2. We plot the loss landscape of [(] 022)) in Figure 2.' For simplicity, we discuss GD and
deterministic versions of Adam. GD's update in this setting is:
puj 11] - 17 (1 and 62] - 0-11 - ")
L2(()p))
(1)
A common simplification of Adam that is more amenable 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 al., 2012) and Adam. Without using the EMA (for both the gradient and second moments of the gradient), Adam's 's update is simplified r7- [L())/VL(() = 7) (oggnVI((0) which is called SignGD:
11] + 01] - 1-- (ig)(L((011)) and 012) + 1|2 17 sign(L2200a))
(2)
Limitations of GD and SignGD (Adam). The optimal learning rate of GD should be proportional to the inverse of the curvature. Let h1 and h2 be the curvatures of L1 and L2 at the local minimum (and thus h1 >> h2). The optimal learning rate for 01j in (1) is x 1/h1. which is much smaller than the optimal learning rate that 022 needs (x 1/h2). The largest shared learning rate can only be 1//a; consequently, the convergence 022 is slow as demonstrated in the brown curve of Figure 2.
The update size of SignGD equals learning rate v in all dimensions. The same update size translates to less progress in the flat direction than in the sharp direction. As shown in the yellow curve of Figure 2, the progress of SignGD in 0/2) is slow because each step only decreases the loss (0)) slightly. On the other hand, along (1]. the iterate quickly goes to the valley in the first three steps and then starts to bounce. To fully converge in the sharp 11) the learning rate r) needs to decay to 0, which will exacerbate the slow convergence the flat [22]. The trajectory of Adam in this example is similar to SignGD, which is plotted as the red curve in Figure 2.
The behavior of SignGD and Adam calls for more aggressive pee-condtionnng--sann dimensions should have relatively smaller updates than flat dimensions so that the decrease of loss is equalized in all dimensions. As suggested by literature on second-order optimization (Boyd & Vandenberghe, 2004) for convex functions, the optimal pre-conditioner should be the Hessian, which captures the curvature of each dimension; the update is the gradient divided by the Hessian of each dimension:
f1g + U| 17 E2 (0t))/h and 1/] + 0p] - 17 L2(0p))/22
Limitations of Newton's method. For non-convex functions. Newton's 's method may not converge to a minimum. In the blue curve of Figure 2, Newton's 's method quickly converges to a saddle point. 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 pre-coordinate clipping 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:
pnj + 81j - clip [41 (mxu p) and 0 + 0/2] - 7 clip( L 12)x2//mxx( p), (4)
where p a constant control the update size, e is a very small constant (e.g., 1e-12), which avoids dividing by 0. When the curvature of some dimension is changing rapidly negative and thus the second-order information is misleading, the clipping mechanism kicks in and the optimizer defaults
CConeretely, in Figure 2, (0m)) = 8(0() - 1)² + 20 + 1) and L2(0(2) = 1/2(0/2) - -4)².
3

---

Published as a conference paper at ICLR 2024
to SignGD. Numerous prior methods such as trust region (Conn et al., 2000), backtracking line search (Boyd & Vandenberghe, 2004), and cubic regularization (Nesterov & Polyak, 2006) tackle the same issue, but the clipping mechanism is 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 the saddle point. Then, in the convex valley, it converges to the global minimum with few steps. Compared with SignGD and Adam, it makes much faster progress in the flat dimension A/z). while avoiding boucing in the sharp dimension [1]-
2.2 SOPHIA: SECOND-ORDER CLIPPED STOChASTIC OPTiMIZaTION
Section 2.1 demonstrates Adam does not sufficiently adapt to the heterogeneous curvatures. On the other hand, Newton's method has a pre-conditioner optimal for convex functions, but is vulnera- ble to negative curvature and rapid change of Hessian. With these insights, we design a new opti- mizer, Sophia, which is more adaptive to heterogeneous curvatures than Adam, more resistant non- convexity and rapid change of Hessian than Newton's 's method, and uses a low-cost pre-conditioner.
We use 0, to denote the parameter at time step At each step, we sample mini-batch from the data distribution and calculate the mini-batch loss, denoted by Lz(Or). We denote by gr. the gradient of (O.).. Gt (O.). Let m, be the EMA of gradients, m4 + - + (1 11 (1))) ).)..
EMA of diagonal Hessian estimates. Sophia uses a diagonal Hessian pre-conditioner, which di- rectly adjusts the update size of different parameters according to their curvatures. We will present two efficient estimators of diagonal Hessian in Section 2.3. To mitigate the overhead, we only esti- mate the Hessian every l steps (k= = 10 in practice). At time step t with mod k = 1, the estimator returns an estimate hu of the diagonal of the Hessian of the mini-batch loss.
Similar to the gradient of the mini-batch loss function, the estimated diagonal Hessian can also have large noise. Inspired by the EMA of moments of gradients in Adam, we also denoise the diagonal Hessian estimates with EMA across iterations. We update the EMA every k steps, resulting in the following update rule for the diagonal Hessian estimate:
he 32he-k + (1 (1- (1-B2)he if mod k = 1; else he -he1-
(5)
Per-coordinate clipping. As discussed in Section 2.1. on nonconvex functions, vanilla Newton's method, which uses Hessian as the pre-conditioner, may converge to local maxima instead of lo- cal minima. In addition, the inaccuracy of Hessian estimates and the change of Hessian along the trajectory can make the second-ordes information unreliable. To this end, we (1) only consider the positive entries of the diagonal Hessian and (2) introduce per-coordinate clipping to the update. For a clipping threshold p > O, let the clipping function be clip(z, p) = max (min 2, p), -p} where all operations applied coordinate-wise The update rule is written as:
Oc+1 - 0 Ti- clip(m./ max(7 h, 1),
(6)
where e >> 0 is a small constant to avoid dividing by 0, and controls the fraction of clipped entries. We present the pseudo-code of the Sophia Algorithm 3.
If he(i] < O. the corresponding entry in the pre-conditioned gradient my max///xxx---// he hx] = m/li//e large and has the same sign as my[0) and thus r clip(m,(j)] [i]/ max(r [ c).1) = n sign(m,.i), [i]). which is the same as stochastic momentum SignSGD. In other words, Sophia uses stochastic Sign Momentum GD as a backup when the Hessian is negative (or mistakenly estimated to be negative.) Also note that clipping controls the worst-case update size in all parameter dimensions to be at most p. which also improves stability. Moreover, because for many parameter dimensions, the clipping is not activated and the update is automatically adjusted, our worst-case update size " can be chosen to be larger than the worst update size v) in stochastic Sign Momentum GD.
Several previous works (Becker& & Le Cun, 1988; Chapelle et al., 2011: 1; Schaul et al., 2013; Yao al., 2021) use diagonal Hessian as a pre-conditioner in optimizers for training neural networks. However, they use more frequent Hessian estimations, which leads significant per-step computation over- head, most likely because of the lack of the clipping mechanism that safeguards against inaccurate and changing Hessian. In general, to the best of our knowledge, there has not been previous reports that showed second-ord order optimizers achieve speed-up on decoder-only large language models in wall-clock time or total compute (see more related work and discussions in Section A).
4

---

Published as a conference paper at ICLR 2024
2.3 DIAGONAL. HESSIAN ESTIMATORS
We introduce two diagonal Hessian estimators, both of which have memory and run-time costs sim- ilar to computing a gradient (up to constant factors).
Option I: Hutchinson's unbiased estimator. For any loss function {()) parameters 0 E RR the Hutchinson's is estimator (Hutchinson, 1989; KoostaKKooassnni & Ascher, 2015; Yao et al., 2021) first draws и E IRd from the spherical Gaussian distribution (0, Ta). and then outputs h = a @ (V2((0)u), where @ denotes the element product, and VPE(O)u is the product of the Hessian with the vector u. The Hutchinson's 's estimator is an unbiased estimator for the diagonal of the Hessian, because EBA] - (((()) It only requires a Hessian-vector product (i.e., (²((O)u), which have efficient implementaaionns in PyTorch and JAX, instead of the full Hessian matrix.
Option 2: Gauss-Newton-Bartlett (GNB) estimator. Suppose f(0, (r) is a loss function on an example (ar, y) of the form E(O, (ar, u)) (f (0, y) where fee is cross-entropy and f(0, RV is the logits, and V is the number of items/classes in a multi-class classification problem (e.g., the vocabulary size in LLMs). First, the Hessian of ε(0, (z, y)) (w.r.1 to variable 0) has the well-known Gauss-Newton (GN) decomposition (Ortega & Rheinboldt, 2000; Schraudolph, 2002).
V3 E(0) (0) S (of(0,
(0,
where 20) is the Jacobian of f w.r. to viewed as a matrix in RR S = (t [(1,)))) E RVV is the second-order derivatives of the loss w.r.t the logits, g = (=)
(t, (0).)))) E RRV the first-order derivatives of the loss w.r.t the logits, and Joo f x) is the second-order derivatives of the multi-variate function f(0, z)) w.r.t O viewed from RV to dddd In neural networks, past works found that the second term f (0, []][] in (11) is often relative smaller than the first term (0, z) S fo,( x)² (Sankar et al., 2021), which is often referred to as the Gauss-Newton (GN) matrix (Dennis Jr & Schnabel, 1996; Ortega & Rheinboldt, 2000; Schraudolph, 2002; Chen, 2011) and used as pre-conditioners in second-order optimizers (Botev et al., 2017; Martens, 2020; Gargiani et al., 2020). Following these works, we build unbiased estimator for the diagonal of the GN matrix, as a biased estimator of the Hessian diagonal.
linear map
First note that S only depends f(0,z) but not y. even though the loss depends on y.2 Thus, S S [ (O) ((=(0e)) Because (oe(t, v) is the negative log-probabilityy of the probabilistic model defined by the categorical distribution Cat(t) with parameter t. by Bartlett's second identity (Bartlett, 1953), we have that, S = Eg~cu(t)) (oe 0)/
(t, f00.)) for any ÿ E (1.....V), which implies that
Eg-Ca((e [al G) (Otcc 1)ttt where the first equality holds for = f(0, x) and the second equality holds for all by Bartlett's second identity. Therefore, the GN matrix satisfies
S-J =
E 1..00..0000..(.,00000000........ g~CC((1)
JoS(0,x)
VVol (oo_
(8)
j~Cal(t)
Vp x)) ()) Hence, the quantity (f), 2) @ (f 9) is an unbiased estimator of the GN matrix for the Hessian of a oxe-example loss ((f(0.z).))
implying diag(Jof(0, S Jof(0,x) Eyncautt) [f 9)
Manc-batch version. Given mini-batch of inputs { Hb) The most natural way to build an estimator for the diagonal of the GN matrix for the Hessian of the mini-batch loss is using
Σ B (f0, (6) @ (f (z).
where 0b's 's are labels sampled from the model on inputs ss'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
LLet x)), softmax([(0, E R the prohability vector obtained from softmax of logits. S diagonal(p(E) x)) - -(0, x)0,z)p(0,4 2)² where diagonal(z) is the matrix with vector z residing on the diagonal.
5

---

Published as a conference paper at ICLR 2024
to the individual ones). Fortunately, by Bartlett's s first identity (Bartlett, 1953) (which generally holds
for negative log- log-likelihood loss of probabilistic models), we have: Vb, Eg. Vl (f (0, Ub) = 0.
Let [(0) B Ei= loolf(e 2b), 96) be the mini- batch loss on the sampled labels (as opposed to the original labels). Öbserving that yo are independent with each other, we have
Bgg
B o = Eijs
z.). O b=1
Σ b (lol 00)
B
b=1 (co()(0,2)):00)) 00) . (lo 900)
(10)
The RHS of (16) is the same as the expectation of (14), which by (13) also equals the diagonal of the GN matrix for the mini-batch loss. Hence, we use B - VoL(0) @ Vol(0) as the estimator. To the best of our knowledge, Wei et al. (2020) first used this estimator. Given the use Bartlett's identities that are central to the estimator, we call it Gauss-Newton Nrettnn-Bartlett (GNB) estimator.
GNB estimator for exponential family. If y is drawn from an exponential family p(y: where the nat- ural parameter 1) is set to be (0) and the loss function (0, .r), y) is the negative log-likelihood loss for the corresponding probabilistic distribution, then all the derivations still hold because (1) S still only depends on f(0, x) but not y, and (2) Bartlett's 's identities still hold.
GNB estimator for squared loss. When y, f(0,2)) E R and (O,00)), - ((((0,)) - y)², the S matrix is identity, and thus one can simply use Jo (0, (0, as the estimator.
Comparisons of Hessian estimators. The Hutchinson's estimator does not assume any structure of the loss, but requires bressian-vector products. The GNB estimator only estimates the GN term and always gives 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).
EXPERIMENTS
We name the algorithm using Hutchinson's and GNB estimator Sophia-H and Sophia-G, respec- tively. We evaluate Sophia on GPT of sizes ranging from 125M to 6.6B. Results indicate that Sophia achieves the same or smaller validation loss than AdamW (Loshchilov & Hutter, 2017) in 50% less number of steps, total compute, and wall-clock time across different model sizes.
3.1 EXPERIMENTAL. SETUP
Language modeling. We train autoregressive models on Open WebText (Gokaslan & Cohen, 2019) and the Pile (Gao et al., 2020) from scratch. Following standard protocol, we 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 C.2.
Baselines. We compare Sophia with Adam with decoupled weight decay (AdamW) (Loshchilov & Hutter, 2017), 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), a first-order adaptive optimizer discovered by symbolic search. For the 30M model, all hyperparameters are tuned with grid search. For other models, all hyperparneters but the peak learning rate are configured 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 search for the largest possible peak learning rate such that the training does not blow up, and ensure 1.25 times the chosen learning rate will lead to a blow-up. For Adam W we found the well-established practice (8) 0.9 and B2 0.95) works consistently better than other choices. For Lion, we use 3 = 0.95 and B2 = 0.98 following Chen et al. (2023). For AdaHessian, we found 33 = 0.92 and 32 = 0.99 works the best. Details on hyperparameter tuning are deferred to Section C.I.
Implementation. We use batch size 480 for GPT-2 and 2048 for GPT NeoX. We use cosine LR schedule with the final LR equaling 0.05 times the peak LR with a fixed 2k steps of LR warm-up following Rae et al. (2021). We use standard gradient clipping (by norm) threshold 1.0. For Sophia, we use 31 = 0.96, 82- 0.99, e = 1e- 12 and update diagonal Hessian every 10 steps. For Sophia-H, we use a subset of 32 examples from the minibatch to calculate the diagonal Hessian. For Sophia-G, we use a subset of 240 examples from the minibatch to calculate the diagonal Gauss-Newton. We
6

---

Published as a conference paper at ICLR 2024
Leaming Rate Schedu
16) CPT2 Nedium (355
[6) GPT Carge 1770N)
240 GPT Nenk, 1.58
100K diese
Adamiii 1CoK
2.2
Adlanfit 200K
700k stas
Aduntii zieK
2.7
Adannty 400k.
400K 90035
Jon 201K
55.0
Stphaa-t 2038
2.44
Sophil G, 2008
3.0
Sochis X01K
Sectia 100K
E2.
244
Soptis-G, 6, 1058
2.
2.7
2.4
Spocsap
XBC
100K
Marberr of Stept
Comoute DPe
Stepi
Number Steps
Figure 3: Comparison of numbers of steps to reach the same validation loss on WebTee. (a) Learning rate schedules. (b) GPT-2 Medium (355M). (c) GPT-2 Large (770M). (d) GPT NeoX 1.5B.
Adammi
AdamWW
Adomty
Adamot
Adann Supt
Vatroton leee 12 1.
Len
3.3
Linn
33.0
Seprea-
Sephad
2.
52.9
22.
2.8
2.7
2.4
2.5
2.2
226 of Stode
LooK
226 516
kumber 2005
130K
Marbbrr Stipe
518
130c
Mimberr of Strps
2004
Mamber of Shipe
Figure 4: Validation loss. (a) GPT-2 Small (125M). AdamW: 2.921, Lion: 2.924, Sophia-H: 2.901, Sophia-G: 2.875 (b) GPT-2 Medium (355M). Adam: 2.691. Lion: 2.678, Sophia-HI 2.645, Sophia-G: 2.627. (e) GPT-
2 Large (770M). AdamW: 2.613, Sophia-H: 2.554, Sophia-G: 2.524. (d) GPT NeoX 1.5B, AdamW: 2.250,
Sophia- G: 2.218.(e) GPT NeoX 6.6B. AdamW: 1.992, Sophia-Gi G: 1.969.
implement the algorithms in PyTorcl (Paszke et al., 2019) and JAX (Bradbury et al., 2018) and train all the models in b 125M and 355M models are trained on A5000 GPUs, while the 770M models trained on A100 GPUs. We use TPU v3-128 slice train 1.5B and 6.6B GPT NeoX.
Hyperparamter tuning strategy. We refer to Section C.1 for the details on hyperparameters and only discuss two key hyperparameters, Y and the peak learning rate 7) 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 7, 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 7 to make the proportion of coordinates where the update is not clipped (i.e., ht//ax/y-aa he, e} < 1) in the range of models with the same architecture and data but different number of parameters. We use 7 = 0.01 for Sophia-H and 7 = 0.05 for Sophia-G G in this paper.
The same y likely can be transferred to
Suppose we already find a suitable 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. One correct (and preferred) way of claim- ing optimizer S is 2x faster than optimizer A is comparing the following two experiments (for a variety of T)s) (1) running optimizer A (e.g. Adam) with T steps, with the optimal learning rate and learning rate schedule (tuned for running for I steps) (2) running optimizer S (c.g.. Sophia) with T/2 steps, with any learning rate schedule, If Experiment 2 achieves loss that is smaller than or equal to the loss of Experiment 1, then we say optimizer S is 2x faster than optimizer A.
Note that modern learning rate schedulers such cosine learning rate (Loshchilov & Hutter, 2016) are highly sensitive to pre-specific total number of steps. Figure 3 (a) shows that with the same peak learning rate, the learning rate run with T/2 steps decays faster than that with T steps. Moreover, the loss of the 7/2-steps run decays faster initially than the 7-steps run but ends up larger. The latter is also not a continuation of the former. Thus, in the proposed comparison above, we insist that Experiment 2 has loss smaller than or equal to Experiment I without any approximations, because even "Adam with T/2 steps" can possibly achieve a loss similar to (but slightly worse than) "Adam with T steps with the same peak learning rate". See Figure 3db)
Technical details. Following the methodology above, we train baselines and Sophia for 100K, 200K, 400K. 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 Super- GLUE (Wang al., 2019). We average the results of 5 prompts (Section C.3).
3.3 RESULTS
Figure 4 illustrates the validation loss curve (log perplexity) on Web or the Pile with the same number of steps. Sophia-H and Sophia-G consistently achieve better validation loss than
7

---

Published as a conference paper at ICLR 2024
Evaluation Medium
GPT: Large
Kammit 100K
72
Acammi 100C
Atamn,, 200k
Lon 100K
Stphia 103K
Sophis 6, 230K
Hs 105K
Siptra- L10K
45
Alonii acee.
G 100K
Apamil 200C
Adummii 2005
7
0350
Figure 5: Fhot evaluation on SuperGLUE. With the same number steps, models pre-trained with Sophia outperforms models pre-trained with AdamW and Lion on most lasks. Models pre-trained with Sophia for 200K steps have comparable performance as models pre-trained with AdamW for 400K steps.
(a) GPT-2 Small (12251)
[b) GPT2 Madiam (355M) 3ed
[c] Sensitivity to and
AdamW
3.2
0.02-
3335
3.349
3.348
6.x
Lion
53K
Soptia
3.3
Aemii Al EOOEA LS54
3.0
0.01
3.345
3.341
3.346
31K
2.9
Nombe
0.005-
3.341
3.353
OK
25%
75K
100K
DK 25C 50K
100K
0.99
Number Steps
Number Steps
Figure 6: (a) AdamW and Lion trigger gradient clipping frequently. With Sophia, gradient clipping rarely happens. (b) AdamW and Lion require re- parameterizing attention with temperature (inverse of layer index). The plot shows the largest LR that AdamW and Lion without the trick allow to be stable, which is much smaller than with the trick. contrast. Sophia does not need this trick. (c) Sophia is not sensitive to ry and B2 choice.
AdamW, W. Lion, and AdaHessian. As the model size grows, the gap between Sophia and baselines also becomes larger. Sophia-H and Sophia-G both achieve a 0.05 smaller validation loss on the 355M/ 770M model (Figure 4 (b) (c)) with the same 100K steps. This is significant improvement since according to scaling laws in this regime (Kaplan et al., 2020; Hoffmann et al., 2022) and results in Figure 3, 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 achieves the same loss in 50% less steps, total compute and wall-clock time. The im- provement in validation loss brought by Sophia can be translated into reduction of number of steps or total compute. In Figure and Figure 3, we evaluate the optimizers by comparing the number of steps or total compute needed to achieve the same or smaller validation loss. Sophia achieve a 2x speedup compared with AdamW and Lion across different model sizes.
The scaling law is in favor of Sophia over Ad In Figure 1 (d), we plot the validation loss on OpenWebText of models of different sizes pre-trained for 100K steps. On OpenWeeTers the gap between Sophia and Adam with 100K targeted steps is larger on 355M / 770M models than 125M models. Moreover, the 540M model trained by Sophia-H has smaller loss than the 770M model trained by A. The 355M model trained by Sophia-H has comparable loss as the 540M model trained by Adam
Few-shot Evaluation on SuperGLUE. As shown in Figure 5, the improvement in validation loss transfers to downstream tasks. With the same number of steps in pre-training, GPT-2 and GPT NeoX pre-trained with Sophia have better few-shot accuracy on most subtasks. Models pre- pre-trained with Sophia have comparable few-shot accuracy as models pre- pre-trained with Adam W for 2x 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 aver- age time per step (T(step)), time spent in Hessian computa- tion (T(Hessian)) and total compute following Chowdhery et al. (2022). Since we calculate the diagonal Hessian esti- mate with a reduced batch size every 10 steps, the compu- tation of the Hessian accounts for 6% of the total compute, and the overall wall-clock time overhead is less than 5% compared with Adam W. In terms of memory usage, our optimizer has two states, and h, which results in the same memory cost as Adam W.
Table 1: Wall-clock time and compute.
Algorithm Model Tisep)
(etssepp//[lllllissseee [[lessiaa) Compute
AdamW
770M
3.25s
2550
Sophia-H
770M
3,40s
0.12s
2708
Sophia-G
770M
3.42s
0.17s
2678
Adamw
355M
1.77s
1195
Suphia.
355M
1.88s
0.0%
1249
Sophia-G
355M
1.86s
0.09s
1255
Sensitivity to p and 82, and transferability of hyperparameters. On 30M model, we perform a grid search to test the sensitivity of Sophia-H to hyperparamters (Figure 6 (c)). All combinations have a similar performance, while 32- = 0.99 and y = 0.01 performs the best. Moreover, this
8

---

Published as a conference paper at ICLR 2024
(a) Choices of k
(6) Diagonal Hessian Preconditooness
(c) Clipping
3.37
AdamiW
3.37
3.39
Witthtton
Soptis G. k-1
Sophia-4, k=100
Sophia-G, k-10
Vallcation
$3.5
AdaniN
Alte:1p
Sophieet
3.35
Adaltessian. k-1
k-2 Sophla
Nermalize
3.33
Sophla-5
3:33
3.31
3.31
3.31
10
2.
30
20
60 Compute soaaflups
Compute/ Mnaflops
Computee saatfoos
Figure 7: (a) Hessian update frequency. (b) Diagonal Hessian pre-conditioner. (c) Element-wise clipping.
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, 52- = 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 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 6 (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.
gradient computed from rare data (Zhang
One trick of pre-training deep Transformers is scaling the product of keys and values by the inverse of layer index as implemented by Mistral (Karamcheti 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 shown in Figure 6 (b). In all the experiments, Sophia-H 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. We study the effect of Hessian update frequency k of Sophia-G G on computational overhead and validation loss on a 30M GPT-2. We consider 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 grid search. We plot the amount of compute and the validation loss of each run in Figure 7 (a). While k = 1 has better validation loss with the same number of steps, the computational overhead is 50% and the speed w.r.t amount of compute is worse than k- = 10. The choice of k 100 still outperforms AdamW, but is not good as k 10.
Diagonal Hcssian pre-conditioners. We compare different diagonal Hessian pre-conditioners (with the same k- 10 and found by grid search): Empirical Fisher (E-F+clip), AdaHessian (AH+clip), Hutchinson (Sophia-H) H), and GNB (Sophia-G). Note that empirical Fisher is the EMA of squared gradients, which differs from GNB in label sampling. Results in Figure 7 (b) indicate that GNB is better than Empirical Fisher, which is consistent with Kunstner et al. (2019). Sophia-H is also consistently better than AdaHessia.. 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 different update clipping strategies in Figure 7 (c): element- wise clipping without pre-conditioners (Clip), update normalization without pre-conditioners (Nor- malize), AdaHessian and Sophia-G without clipping (GNB). The learning rate is found by grid search. Note that clipping without pre-conditioner essentially the same as sign momentum, or Lion with B1 = 32. Without element-wise clipping, we find that AdaHessian will diverge with == 2 and GNB will diverge with A = 5. thus we use A = 1 for AdaHessian and 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.
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 2x speed-up compared with Adam Wi in the number of steps, total compute, and wall-clock time.

---

Published as a conference paper at ICLR 2024
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv preprint arXiv:2002.09018, 2020.
Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kronecker-factoree approximations. In International Conference on Learning Representations, 2017.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. International Conference on Machine Learning, pp. 404-413. PMLR, 2018
MS Bartlett. Approximate confidence intervals. Biometrika, 40(1/2): 1953.
Sue Becker and Yann Le Cun. Improving the convergence of back- pack-propagation learning with. 1988.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569, PMLR, 2018.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et Gpt-neox-20hh An open-source autore- 2022.
gressive language model. arXiv preprint arXiv:2204.066
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. International Conference on Machine Learning, pp. 557-565. PMLR, 2017.
Stephen Boyd and Lieven Vandenberghe. Convex optimization Cambridge university press, 2004.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Miine, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL gåttttp:// com/ com/googles google/ jax.
Heinz Braun and Martin Riedmiller. Rprop: fast adaptive learning algorithm. Proceedings of the International Symposium on Computer and Information Science VII, 1992.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamners. Advances in neural information processing systems, 33:1877-1901, 2020.
Charles George Broyden. The convergence of class of double-rank minimization algorithms 1. general considerations. IMA Journal of Applied Mathematics, 6(1):76-90, 1970.
Olivier Chapelle, Dumitru Erhan, et al. Improved preconditioner for hessian free optimization. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, volume 201. Citeseer, 2011.
Pei Chen. Hessian matrix vs. newson hessian matrix. SIAM Journal Numerical Analysis, 49(4):1417-1435, 2011.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675. 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 2022.
Andrew R Conn, N Gould, and Ph L Toint. Trust-region methods, siam. MPS, Philadelphia, 2000.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. arXiv preprint arXiv:2208. 11195, 2022.
10

---

Published as a conference paper at ICLR 2024
John E. Dennis Jr and Robert B Schnabel. Numerical methods for unconstrained optimization and nonlinear equations. SIAM, 1996.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. arXiv preprint arXiv: 2018.
Timothy Dozat. Incorporating nesterov momentum into adam. 2016.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(lu):2121-2155) 2011.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 000027, 2020.
Matilde Gargiani, Andrea Zanelli, Moritz Dichl, and Frank Hutter. On the promise of the stochastic generalized gauss-newten method for training dnns. arXiv preprint v7Xi::2006002409 02409, 2020.
Rong Ge and Tengyu Ma. Decomposing overcomplete 3rd order tensors using sum-of-squares algo-
rithms. arXiv preprint wrXiv:1504.05285 2015.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approx- imate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pp. 2232-2244 PMLR, 2019.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, 2019.
Roger Grosse. Neural Network Training Dynamics. 2022.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lec- ture 6a overview of mini-batch gradient descent. Ciled on, 14(8):2, 2012.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications sitiistics-Simlattionn and Computation, 1803):0559-1066. 1989.
Peter Izsak, Moshe Berchansky, and Omer Levy. How to train bert with academic budget. arXiv preprint arXiv:2104.07705, 2021.
Majid Jahani, Sergey Rusakov, Zheng Shi, Peter Richtárik, Michael W Mahoney, and Martin Takáč. Doubly adaptive scaled algorithm for machine learning using second-order information. arXiv preprint arXiv:2109.05198 2021.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How escape saddle points efficiently. In International Conference on Machine Learning, pp. 1724-1732. PMLR, 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
11

---

Published as a conference paper at ICLR 2024
Siddharth Karamcheti, Laurel Orr, Jason Bolton, Tianyi Zhang, Karan Goel, Avanika Narayan, Rishi Bommasani, Deepak Narayanan, Tatsunori Hashimoto, Dan Jurafsky, Christopher D. Manning, Christopher Potts, Christopher Ré, and Percy Liang. Mistral - journey towards reproducible lan- guage model training. https:: ctr...r.cm///////////////// stanford. edu/ wwtpt/eet/////e/eeeeete http:///mmmmmmmmmmmmmmmaaa .tmll, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 69800 2014.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approxima- tion for natural gradient descent. Advances in neural information processing systems, 32, 2019.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304. / /3660, 2023.
Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent converges to minimizers. arXiv preprini arXiv: 1602.04915, 1050:16, 2016.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249, 2020.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv
preprint (60:::660000038333 2016.
lyya Loshchilov and Frank Hutter.
Decoupled weight decay regularization. arXiv preprint
1rXiv:1711.05101 2017.
Vien V Mai and Mikael Johansson. Stability and convergence stochastic gradient clipping: Beyond lipschitz continuity and smoothness. In International Conference on Machine Learning, pp. 7325- 7335. PMLR, 2021.
James Martens. New insights and perspectives the natural gradient method. The Journal of Machine Learning Research, 21()):576-5881, 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference machine learning, pp. 2408-2417. PMLR, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
James Martens et al. Deep learning via hessian-free optimization. In ICML, volume 27, pp. 735-742, 2010.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing 1stm lan- guage models. arXiv preprint arXiv: 7708,0218.. 2017.
Yurii Nesterov and Boris T Polyak. Cubic regularization of newton method and its global perfor- mance. Mathematical Programming, 108(1):177-200 2006.
OpenAl. Gpt-4 technical report. arXiv, 2023.
James M Ortega and Werner C Rheinboldt. Iterative solution of nonlinear equations in several variables. SIAM, 2000.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.330 2013.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle. A. Beygelzimer, F. d'Alché-Bun Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8033.. Curran Associates, Inc., 2019. URL http: httpp/aappraase/aa/ee neur wtt.riss..ccc.aeeee cc/pcpcp//////ccccccpcc 9015-pytorch- an- -nt---ann-mmm--aaeeee iemmmmeeiiieee oeeereeeeeiiii--eeroereeeeeeeeee -eg pdf.
12

---

Published as a conference paper at ICLR 2024
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)9, 2019.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112. 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv:1904. 09237, 2019.
Farbod Roosta-Khorasani and Uri Ascher. Improved bounds on sample size for implicit matrix trace estimators. Foundations of Computational Mathematics, 1565)11187-122,, 2015.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singu- larity and beyond. arXiv preprint arXiv:1611.07476, 2016.
Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization, In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 9481-9488 2021.
Tom Schaul, Sixin Zhang, and Yann LeCun, No more pesky learning rates. In International confer- ence on machine learning, pp. 343-351. PMLR, 2013.
Nicol Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neu- ral computation, 14(7):1723-1738, 2002.
2002.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604, PMLR, 2018
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958 2014.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint ouXiv:/706.03// 2017.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general- purppss language un- derstanding systems. Advances in neural information processing systems, 32, 2019.
Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. arXiv preprint arXiv:2002. 12915, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the tre- art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38 45, Online, October 2020. As- sociation for Computational Linguistics. URL https: www..w ac web. org/ aanthooooooyy/ 2020. momnlpdemmce 6.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pp. 581-590. IEEE, 2020.
13

---

Published as a conference paper at ICLR 2024
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 10665-10673, 2021.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learn- ing: Training bert in 76 minutes. arXiv preprint arXiv: /00::90044 1904.009642 2019.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint 190. 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33: 15313-153333, 2020.
Lin Zhang, Shaohuai Shi, and Bo Li. Eva: Practical second-order optimization with kronecker- vectorized approximation. In The Eleventh International Conference on Learning Representations, 2022a.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint AXXv:22005.00666 2022b.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvoreek, Xenophon Pa- pademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33:18795-18888 2020.
14

---

Published as a conference paper at ICLR 2024
A 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., 201 1) adapted the learning rate of features by estimated geometry and assign larger learning rate to infrequent fea- tures. RMSProp (Hinton et al., 2012) generalized RProp and is capable work with smaller batch sizes. Adam (Kingma & Ba, 2014) improved RMSProp by introducing a running average of gradi- ents, and has 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 vari- ants 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 a improved version of sign momentum SGD.
Second-order Optimizers in Deep Learning. Second-order optimizers are believed to have the po- tential 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 diag- onal 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 's approximation of Hessian and Fisher information matrix. Martens & Grosse (2015) and follow-up works (Ba et al., 2017; George et al., 2018; 8; Martens et al., 2018; Zhang et al., 2022a) proposed to approximate the Hessian based on the structure of neural networks. Yao al. (2021); Jahani et al. (2021) proposed 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 /atttttiionnmmmrrrrra overhead due 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 speed-up on large language models in total compute.
Gradient Clipping. Global gradient clipping has been standard practice in pre-training language models (Merity et al., 2017; Radford et al., 2019; Izsak et al., 2021; Zhang et al., 2022b). helps stabilizes training and avoids the effect of rare examples and large gradient noise. Zhang 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.
O Ottiiization 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 outper- forms SGD by 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 Transform- with AdaFactor (Shazeer & Stern, 2018), which is 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.
15

---

Published as a conference paper at ICLR 2024
B ADDITIONAL EXPERIMENT RESULTS
Dynamics of Sophia in training. We measure the &g norm of the EMA of the diagonal Hessian he and the proportion of parameters where clipping happens (that is, my/he is larger than 7) during pre- training Figure 8. 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 the algorithm.
1.0
100
Sophik
Proprtoo
0.8
0.7
0K
25k
50K
7SK
100K
0K
25K
50K
75K
100K
Stops
Figure 8: Visualization of training statistics. clipped. (b) l2 norm of the EMA of Hessian hT
(a) The proportion of parameters whose update is
Results with different number of steps. Due to space limit, runs with different number of steps and their comparison provided in Figure 9. Across different total number of steps, Sophia outperforms AdamW W and Lion with a large margin as the main experiments we presented in Section 3.3.
(a)
Small (125M)
(b) GPT-2 Medium (355M)
3.3
AdamW
3.4
AdamW, 100K
1.2
Lion
AdamW,, 200K
Sophia
3.2
Saphia, 100K
Validation
Sophia, 200K
3.0
8.0
2.8
2x Speedup
Spee
200K Number
100K Number
Figure 9: Results of training for different steps.
GPT NeoX, 6.6B
2.8
AdamW, 2e-4
Validation Loss
AdamW, 1.5e-4
2.6
Sophia-G
2.4
2.2
2.0
0K
25K Number Steps
50K
75K
100K
Figure 10: Results of 6.6B GPT NeoX tuning. We only carried two runs for AdamW and one run for Sophia on the 6.6B model. For AdamW, we tried peak learning rate 1. 1.2e-4 (black solid curve) and 1.5e-4 (black dashed curve). The loss of the 1.5e-4 run blew up at about 17K steps. This indicates that 1.2e-4 close to the largest peak learning rate for AdamW that setup. Note that with the unstable 1.5e-4 learning rate, Adam is slower than Sophia-G before blowing up.
C ADDITIONAL EXPERIMENT DETAILS
C. HYPERPARAMTER TUNING
The hyperparameters we consider for baselines are as follows: peak learning rate, B1, 82 and weight decay, All hyperparameters except the peak learning rate are tuned with grid search on a 30M GPT-
16

---

Published as a conference paper at ICLR 2024
Table 2:
Model Configurations and Peak Learning Rate.
Acronym
Size
d model
head
depth
AdamW Ir
Lion Ir
Sophia- H Ir Sophia-Gr G Ir
30M
384
6
6
1.2e-3
4e-4
1e-3
le-3
Small
125M
768
12
12
6c-4
1.5c-4
6c-4
6c-4
Medium
355M
1024
3e-4
6e-5
4e-4
4e-4
540M
1152
18
30
3e-4
4e-4
4e-4
Large
770M
1280
20
36
2e-4
3e-4
3e-4
NcoX 1.5B
1.5B
1536
24
48
1.5c-4
1.2c-4
NeoX 6.6B
6.6B
4096
32
32
1.2c-4
6c-5
2 trained for 50K steps. The peak learning rate is tuned on models of different sizes with grid search separately. We search 31 in [0.8, 0.9, 0.95, 0.96, 0.99] and f8 in [0.9,0.95, 0.98, 0.99, 0.995]. Weight decay is chosen from 0.05, 0, 0.2, 0.5. For Lion, we also include 31 = 0.95, 82 = 0.98 as suggested by Chen et al. (2023). On 30M models, we found AdamW is sensitive to the choice of 31 but not 82. 31 0.9 works the best for AdamW while 3 = 0.95 works the best for Lion. We use 322 = 0.95 for Adam W since this is the dominantly used configuration in the LLM pre-training literature, and B2 0.98 for Lion as it is recommended by Chen et al. (2023). We found that weight decay 0.1 works the best for Adam W, while 0.2 works the best for Lion and Sophia.
For peak learning rate 125M and 355M, we perform a grid search. For larger models, 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. 1.5e-4, 1.2e-4, le-4, 8e-5, 6e-5, 4e-5]. For example, 6.6B GPT NeoX model work with I 1.2c-4 peak learning rate, but the loss will blow up if we increase the learning rate 1.5e-4 shown in Figure 10. The result of grid search of peak learning rate is provided in Table 2.
We use 31 = 0.96, B. 0.99, =1e-12 and ls 10 for Sophia. We adopt the following procedure to obtain these hyper parameters. We first fix 7 = 0.01, k = 10, and tune 31 and 33 with grid search on a 30M model, and directly use 3n and 3g from the 30M model on models of larger sizes. Similar to AdamW, we find that Sophia is not sensitive to 32. We then fix B1 0.96, 82 = 0.99 and tuning k 10. As shown in in Figure 7 (a), = 10 is better than h = 1 or k = 100 in terms of the balance between convergence speed and the computation overhead.
After finding out 31 0.96, 332 0.99, le 12 and k: = 10 with the method above, we tune y and peak learning rate jointly. We first tune 7 to make the proportion of coordinates where the update is not clipped (i.c., |m./ max(y- h,, e]| < 1) in the range of 10% - 50%. We search for y in the 00t.0.0....22..05....0.
list of while y 0.05 works the best for Sophia-G. We then fix 31, .2 = 0.99, 6, 7, k for all larger models.
As result we find out 0.01 works the best for Sophia-H
To tune the peak learning rate, adopt the same procedure as result of grid search of peak learning rate is also provided in Table 2.
for baseline methods. The
C.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 540M model for scaling law visualization. We provide the model specifications in Table 2. We use the nanoGPT (https:/ /tsp::/gtt /gi thub. , com/ hinn/aaze karpathyyyy/... /tthy//aanm /nanoGT//)// 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 Open WebText (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 69loat16. The 125M and 355M models are trained on machines with 10 A5000 GPUs, while the 770M models are trained on an AWS p4x.22xl instance with 8 A100 GPUs.
We consider 1.SB 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: httg//tttuit com/ stanford wwwrd-crfmleeeaatteee wwtpnnterrrreeeeeeeeee
for GPT NeoX. We use fully sharded data paral- lel with gradient accumulation to enable batch size of 512 for the 1.5B model and 1024 for the 6.6B model. These models are trained on a TPU 128 slice.
17

---

Published as a conference paper at ICLR 2024
The contest a paseges contaming some informtioon Givnn quistion about the content, use the information arswer the question with wther "te'' hoi.
Given premise and hypothesis, anower whather the hypothecs legienlly to from the premise with True' Taker NMather"
Context: 3 way limp The center contact of the bull typically connects to the medium power flamnen, and the carnects the lowepower filemen.. this if s-may bubt serewed into standard light socket that only center contact, only the medium power fllmeet opurates. this bubh regutor socket will resultin betoxing like normal 10Mk buth Queetion Trony bullss
Cantect: she that when her hualand died that
of the 110 150 w butb, putting
unds had said that would never put her
home. So
kind itt don't know iiee,, don''t thask
work in any tamp
would but sho getting pretty bas like she has have like inttl tollie right heribed and, it', Uh huh, B and mom don't linass
Answer: Yes
of her pretoy much hard decision, butt don't think would do to my parents
to
Perfume: The Story of Munderer (fiim) Pertume The Story of Strrderer a 2005 Geman periodt poychological crime thriller film directed by fom Tylwer starting Ban Whishhm,, Alan Rekmann Rachal Hurd Wood, and Dusin hoffman. Tykwer with Johnrr dimek and Roinhold Hall alo composed the music, The scneerpiay by vyowe,, Andhow linknn and tlend richinge is based on Pairic sose nc's 1905 masel Perfume, see 18th century France, the fim telss the story of Jean Bapeste Grenoullle (Whishow), on offactory genus,, and his homidical quest for the perfect scent. Quvestion the film perfume based a true story Answar: ho
Contest:
personally, Question: she would de to her parents Answer No
Contect No, was, I dion'' like the way it ended. A know, well the only resson know shbay it ended is Arsenio Hall one night, Chr stopher Recves told, that, you know, 8: Uh huhh A: tan't bellase they billed them. Qu Quentton they iolled them Ansser: Ves
BaolQ
CB
premise and
Choose the coreet ending for the sontect.
the
rewealed and wite advied local Hill Bank Antorveneto agalnst
talken mergers,
Choles?: the women made him barch.
Cholce the woman cissed Nim.
Question:
Antomvnneaa.
Ansoors Chorcdt
Contest: The man had lipetic his cheek Secause
Contest:
The reapened public July afterite exennshe refurbitiing 1936 calendar.
Cholos1: uttanded voga class.
Cholco2: boought fruits and vegetables.
Queettonn The Statute liberty buift 1986.
Wednesday
Answer: Choice2
Contest:: made resolction eat healthy diet
Answer: hone
RTE
COPA
Figure 11: Prompts for SuperGLUE downstream evaluation.
GPT
GPP2 Large (77)M)
Scaling
3.4
AdammV
Adsmw, 100K
2.8
Adamli Soptra-G
Adamle
Scotil-
2.7
Sosh a
Soptia-G
3.0
Sosh a
2:9
2..5
2.3
Spacdop
2.5
Speedip
Compute cKIPLOPS
Number of Stops
Numbbor of Stops
Model Size IN
Figure 12: An enlarged version of Figure 1.
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 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.
C.3 DOWNTREAM EVALUATION
We perform fho- evaluation of the models subtasks of SuperGLUE. We use 2-shot prompt- ing 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.
18

---

Published as a conference paper at ICLR 2024
D A DETAILED EXPLANATION OF GNB ESTIMATOR
We leverage the structure of the loss to design a biased stochastic estimator for the diagonal Hessian, following Schraudolph (2002); Martens (2020); Wei et al. (2020). Suppose E(O, (x, v)) is loss func- tion on an example (I, y) of the form E(O, (z. v)) = (S(, x), y) where foo is the cross-entropy loss and f(0,zz) RV is the logits, and V is the number of items/classes in multi-class classification problem (e.g., the vocabulary size in LLMs). First, the Hessian of E(O, (æ, y)) (w.r.t to variable Ø) has the well-known Gauss-Newtoo decomposition (Ortega & Rheinboldt, 2000; Schraudolph, 2002) (which is a simple consequence of the chain rule),
Vi ((0) (0, S (o)( x) + J lq
(11)
where (0, is the Jacobian of f w.r.t 0 viewed as a matrix in Rd S 8 (t..v)))) 812
e
t= (=)(0,))
RVV is the second-order derivatives of the loss w.r. to the logits, g (1.((..))
E R is
(0,z)
the first-order derivatives of the loss w.r. to the logits, and Joo (0, x) is the second-order derivatives of the multi-variate function f(0, æ²) w.r.t 0, viewed as a linear map from RV Rdddd where d is the dimension of the parameter 0.
In the context of neural networks, past works have found that the second term [O, x) [)] in Equa- tion 11 is often relative smaller than the first term JeJ [0, x) S (0, (Sankar et al., 2021), which is often referred as the Gauss-Newton matrix (Dennis Jr & Schnabel, 1996; Ortega & Rheinboldt, 2000; Schraudolph, 2002; Chen, 201 1) and used as pre-conditioners in second-order op- timizers (Botev al., 2017; Martens, 2020; Gargiani et al., 2020). Following this line of work, we build an unbiased estimator for the diagonal of the Gauss-Newton matrix, which is a biased estimator for the diagonal of the Hessian.
We first claim that S only depends f(0, .)) but not y. even though the loss depends on y. Thus, S 8 (t.0)
(tt..) 012
for any ý E (1,..., V}. which implies that S = Egg-p(e.))
(8,()).) Because loo (t, 3) is the negative log-probability of the probabilistic model defined by the categorical distribution Cat(t) with parameter t, by Bartlett's second identity (Bartlett, 1953), we have that,
f(0
(1, 9) alo (1, Ot
S
E (a-C((t)
E Catt))))
(12)
0t2
Ot
where the first equality holds for f00 and the second equality holds for all by Bartlett's second identity. Therefore, the Gauss-Newton matrix satisfies
JoS(O, Olc (1. 9)
(l.
Jos(0,x) S. JoS(0x)) = E
Jo)
)~Cal(()
Øt
Øt
E g-Cal(t)
( x) ())
(13)
which implies that diag -) x) By-Caat) (f o Pf x). D) Hence, the quantity oo (0, x). 9) (f 9) is an unbiased estimator of the -ewssnNewton matrix for the Hessian of a onc-example loss {f v).
Mini-batch version. Given a mini- batch of inputs {(16, y6) H 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
B
1
M B 6-
(lo(f((,,z)).
96) 2 (f
(14)
DDnotte by p(0,z)) softmax(f(0,22)) 2)) E RV the probability vector obtained by applying softmax on the logits. Indeed, a simple derivation shows that S = diagonal(p(O, x)) - p(0,=))(( =))(0,=)) where diagonal(p(O,, .r) is the matrix with the vector p(0, .x) residing on the diagonal. In fact, this is a general prop- erty of exponential families-the Hessian of the negative log-likelihood of any exponential family distribution only depends on the parameters of that exponential family, but not on the example on which the likelihood is evaluated.
19

---

Published as a conference paper at ICLR 2024
where Db is are labels sampled from the model on inputs zg's respectively. However, as noted by Grosse (2022), implementing this estimator inconvenient under the current nuto-differentiation frameworks, where the users only have access to the average gradient over 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:
vb, Egg (f xb), (6) = 0.
(15)
Let [(0) E (f mx), 96) be the mini-batch loss on the sampled labels (as opposed to the original labels). Öbserving that yb are independent with each other, have
Eg B VoD( e vo Egg
1
V (f(0,
Σ (fo,
B
B
1
= -Eg
Σ
(f 26), O (lcl(((0,
(16)
B
b=1
Note that the RHS of Equation 16 is the same as the expectation of Equation 14, which, by Equa- tion 13, also equals to the diagonal of the Gauss-Newton matrix for the mini-batch loss. Hence, use B (() @ voī( as the estimator.
To the best of our knowledge, this estimator of Gauss-Newtor matrix was first used in Wei et al. (2020). Given the use Bartlett's first and second identities that are central to the estimator, we call it Gauss-Newton Bartlett (GNB) estimator.
E LIMITATIONS
Scaling up to larger models and datasets. Although Sophia demonstrates scalability up to 770M models and Open WebText, and there is essential constraints from further scaling up, we do not compare with Adam V and Lion on larger models and datasets due to limited resources. We believe Sophia is faster than Adam W and Lion on larger models given the improvement in scaling laws and better pre- pre-training stability.
Holistic downstream evaluation. We evaluate pre-trained checkpoints 4 SuperGLUE subtasks, which only demonstrates the improvement in downstream performance several datasets. While a holistic evaluation of language models itself is open research topic, better downstream evaluation is still important. The limitation in downstream evaluation is also due to the limited model size, because language models at this scale do not have enough capabilities such in-context learning, and mathematical reasoning.
Evaluation on other domains. While this paper focuses optimizers for large language modeling, a more general optimizer should also be evaluated in other domains such as computer vision, rein- forcement learning, and multimodel tasks. Due the limitation of computation resources, we leave the application to other domains and models to future works.
F 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 mini- mum) 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 F.1. L d + R is twice continuously differentiable, strictly convex function with 0* being its minimizer. For convenience, we denote Amiin 2 L((*)) by p.
The following assumptions state that the Hessian has a certain form of continuity-wittee within neigh- borhood of size It, the ratio between the Hessians, 2²0(0') 2 L()), assumed to be bounded by a constant 2.
Assumption F.2. v0, 0' E Rd 110
There exists a constant R > O, such that
112 < R
||22100)~~~>100) $2
(17)
20

---

Published as a conference paper at ICLR 2024
We analyze the convergence rate of the deterministic version of the Sophia on convex functions,
P 0. vV,'' clip((((==(()))><V((()) p),
(18)
where (0) = V.V, is an eigendecomposition of (2), Here, we use the full Hessian the crn-connit because the diagonal Hessian pre-conditioner cannot always work for gen- eral functions which may not have any aligmment with the natural coordinate system. Moreover, the matrix V, transforms (V2)().)) (L. into eigenspace and thus the clipping element-wise in the eigenspace. We do not need the max between Hessian and in the original version of Sophia because the Hessian is always PSD for convex functions. Finally, the matrix V, transforms the update back to the original coordinate system for the parameter update.
be done
Theorem F.3. Under Assumption F.1 and Assumption F.2, let 1j = 1/2, p = 2V8 the update in Equa-
Rt
tion 18 reaches loss most in T S d- L(Pu)-min L
n + In RR² 32dr
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/c 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 cigenvalue of the Hessian (or the smoothness parameter) in standard analysis of gradient descent in convex optimization (Boyd & Vandenberghe, 2004). Moreover, quadratic functions, the convergence rate of simplified Adam (SignGD) depends on the condition number (Appendix G.I). This demonstrates the advantage of Sophia in adapting heterogeneous curvatures across parameter dimensions.
on simple
G THEORETICAL ANALYSES: DETAILS OF SECTION F
Theorem F3 is direct combination of the Lemma G.10 (Descent Lemma), Lemma G.9 and Lemma G.11. In the analysis, there will be two phases. In the first phase decrease loss to up" in 8 L(8(0))- steps. In the second phase, there will be an exponential decay of error.
Lemma G.1. Under Assumption F.I, we have that L(0) + oo whenever |01|| + oc.
Proof of Lemma G.I. By convexity of L, we have v0 E E2 with 10
0" 112 1,
1
|10
8 11 L(0") 2 L(0"
0 0+
L(0)
2
min L(O0 + 0).
(19)
00 - 0- 10-0-11
1
0* l1
10
0+
l
Since strictly convex, A min | L(0 0) L(0") > 0. Thus we conclude that
L( Z ||0 81112 A L( 2 (110)] - ||++ L(0").
(20)
Therefore when ||0]| > oc, (())) oo as well.
Note that we don't 't assume the Hessian of loss is Lipschitz. Assumption F.2 only assumes the Hessian in neighborhood of constant radius only differs by a constant in the multiplicative sense.
Lemma G.2. For any 0 E Rdd satisfying L(0) - min < RR² it holds that ||e 1
ll S
L(6) [10)-miul < R. 2
Proof of Lemma G.2. We will prove by contradiction. Suppose there exists such 8 with L(0) aR²
but ||0 8* 1|2 > 2 L(e) L(o)-min We consider 0' & 0 + V 2L(0) 8-0" and 0* and that L is strictly convex, know that L(O') L()). However, by Taylor expansion function f(t) A L( t(0' 0*)), we have that
Since 0' is between 1-
f(1) = f(0)+ + (() f*(
for some E [0, 1].
(21)
2
21

---

Published as a conference paper at ICLR 2024
Note that 106 12 < 10 - 12 S R. by Assumption F.2 and Assumption F.I, we have (t)
(0' /-) L(1'' (1 - (0)0) 0*) 7 (0' 0²) L 0*) 5 110 0-1 /
2(((0) - min L)). Also note that f(1) = L(0 J. f() L( and f (0) = 0, we conclude that
L(0') - ((0)) 2 L(0) - L(0*), namely (0") 2 (0). Contradiction!
Lemma G.3. For any 0 E Id satisfying that [[
1 2
Ru is holds that l|e - 8* ll2 1
2| < R.
Proof of Lemma G.3. We will prove by contradiction. We consider function f(t)
8-0' VL(O" + 8 8-6' T0-0 Because of the strict convexity of L, f is strict monotone
increasing function. If |00 8
2 but [|VL(0)]|2 I 2. then we have f(R) <
f(le- - 8 1l)) S [VVL())] On the other hand, by Assumption F.2 and Assumption F.I, f'(t) 2 2
for E [0, Rl. Thus f(R)= 2 f(0) Se= 21V1(111
(((()tt= |/ [.)]] Contradiction!
Lemma G.4. For any 0 Rd the following differential equation has at least one solution on interval [0,
d0(t)
dt =V(V)L(()))~>
8(0)= 0
(22)
and the solution satisfies that
all t [0, 1] and ((0) 0"
Proof of Lemma G.4. Since L is continuous and positive definite by Assumption F.1, (V2T. -1 continuous and thus the above ODE (47) has a solution over interval [0, for some positive T and we let Taxx be the largest positive number such that the solution exists (or T = oo). Now we claim Taax 2 1, otherwise [[O(t) - ll2 must diverge to infinity when f + Tmax. However, for any S 1, we have
dV LV)())
(23)
di
which implies that VL(O(()) = -1 (1-1))100 for all 1 E [o. 1]. Therefore, aL(O()) (V7,000))) dt
(1 (()))))====((((())(0))()(() (VVIL))
S 0.
(24)
Thus L(((t)) < L(O(0)). By Lemma G.I. we know that [[((t]] remains bounded for all e [0, maxx thus Tmax 2 1. Note that 0(1) has zero gradient, ((1) must be 0" This completes the proof.
Lemma G.5. For any 0 E Rdd satisfying (1) T.(0) min L. S that
AR² 16
(2) [V(()) S
Ry it holds 4.
L(0) min L S VL (V VL(() < 4(L(0)
min L).
(25)
Proof of Lemma G.5. Let {0(t)] be the solution of Equation 47. We know that VL(((t))
(1 (() for all t E (0, and that 0(1) = 0 by Lemma G.4. For case (1), by Lemma G.2,
know that for any 1 1. [|O(t)- - 0111 R/2. For case (2), by Lemma G.3, know that
for any t E [0, 1). [|0(t) 0=]12 < R/2. Thus in both two cases, [[0(t)) A1g = [[((t) 0(0)ll2 =s
[|0(t) 0- + 10(0)) 0*1 < R. By Assumption F.2, it holds that
L(V I(0))) _ (V
(26)
for all t E [0, 1]. Therefore, we have that
/' (ZL(001) 1=
L(0)- min L = L(O(0)) - L(((1))
(1 -
(27)
The proof is completed by plugging Equation 26 into Equation 27 and noting that Fi - 1) = 1/2.
22

---

Published as a conference paper at ICLR 2024
Lemma G.6. For any 0 Red satisfying (1) I(0)) min S I that
uR² (2)
Ry
holds
2
L( - min L S P V
(28)
Proof of Lemma G.6. The proof of Lemma G.6 is almost the same as that of Lemma G.5 and thus omitted.
Lemma G.7. For any Rd satisfying L(0) min S
PR 16 it holds that
\V- S 8(T(0) min T.)
(29)
P
Proof of Lemma G.7. By Lemma G.2, we have that 110 - 81112 S By Assumption F.2, we have (2) = ( = " Ta. By Lemma G.5, we have that
((L(0) min L) 2 =V(0)
(30)
(31)
(32)
This completes the proof.
Lemma G.8. For any Rdd satisfying that [(V"100)))) S 2' i holds that
L(0) - minmm S VL(O) (FLL0))~)0 - 4(L(0) min L).
(33)
Proof of Lemma G.8. Let {0(t)) be the solution of Equation 47 and we claim that for all t E
[0, 1]. [[((t) O|l2 < R. Otherwise, let T the smallest positive number such that [[( - 0l1
R. Such 7' exists because [[0(t) 0|1 is continuous in & and [|0(0) Oll2 = O. We have that
dd(t) dt 2
R= (((()) - ((0) <
dt
(34)
((V=L(00)))~~<<<() dt
(35)
dt
(36)
<2 [[(v'=(0))~~((0)1111
(37)
<2T 2 = RT,
(38)
which implies T = 1. Here in Equation 37, we use Assumption F.2. Thus we conclude that for all
te [0, 1). [l((t) 0|1 < R. By Assumption F.2, it holds that
2()2 2 (100))))0000000))))0000)10""0) 1 2 (V
(39)
Therefore, we have that
((0) min L L(((0)) - L(((1)) (FL(01))
(1 - 1111)V000)))TV>
(40)
=
The proof is completed by plugging Equation 39 into Equation 40 and noting that fi - 1) 1/2.
23

---

Published as a conference paper at ICLR 2024
Lemma G.9. p S
R 2/8
then for any A <
Rpn and any 0 E Rd satisfying
10
min VL()) oo |0 ([(O)) SA
(41)
where (0) EV is the eigendecomposition of V2)(0), 2s is the ith row of and Σ diag(O1.... Jd). it holds that
25A
L(0) min L v A +
(42)
p²y
In particular, if we set A - up²
20 we have L(0) min v
up² 8
Proof of Lemmna G.9. Let Lg - fie does not happen. Then we have that
[d]
VL())) o, S p} be the set of indices where clipping
Σ
oi 1 (L
<A
(43)
Σ i
ll VL( SA
(44)
Now we consider a new strictly convex loss function in Rlol,, which restricted on the space of {0 Eielo ihihh | glal], that is, Le(w)= L(0 EiEL, Wpgjoa). This new loss function Lg clearly satisfy Assumption F.2 since it a restriction of L into some subspace of Rdd By Lemma G. know that inf,, (w) be attained and denote by w" By Assumption F.1, we know that Le is strictly convex and thus Le > 0, which means Assumption F. also holds for Le.
Next we will apply Lemma G.8 on Lg at w = O. We use V1, E l xd to denote the submatrix of V containing rows in for any C [d]. One can verify by chain rule that VLs VI, V Vi, w) and that 2 Lo (w) VI, 2 L(0 + VT, V), Thus we have that
(V Lo)) (0) LV))
(45)
By the definition of 1g. we know that |1, (VL(())
x < p.
Thus
(0) (L)( < vd [V, (V2L(0))
(VL Vd.
R.
Thus apply
Lemma G.8 on Lg at w = and conclude that
(o Lo(w")
= E0 | AA
(46)
ieto
Thus L(0) - [. + V7, w²) = (0) - (1o))) A.
To do so, our strategy is to first show that [| + Vi (0))== is small and then to use Lemma G.6. We will use Iī to denote the complement of Ia in [d] and Viş E (x to denote the submatrix V which contains all the rows that do not belong to Ig. Note that w" is the minimizer of Le, know that V1, L + VJJ w²) 0 and that (Q V w*) I|2 V + V,J w=)]lp-
remains to show that L(0 + VIJ w²) - L(0') < 25 مap
Now consider the following ODE
due(t) - df
w(0) = 0.
(47)
By Lemma G.4, we know this ODE has solution w()) over interval [o. 1] with w((1 = w=. With the
same argument in the proof Lemma G.8, we know that [t]]] < R for E [0, 1]. Thus we
24

---

Published as a conference paper at ICLR 2024
have for any / E [0, 1)
LV + Via wo())
V1g
(48)
dt
V15 V2L(0 V1, (o
(49)
(50)
V13 V
(51)
vr,00++),w (V2 V., 10 (1,0
(52)
For the first term (Equation 51), by Assumption F.2, we have that
(L(O)
2A
ViEVV=L(0+ < V =
I
(53)
ig
igTo
For the second term (Equation 52), by Assumption F.2, we have that
VV + VI, VVPLL + v),
(54)
(55)
-VL((0) VI,
(56)
-8 Σ o.
1 (L()) 2 88A
(57)
PE
Thus we conclude that
V1
LV w())
1A which implies that
dt
VL + ViO xx
Vrg V L(O + VI w²) w²)ll
(58)
LV + V1, wt)) dt
Vig VL(O)
Vr8
(59)
dt
1
VL(())1 + [']
dVL(0 + V1, w())
V13
V15
dt
(60)
dt
A
44
54
(61)
P
P
Applying Lemma G.6, we have that
L(0 ViJ min L S 1 Vi, VIT
25A2
(62)
p²²
This completes the proof.
Lemma G.10 (Descent Lemma). For any 11, 0 with np < R//dd, 0 E Id and any eigendecom- position of L()), where V.V,T Ia. oi is diagonal (²) = V ΣV, define
0 A 0 nV (ipp(((VVVL()))>VVVL1))
(63)
it holds that
L(04) L(0) < -(1 772) Σ min(ple, (L .. 1o (
(64)
where Us the ith row of matrix V.
Proof of Lemma G.10. Let u e -Ip((((V2L()) By the definition of clip operation, we know that TTull, = ||ul|2 I Vdp. Thus we have ||0+ - Oll = Tull2 < npvd. Define
25

---

Published as a conference paper at ICLR 2024
((() T(10+ (1 - 1)0). By Assumption F.2, we know that (1) - (0) for all E [0, 1] and thus
[ /
(() = [(0)+
1=0 ("())d [
(65)
It remains to show that
((()) = - E min(p ( 1/ (.. '):
2. (0) 1 n" nd minfple VL(O). .0, 1
VL(0))
First, by chain rule, we have (0) = (VL(0), -n( = (VVL(0). -mu) =
-n (VVL(0),
= -1 Σi= min b, VL(0)) / VL()) }
Second, again by chain rule, we have (0) - 17² (V u, V2L(0) u) vf (u, Eu)
Eu.lu Note that by definition Wul = min(ly; (L(0)) /0,, p}. we have |u o; S
minfle; VL(O)) p} 11, |o, (L()) /0t a, min(/u [ /o. P b VL(0)[3, which com- pletes the proof.
Lemma G.11. np S R// Vd and for some T E L(Or) - min L S upp t t2 T,
8 then if holds that for all
1. 01+1 =0 -
2. L(O() min L (1 (111 min L).
Proof of Lemma G.I First by Lemma G.10, we have for all 2 T, (0.) min L S L(Or) min T. < po² "g² therefore by Lemma G.7, have [(V=T(0))- implies clipping will not happen. This completes the proof of the first claim.
S p for all 2 T, which
For the second claim, by Lemmas G.5 and G.10, we have that
L(0m-1) L(O.) s - (n n7 Σ o, 1 (()))
(66)
1=1
(1) - S - n1( (L)((()))))))) - min L),
(67)
(68)
which completes the proof.
G.1 LOWER BOUND FOR SIGNGD ON 2-DIMENSTONAT. QUADRATIC LOSS
Define L : I2 + R as quadratic function with parameter μ4, 3 as Lp,s - 6 + f0 We have the following lower bound, which shows signGD's convergence rate has to depend on the condition number B/p.
Theorem G.12. For any μ. 3, A, > 0, suppose there exist a learning rate and a time such that for all Go satisfying that Lu, A. signGD reaches loss at most at step T 1 and T (in the
sense that (0rr Scaa and ( S c). Then, must satisfy 2 2( -2) 3
Proof of Theorem 1. We consider two initialization: 0g = (0, and 06 = (v/ 24 0), and
let 0, and 6% be the iterates under the two initializations. For each coordinate i E (1,2), 2), because
[(0:)) (6:+ [[99] = 17. we have that (0.)(a] (60 )pp] 2 Thus 2c 2 L (0T) + Ly, (OT-)) 2
((0T) + (0r- 1)pp) z Sn² which implies 17 S
V
8c 3
The fact that ( L S 2c implies (04) S V Because SignGD can only move
each coordinate by at most, we have (T 1)n 2 v 24/H V
As Using the fact that 17 S
have that 2(T 1) (/3 v2) /. which completes the proof.
26

---

1. When the Hessian is much bigger than the gradient, then the update made by the second order algorithm (=) -1 L'()) second-odes algorithm cannot work. This oftentimes happens when the third-order deriva- Published as a conference paper at ICLR 2024
G.2 1D NON-CONVEX FUNCTION
In this subsection, we consider non-convex case. We start with an ID non-convex function L(0) where 0 E R.
For simplicity, we take = 0 and consider the deterministic version of Sophia,
0 + 1j -lippmxx("" (L(0)))))___(()))pp))
(69)
where we use the convention that for any I E I, x + 0, 2/0 = sign(x) oo and 0/0 = 0.
For simplicity, we first assume that L has single minimizer (When there are more than one local minima, the algorithm will simplify converges to the closest local minima.)
Assumption G.13. We assume that L is a three times differentiable funtitioo with only one minimizer, 0", where _ (07) > 0.
The key challenges for recond-order methods are two-fold as described below.
is too small. For example, when L(0) = (xp(cxp(0));
tives is much larger than the second-order derivative, and as a consequence the second-order derivative is larger than the first-order derivative.
2. When there exists a saddle point where the gradient is zero, and but the Hessian is very close to zero. Such saddle points are often called non-strict saddle points Ge & Ma (2015).
The first situation is a pathological situation and past works on Newton's 's method oftentime assume that it does not occur. Situation 2 is also common challenging situation that most prior works non-convex optimization excludes Ge & Ma (2015); Lee et al. (2016); Jin et al. (2017) because SGD and many first order methods cannot escape such saddle points. Our assumptions below assume that these two cases do not occur.
Definition G.14. For any 7, G. R*+, we say that function L IR - IR satisfies the (7, G. condition iff
srrict-addde--na andlee-an---a bee--buunddditee
G suu inf |IZ((0) c L' (0))) S sa) e>>
(70)
{ [1L"(00 I!(())
T."(0) 2 a/2
sup
(71)
Note that 7 > O means that any stationary point has a decent Hessian, that is, is either a local minimum or local maximum with reasonably strong curvature. This condition will matter more for the high-dimensional cases. Our bounds will depend on G and (though G and depends implicitly.) We also note that the bound (71) cannot hold over 0 E R (as opposed to over all 0 such that [L"(0)] 2 a/2) for non-convex function, because when (0) changes sign, i.e., from positive to negative, L"(B)
(0) has to unbounded (or even undefined when T!) = 0.)
Theorem G.15 (Main, ID Non-convex), Assume one- dimenninsio function L satisfies Assump- tion G.13 and (y, a.)
property (defined in Defi- nition G.14). Then, picking 7 = and p = 2-'' we have that the algorithm with in update Equa- tion 69 reaches loss most in at most ((T) steps where
T= L(Oo) - min T min(4a/72, G²/3 + In I((O)) min L
(72)
G.3 DECOMPOSABLE NON-CONVEX FUNCTIONS
Because our algorithm uses only element operations, makes sense to consider a non-convex func- tion [()) that be decomposed into of non-convex 1-D functions each individual pa- rameter:
d
((
(1))))
(73)
27

---

Published as a conference paper at ICLR 2024
As have been done in Equation 69, for simplicity, we take = and consider the deterministic version of Sophia, Vi E [c]00p] + 01--111-11 12 (i -0000))) L.)
(74)
Theorem G.16 (Non-convex, Decomposable). Assume function L is decomposable as in Equa- tion 73, Suppose each Li satisfies Assumption G.13 with minimizer oi and L (0 -ps. has har---nnddbuunddd-----------
(y, G,a)
property (defined in Definition G.14).
Then, picking = 2 and 2a we have that the algorithm with in update Equation 74 reaches a loss most in at most (T) steps where
L(6u) - min L min(4a/ min(4a//2,CP/a) + II L(0u) - min L
T=
(75)
We note that Theorem G.16 is direct corollary of Theorem G.15, The Theorem does not depend on the condition number (the ratio between the largest and smallest curvatures among all dimen- sions). Instead, assume that all dimension satisfy the same strict-saddle and bounded third-order derivative properties as stated in Definition G.14.
A strong theoretical analysis of any second-order optimizers' convergence for non-convex cases is of its own interest and we hope that our work can motivate more theoretical work on second-order algorithms in the future.
G.4 PROOFS THEOREM G.15
We state a slightly more general version of Theorem G.15 for generality.
Theorem G.17. Assume mne-immenioonal function L satisfies Assumption G.13 and (7, G, a) striet- saddle-and wewpendddd-iidddddr---reeeeeeeee
property (defined in Definition G.14). Then, for any nj < Oo E IR and c > 0, it holds that
1 - 24y
L(Oo) - min L (mpp G2/a3 + In L(Ob) min L
)
Y, T2 111 v)
L(OT) Se.
(76)
In particular, if we pick 17 = and 1 we have that 27''
L((b) - min In L(0o) min L
) L(OT) Se. (77)
V, T24
min G2/a3
6
Lemma G.18. Under Assumption G.13, for all e R. (0 - 0", L'())) 2 O. and the equality is only attained at 0 8"
Proof. Lemma G.18 It suffices show that there is only one stationary point, which is 8*. Suppose there is another stationary point which is larger than 0, let 6' be the smallest stationary point larger
than 8* We have that L'(0") L'(0*)= S So-p- L"(#) = 0 and for any A E (0",")), 1 ("(s)ds has the same sign. The sign must be positive because (O)(0) is continuous and positive 0" by As- sumption 131. Thus we conclude that T."(O') < 0. However, this implies 0' is a local maximizer. which contradicts with Assumption G.13.
Lemma G.19. Under Assumption G.13, for any 0' E IR such that L"(O") < 0, it holds that
(''(0') 2 sup [L'(O)) - ()) <
(78)
e>0
Proof of Lemma G.19. Since L" is continuous, and L"(O") < 0, L"(0*) > 0, we know for any e >> O, the set [['((0] - < (0) < 0) not empty. Now we claim that for any e > 0, there d such that [L.( S [L(('') and S (².).) O, which would imply that [L' [L'()') 2 sup.>o [T' (0.)) supeso inf { [IZ'(0] - (0) <0}.
Now we prove the above claim. Without loss of generality, we assume 0' > 0" By Lemma G.18, we know L'()) is positive on oo). Now let be the infimum of positives number such that
28

---

Published as a conference paper at ICLR 2024
(0" + A) = O. (if no positive A makes (0" + A) O. we set A = co. ) It is evident
that (0) is positive and decreasing on (0", f + A). We proceed in two cases: (1) If A is not
oo, then since I." is continuous, we know L" + +) = 0 and A is the smallest number such that L" + -) = Thus for any > O, there is some point sufficiently close to 0 + A which satisfies the requirement in the claim. (2) If A oo, note for any 0 > O', we have that -(0)) S
L'(0) - ('(') = S (s(s) & (0 - 0') supoelør,e)) L" Taking + oc, we know that supee[o oo) L²) = O. Thus for any > O, there is some point larger than 6' but with arbitrarily small Hessian (in absolute value) which satisfies the claim.
As a direct corollary of Lemma G.19, we have Corollary G.20.
Corollary G.20. Under Assumption G.13, for any 7, G, a E IR,, and function L : RR - IK satisfies the (7, G,a) strict-saddle-ande -nd--aandeddttd toreeeedd-iir--rr-rree thind-orderrderivatiie condition
G inf (T'(O)) - - <..
(79)
( 21""(01 [0) 2 a/2
7 sup
(80)
(.((0)
Lemma G.21. In the setting of Theorem for any 0, 0 E R satisfying 10 1/7(L// = 00 y = 0), it holds that
1. (0) 2 = 1,)(0)= 2 L"10)
= 2. (0) 2 a/2 = L')()) S (0):
Proof of Lemma G.21. We first define f(t) & f(0 + sign(('- We will prove the first claim by contradiction. Suppose the first claim is not true. then there exists t E (0, 1/yl such that ["(f(t) < L"(0), We let t' be the smallest number in (0. 1/7l that L'(f(tt) = ="()), Such t" always exists because L"()) is differentiable by Assumption G.13 and thus continuous. By definition of l",
2
2
we know that for all t E [0, 7/2 2 (L.)(((O) 2 [l)+t. Thus In2 2 L)f(tt)) [T([()
In L'((()) = S d'111"(1(0) dt 2 t -12.. However, - In2 < 1/22. Contradiction! dt
Suppose ['))(t) 2 a/2 for all t E [0. 1]. then we have In [' - 81)) - In L" (f())) =
- d (lnt"()())))) dt < 10" 0/// < 1/2 2. Otherwise, let t* be the largest number [0, 1) such that L" de [()) a/2. Applying the previous argument we have ("(d") < 2L"
<
(O))). This completes the proof.
Lemma G.22. In the setting of Theorem for any 0 E R satisfying that [((''(0))~<(0))
1 and that (0) > a, it holds that
24
L(0) min L < (L"(0))'1' < 4(L(0) - min L).
(81)
Proof of Lemma G.22. By Lemma G.21, we know that L"((") 2 (O) for all 10 0 < 1/7-
Thus we must have
(L)(0), I'(0 - sigg(I"((0) 4)) < [0) [T)(0) L' LL)
2y < 0. Thus
1 01 < 1/7 and L is
-strongly convex between @ and d*, which implies that 2
L((") 2 L(0) + (0" ()) + 10- 1 L" 2 L(0)
1 ['')))² (L"()))~*.
(82)
4
Similarly, by Lemma G.21, we know that L'())) < (0) for all 10 - 0' < 1/7. Thus we have
L(0") < L(l) (0* ()L + 00 2 (0) < L(Ø) ( (L' (L)
(83)
The proof is completed by noting that L((*)= = min L.
29

---

Published as a conference paper at ICLR 2024
Lemma G.23 (Descent Lemma for ID non-convex loss). In the setting of Theorem G.17, for any ?) > with np I 1/7 and 0 E Define
L'(0)
0+ _ -jclip
max (VLL(0),))'
(84)
holds that
[0)(0))
L(0+) - L(0) 1 (1 (n-1)) -²(²) min [L)()), max(LL"0),
(85)
I - (m - n min { G² /a, L(0) min L}
(86)
Proof of Lemma G.23 We first prove (85). Let clip
max V L000000 PP) By the definition of
(0)
clip operation, we know that |u| - p. Thus we have 10_ 0| ju| S np. We will proceed by
discussing two cases respectively: (0) 2 and (") S a. Define f(t) L(t0+ + (1 t)0).
First, for the case (") 2 By Lemma G.21, we know that (t) - 2f" for all E [0, 1] and thus
((1) = ((0) 1 [ l ('))dd & 5()) (0 f"
(87)
remains to show that
1. = 11 (m01 ('(0))
2. (0) S min (L'(0)-
The first claim is immediate by chain rule. The second claim holds because uL" [1'(0) by definition of u. This completes the proof of the first case.
For the second case, we have (,) S which implies that |u| 2 |u'] where w' A minf L'(0)
We have L(0+) L(0) f L'(0") )dd' S $S00 L'(8") due to Lemma G.18. Furthermore, we have that
0- J0-0 ( L'(O')d0'= = --n''''' +
O-mp
S" I""(sddsdl < -pp'''(00) + a
(88)
t=0
where the last we use (s)s) S 2a by Lemma G.21. The rest of the proof of the second case is the immediate and the same that of the first case.
['(
Now we turn to the proof of the simplified version (86). When (L((0) < max [''CO)
we have
(0)
L'(0) H
[L'(O) 2 pax. Thus min p [)())))), max (L"(),a
min
par
max L" (o).a
For the term
|='(0)' max ()) if (0) < a, then by Corollary G.20, we know [L'(0) 2 G Gnd thus (I))
L'(0))²
Otherwise, if L" 2 a, then by Lemma G.22, we have max (L)())))) 2 L(0) min L. This completes the proof.
Now we are ready to prove Theorem G.17.
Proof of Theorem G. 17. The theorem is direct consequence of Lemma G.23 (Descent Lemma for 1D loss).
30


Reviews:
[{"review_id": "7qK9rgwKce", "reviewer": "Reviewer_R1G3", "summary": "The paper introduces a new optimization algorithm for pre-training LLMs. Their algorithm combines several standard techniques in an interesting way to come up with a novel update rule. Specifically, the techniques combined include: Hessian pre-conditioning of the gradient, coordinate-wise clipping of values in the update step, the use of the Gauss-Newton-Bartlett estimator for the Hessian's diagonal entries (justified by previous works the authors cite), and the idea of performing approximate computations each time and more correct computations only periodically. While each of these techniques is itself quite classical, the paper's combination of these ideas to give a new algorithm that does quite well experimentally is very impressive.", "questions": "Thank you for an overall interesting and nicely written paper. I'd be interested in understanding the following. \n\n**Question 1.** The theoretical analysis (Section E, F) of the submission seem to make assumptions of convexity and slow-changing Hessian on the loss function. This, to me, seems somewhat at odds with the motivation for the key novelty of the algorithm, which is to clip the value of the positively clipped Hessian (from my understanding, based on Page 3, \"Limitations of Newton's method\" paragraph, it looks like the role of the clipped Hessian is precisely to \"mitigate the rapid change of Hessian\" and be less \"vulnerable to negative curvature\" (cf. Section 2.2)). From this, it seems to me that the theory in Sections E and F are perhaps justifying only the gradient clipping, not the clipping under nonconvexity or negative-curvature-Hessian. Do the authors have any additional theory for the algorithm when the loss is nonconvex and/or when the Hessian has a negative curvature or is rapidly changing? This would, in my opinion, truly complement what's in the experiments and greatly strengthen the paper. \n\n**Question 2.** The authors mention that their algorithm avoids local maxima and saddle points. My question is, would they be able to comment on what kind of local minima their algorithm converges to? Specifically, are these flat minima or sharp minima? Given that there has been a fair amount of work on the varying generalizability of these two types of minima, perhaps it might be important to understand this question so as to avoid overfitting.\n\n**Comment 1.** I found Figure 1 too small to see the details. Since the authors use this figure to explain much of the initial intuition, it would help to have it bigger. (I understand space is limited, so I'm definitely not penalizing this aspect; but perhaps there could be a bigger version of the same figure in the appendix, or perhaps the authors could use another visual medium to show the phenomenon of Fig. 1.) \n\n**Remark 1.** I would be happy to increase my score of \"Soundness\" and my \"Confidence\" were Q1 and Q2 to be satisfactorily answered. Regardless, I think this is a very good contribution to the ICLR community.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "I believe the novelty and creativity of the paper's algorithm for what is currently a very relevant topic of study is its biggest strength. What I mean by novelty/creativity is the update step of  $$\\max( \\min ( \\frac{\\hat{g}}{\\max( \\hat{H} , \\epsilon)} , 1 ), -1  ),$$ where $\\hat{H}$ and $\\hat{g}$ are the approximate Hessian and gradient, respectively. This formulation in some sense interpolates between using a Newton step and using SignSGD based on the local geometry of the loss function, thereby choosing the \"right technique\" between these two and avoiding the other's disadvantages (for example, avoiding getting too slow under heterogeneous curvatures due to Adam and avoiding getting stuck at a local-non-minimum due to Newton). In order to achieve this goal, the authors come up with computationally light estimators of the diagonal entries of the Hessian. \n\nThrough their above innovation, the authors make *second-order* methods computationally tractable for pre-training LLMs, which is a big achievement.", "weaknesses": "Please see Q1 and Q2 below. In my view, the paper could be stronger were these questions to be satisfactorily addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BSHAWged98", "reviewer": "Reviewer_x5VB", "summary": "The paper describes a second order optimizer Sophia. Sophia works by estimating the diagonal of the Hessian, and using it to precondition the gradient, in the family of Newton optimizers. The two main differences with previous similar work (Adahessian) are that Sophia applies per-parameter clipping to the Hessian to limit the effect of Hessian noise, and uses an exponential moving average (like Adam) to effectively compute diagonal Hessians over a larger batchsize, to reduce noise in the approximation.\n\nThe authors mention two methods to estimate the diagonal Hessian --- the standard Hutchinson method that uses Hessian-vector products, and a relatively new Gauss-Newton-Bartlett method that uses a second forward-backward pass. The latter produces better results in their experiments. To mitigate the extra time needed for the diagonal Hessian computations, they use the standard heuristic of computing the diagonal Hessian infrequently, they compute it once every 10 steps.\n\nThe authors test their method by training GPT2 models with up to 6.6B parameters, showing that Sophia is able to achieve significant savings over Adam, which is commonly used to train GPT models. They also show that the model trained with Sophia for the same number of steps gets better accuracy on downstream SuperGLUE tasks than the model trained by AdamW for the same number of steps.", "questions": "You said in the text that \"the gap between Sophia and Adam with 100K steps increases as the model size increases\", but the results in Figure 4 do not seem to support this claim --- the loss reached by Adam in 200K steps was reached by Sophia in about 160K steps in the 1.5B model, and similarly for the 6.6B model.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The novel contribution of this paper is to combine diagonal Hessian estimates (used by others, notably Adahessian), with exponential moving average of the Hessian and per parameter clipping of the Hessian (both also used before). The clipping allows them to reduce the effect of Hessian noise, and allows them to compute the Hessian infrequently, so the overhead of Hessian computation is not very large. The EMA improves stability of the Hessian computation.\n\nThe experiment with various versions of GPT seems quite clear, in all cases Sophia reached the same validation loss as Adam in fewer steps.", "weaknesses": "The main weakness is that the experiments are only in one problem domain --- decoder only LLMs of various sizes, although this is indeed a big application. Other problems and models should also be considered to see if the optimizer produces good results on different tasks.\n\nThe section 2.3 on GNB was not so clear, perhaps a more complete explanation in the appendix would help.\n\nThere is only weak theoretical justification for the optimizer, so it is justifiably relegated to the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sQ400UEv6I", "reviewer": "Reviewer_knid", "summary": "The authors proposed Sophia, a class of optimization methods making use of second-order information without dramatically increasing the computational cost. The paper focuses on two instantiations: Sophia-H and Sophia-G. In each algorithm, the second-order information comes from an estimation of the Hessian diagonal, which is updated by the moving average of the gradients divided by the moving average of the estimated Hessian, followed by a novel element-wise clipping procedure. In extensive numerical experiments, Sophia outperformed Adam by a 2x speed-up, making Sophia the first second-order optimizer achieving a speed-up on decoder-only large language models in wall-clock time or total compute.", "questions": "Adam is essentially RMSProp + momentum, where RMSProp is a method for adjusting the learning rate.\nIf my understanding is correct, Sophia is an entrywise learning rate scaling method, so a natural question is: have you considered combining Sophia with momentum?\n\nHow well does Sophia work for (potentially parameter-efficient) fine-tuning tasks instead of training a foundation model from scratch?\n\nThe idea of estimating the Hessian diagonal is not new, and both estimators proposed in the paper are not super novel. However, no second-order optimizer has performed as well as Sophia according to your experiments. How do you explain the success of Sophia? Is it due to regularization via gradient clipping?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper is well-organized and well-written. The contributions of this paper are significant. It is an interesting idea to use clipping to regularize the update size, in contrast to some prevalent choices like backtracking line-search. The method of updating the Hessian matrix, in which the matrix is updated every $k$ steps rather than every step, offers a tradeoff between Hessian estimation accuracy and computational cost. Moreover, Equation 10 provides a practical implementation of the GNB estimator. The paper also has extensive numerical experiments and some theoretical analysis, demonstrating the benefits of Sophia over other competitive methods. Furthermore, according to Figures 6 and 7, Sophia is insensitive to the choice of hyperparameters.", "weaknesses": "In the nonconvex setting, saddle points because very attractive to Newton-like optimizers, but it is non-obvious if Sophia can converge to a local minimum instead of a saddle point. While the authors provide an intuitive argument in Figure 2, it is not a rigorous one. The saddle-free Newton method [1] might be of some interest: it can additionally leverage the absolute curvature information when the Hessian diagonal is negative, in which case we also want to move fast.\n\n[1] Identifying and attacking the saddle point problem in high-dimensional non-convex optimization (https://arxiv.org/abs/1406.2572)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "azOx90PZqe", "reviewer": "Reviewer_Gbmx", "summary": "The paper proposes Sophia, a scalable optimizer with diagonal pre-conditioner and coordinate-wise clipping. Two estimators of the diagonal Hessian, Hutchinson and Gauss-Newton-Bartlett, are discussed. The authors illustrate the intuition behind Sophia via a simple yet convincing example. Some theoretical results are also derived. Numerical experiments are conducted on the pre-training of large language models to demonstrate the significant speed-up over other optimizers like AdamW and Lion.", "questions": "Overall, I think this is an excellent work. Though the theoretical results are weak, the numerical performance of Sophia is significant and impressive. Here are several questions for clarification and further improvement.\n\n- The authors mention that there is an efficient implementation in PyTorch and JAX for Hessian-vector product. What is the actual computational cost of this Hessian-vector product, compared with the computation of the (stochastic) gradient?\n\n- Does the 2x speedup still exist when considering larger models typically with billions of parameters? Can authors add more explanation for Figure 10? I am also curious about the performance of Sophia on other models like LLaMA and BERT.\n\n- This question might be a little beyond the scope of this work. Does Sophia still work well on other tasks beyond pre-training, such as fine-tuning? Can the authors provide some comment on this?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "Sophia proposes several enhancements on top of usual first-order optimizers. The use of diagonal Hessian exploits the curvature to accelerate convergence while keeping per-iteration cost comparable to first-order methods. Furthermore, the coordinate-wise clipping safeguards the iterates from unreliable second-order information and improves stability. The example and discussion in Section 2.1 illustrate the intuition behind Sophia in a very clear manner. The numerical experiments are comprehensive and the performance of Sophia looks quite strong.", "weaknesses": "- The writing and presentation have room for improvement. For example, the use of $L()$ and $l()$ are kind of messy. And function $\\text{clip}()$ is used before its definition. \n\n- There is a huge gap between Sophia and the theoretical results derived in the paper. The theoretical guarantee is for an algorithm which is (almost) completely different with Sophia under a fairly strong assumption.\n\n- Further studies are required to show if Sophia's speedup still exists on real large language models with billions of parameters and models other than GPT. See Questions for more discussion on this point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)