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

• 50% fewer steps, 50% less total compute, and 50% less wall-clock time than Adam achieve the same validation pre-training loss; Figure 1 reports 2x speedup over AdamW.Source: Figure 1 · p. 2, §1 · p. 2
• At 100K steps, Sophia on 540M matches Adam’s validation loss on 770M, whose training time and inference cost are each 40% higher; Sophia’s gap grows with model size.Source: §1 · p. 2, Figure 1 · p. 2
• Sophia refreshes stochastic diagonal-Hessian estimates every k=10 steps, averages gradients and curvature, divides then elementwise clips updates; either estimator adds 5% average per-step overhead.Source: §1 · pp. 2–3
• The curvature preconditioner penalizes updates more in sharp than flat dimensions, while clipping limits worst-case updates and guards against inaccurate or rapidly changing Hessians and non-convexity.Source: §1 · pp. 3–4
• In simplified convex analyses, Sophia’s runtime bound is independent of both local condition number and worst-case curvature, reflecting adaptation across heterogeneous parameter curvatures.Source: §1 · p. 4

2.1 Motivations

• Sophia’s diagonal-Hessian preconditioner adapts updates to heterogeneous curvature, while per-coordinate clipping limits unreliable updates under negative curvature or rapidly changing Hessians.The clipping bounds update size; when curvature is negative or too small, the update defaults to stochastic momentum SignSGD.Source: §2.1 · p. 5, §2.2 · p. 6

2.2 Sophia: Second-order Clipped Stochastic Optimization

• Sophia updates an EMA of gradients using an EMA of diagonal-Hessian estimates, computing Hessian estimates every k steps (k = 10 in the implementation).The Hessian EMA is updated only when a new estimate is computed, reducing estimation frequency while smoothing noisy estimates.Source: §2.2 · p. 6
• Sophia uses only positive diagonal-Hessian entries and clips each preconditioned coordinate, controlling worst-case update size when curvature estimates are unreliable.A small ϵ avoids division by zero; the clipping threshold ρ sets the per-coordinate limit.Source: §2.2 · p. 6

2.3 Diagonal Hessian Estimators

• Hutchinson estimates the Hessian diagonal unbiasedly using a random-vector Hessian-vector product; GNB estimates the Gauss-Newton term and yields a positive-semidefinite diagonal estimate.GNB is biased for the Hessian diagonal, but its nonnegative estimate ensures the preconditioned update is a descent direction.Source: §2.3 · pp. 7, 9

3.1 Experimental Setup

• Experiments trained GPT-2 models from 125M to 770M and GPT NeoX models at 1.5B and 6.6B parameters on OpenWebText and the Pile, comparing Sophia with AdamW, AdaHessian, and Lion.For models other than 30M, most hyperparameters were held fixed from the 30M configuration; peak learning rates were searched separately for 125M and 355M and increased cautiously for larger models.Source: §3.1 · pp. 9–10

3.2 Evaluation

• The 2x-speedup test compares an optimally tuned T-step baseline with a candidate run for T/2 steps, requiring the candidate’s loss to be no greater; main comparisons used 400K versus 200K and 200K versus 100K steps.Baseline hyperparameters must be tuned for the specified T-step budget because learning-rate schedules and optimal settings depend on the total steps.Source: §3.2 · pp. 10–11

3.3 Results

• Across model sizes, Sophia achieves a 2x speedup over AdamW in steps, compute, and wall-clock time; at 100K steps, its 355M and 770M validation losses were lower than AdamW’s.On GPT-2 Medium (355M), Adam: 2.691, Sophia-H: 2.645, and Sophia-G: 2.627; on GPT-2 Large (770M), AdamW: 2.613, Sophia-H: 2.554, and Sophia-G: 2.524.Source: §3 · p. 9, Figure 4 · p. 12, §3.3 · p. 12
• Sophia’s validation-loss advantage grows with model scale, and its 200K-step models achieve few-shot SuperGLUE accuracy comparable to AdamW’s 400K-step models.With equal pre-training steps, Sophia models also have better few-shot accuracy on most subtasks for GPT-2 Medium, GPT-2 Large, and GPT NeoX 1.5B.Source: §3.3 · p. 12, Figure 6 · p. 13

3.4 Analysis

• Hessian computation every 10 steps accounted for 6% of total compute, while Sophia’s wall-clock overhead was less than 5% versus AdamW and its two optimizer states matched AdamW’s memory cost.The 30M-model choice γ=0.01 and β2=0.99 transferred to GPT-2 models from 125M to 770M; tested Sophia-H hyperparameter combinations had similar performance.Source: §3.4 · p. 13
• Sophia-H seldom triggered gradient clipping, versus more than 10% of steps for AdamW and Lion, and needed no attention rescaling; without rescaling, their stable learning-rate ceiling was 1.5e-4 versus 3e-4 with it.The clipping comparison used a threshold of 1.0; the learning-rate comparison was on GPT-2 Medium (355M).Source: Figure 7 · p. 13, §3.4 · p. 14

3.5 Ablation Study

• At k=10, Sophia-G offered a better compute–validation-loss tradeoff than k=1 or 100; GNB beat empirical Fisher, Sophia-H beat AdaHessian, and adding GNB to coordinate clipping improved on clipping alone.For k=1, computational overhead was 50% and convergence by compute was worse than k=10; without element-wise clipping, AdaHessian diverged at k=2 and GNB at k=5.Source: §3.5 · pp. 14–15, Figure 8 · p. 14

Theoretical Analysis

• The analysis assumes a twice continuously differentiable, strictly convex loss with a minimizer, and bounded Hessian variation within a neighborhood of radius R.The Hessian-ratio bound is stated as a constant 2 in the supplied passage.Source: §4 · p. 15
• For general convex functions, the analyzed update uses the full Hessian, clips element-wise in its eigenspace, and transforms the update back to the original coordinates.The full Hessian is used because a diagonal preconditioner may fail when curvature is not aligned with the natural coordinate system; convexity makes the Hessian positive semidefinite, so the original Sophia max with ϵ is unnecessary.Source: §4 · p. 15
• Sophia’s runtime bound is independent of the condition number and worst-case curvature; after a burn-in period, error decays exponentially, while simplified Adam (SignGD) depends on condition number even for simple quadratics.The theorem’s bound includes a burn-in term for reaching a local region and a term logarithmic in 1/ϵ; the supplied theorem passage is incomplete, so no full formula is reported.Source: §4 · pp. 15–16

Stochastic Adaptive First-order Optimizers in Deep Learning

• Adaptive first-order methods developed from RProp through AdaGrad and RMSProp to Adam, which became dominant in deep learning; later work proposed Adam variants and Lion.Lion was discovered through a search over adaptive first-order algorithms and described as an improved sign-momentum SGD (Chen et al., 2023).Source: §5 · p. 16

Second-order Optimizers in Deep Learning

• Deep-learning second-order work includes diagonal-Hessian preconditioning (Becker & Le Cun, 1988), conjugate-gradient and Fisher/Gauss–Newton approximations, neural-structure methods, and EMA diagonal-Hessian preconditioners.Examples include Martens et al. (2010), Pascanu & Bengio (2013), Martens & Grosse (2015), and Yao et al. (2021); Jahani et al. (2021).Source: §5 · p. 16
• Prior second-order optimizers may incur frequent-Hessian computation costs or depend on specific architectures or hardware; the authors knew of no reported total-compute speed-up on large language models.Examples include CPU-offloaded Hessian computation and Fisher approximation requiring ResNets and very large batches (Anil et al., 2020; George et al., 2018).Source: §5 · p. 16

Gradient Clipping

• Sophia is presented as the first second-order method to clip updates coordinate-wise, addressing changing Hessians and inaccurate Hessian approximations; earlier work studied global and per-coordinate gradient clipping.Global clipping is standard in language-model pre-training, while per-coordinate clipping can function as adaptivity (Crawshaw et al., 2022).Source: §5 · p. 16

Optimization Algorithms in LM Pre-training

• Adam with decoupled weight decay dominates language-model pre-training and outperforms SGD on Transformer language modeling, while AdaFactor and per-layer-parameter-norm-scaled Adam are alternatives.AdaFactor is a low-rank version of Adam (Shazeer & Stern, 2018); You et al. (2019) proposed scaling Adam’s update by per-layer parameter norm.Source: §5 · pp. 16–17

• Sophia is a scalable second-order optimizer for language-model pre-training that converges in fewer steps than first-order adaptive methods at almost the same per-step cost.Source: §6 · p. 17
• 2x speed-up over AdamW in steps, total compute, and wall-clock time on GPT language modeling.Source: §6 · p. 17

A Additional Experiment Results

• After the initial stage, the EMA diagonal-Hessian norm grows steadily while clipping approaches 60%; effective clipping is generally 50–90%, so users are advised to tune γ.The paper says per-coordinate clipping corroborates its importance in Sophia.Source: A Additional Experiment Results · p. 23, Figure 9 · p. 23
• Across different step-count choices, Sophia outperforms AdamW and Lion by a large margin.Figure 10 presents the results for runs with different total numbers of steps.Source: A Additional Experiment Results · p. 23, Figure 10 · p. 23

B.1 Hyperparamter Tuning

• Baseline tuning grid-searched peak learning rate, β1, β2, and weight decay, using a 30M GPT-2 for 50K steps except for model-specific peak learning rates.The β1, β2, and weight-decay grids were [0.8, 0.9, 0.95, 0.96, 0.99], [0.9, 0.95, 0.98, 0.99, 0.995], and [0.05, 0.1, 0.2, 0.5], respectively.Source: §B.1 · p. 23
• AdamW worked best with β1 = 0.9 and β2 = 0.95, while Lion worked best with β1 = 0.95 and used β2 = 0.98 as recommended by Chen et al. (2023).AdamW was sensitive to β1 but not β2; β2 = 0.95 was chosen as a commonly used LLM pre-training setting.Source: §B.1 · pp. 23–24
• Weight decay 0.1 worked best for AdamW, while 0.2 worked best for Lion and Sophia.Source: §B.1 · p. 24
• For 125M and 355M models, peak learning rates were grid-searched; for larger models, the largest stable rate was selected from a candidate list.The peak-learning-rate search results are provided in Table 2; fine-grained trials on 355M GPT-2 and 1.5B GPT Neo-X supported the listed choices.Source: §B.1 · p. 24, Table 2 · p. 24
• Sophia used β1 = 0.96, β2 = 0.99, ϵ = 1e-12, and k = 10, with β1 and β2 selected on a 30M model before reuse at larger sizes.Sophia was not sensitive to β2.Source: §B.1 · p. 24
• With γ = 0.01 fixed, k = 10 balanced convergence speed and computation overhead better than k = 1 or k = 100; tuning γ gave 0.01 for Sophia-H and 0.05 for Sophia-G.γ was selected to put the proportion of unclipped update coordinates in the 10%–50% range, then jointly tuned with peak learning rate.Source: §B.1 · p. 24

B.2 Model and Implementation Details

• Experiments use GPT-2 small, medium, and large, plus 30M and 540M models for hyperparameter search and scaling-law visualization, respectively.Model specifications are provided in Table 2.Source: §B.2 · p. 24
• GPT-2 models use OpenWebText, the GPT-2 tokenizer, and nanoGPT’s train-validation split, with 9B training tokens and 4.4M validation tokens.The implementation uses GELU activations and disables bias and Dropout during pre-training.Source: §B.2 · p. 24
• GPT-2 runs use batch size 480 with distributed data parallel and gradient accumulation; 125M and 355M models use 10 A5000 GPUs, while 770M uses 8 A100 GPUs.All models are trained with bfloat16; the 770M run uses an AWS p4d.24xlarge instance.Source: §B.2 · p. 25
• GPT NeoX models of 1.5B and 6.6B parameters use the Pile and GPT NeoX tokenizers, with levanter and fully sharded data parallel on a TPU v3-128 slice.Gradient accumulation enables batch sizes of 512 for 1.5B and 1024 for 6.6B.Source: §B.2 · p. 25
• Above 355M parameters, AdamW and Lion became unstable near Radford et al. (2019)’s learning-rate choice; inverse-layer-index attention scaling addressed this, while Sophia did not need it.The paper notes that Sophia does not require the scaling trick (Section 3.4).Source: §B.2 · p. 25

• The evaluation covers four SuperGLUE subtasks using 2-shot prompting.Source: §B.3 · p. 25
• Each prompt contains an instruction followed by two examples sampled from the training split, and decoding is greedy.Source: §B.3 · p. 25
• Accuracy is reported on the validation split, averaged over five exemplar selections; example prompts are shown in Figure 11.Source: §B.3 · p. 25

Scaling up to larger models and datasets

• Despite demonstrating scalability to 6.6B-parameter models, Sophia is not compared with AdamW or Lion on larger models because of limited resources.The authors believe Sophia may be faster than both on larger models, citing scaling-law improvements and better pre-training stability.Source: C Limitations · p. 26

Holistic downstream evaluation

• Evaluation covers pre-training validation losses and only four SuperGLUE subtasks, partly because models at this scale lack capabilities such as in-context learning and mathematical reasoning.Source: C Limitations · p. 26

Evaluation on other domains

• The paper evaluates optimizers only for large language modeling and leaves computer vision, reinforcement learning, multimodal tasks, and other models to future work because of limited computation resources.Source: C Limitations · p. 26

Theorem 4.3

• Theorem 4.3’s analysis has two phases: reduce the loss to a threshold, then obtain exponential error decay.The theorem is described as a combination of the descent lemma and Lemmas D.9 and D.11.Source: D Theoretical Analyses: Details of… · p. 26

Supporting lemmas

• Under Assumption 4.1, the loss tends to infinity as the parameter norm tends to infinity.This coercivity property is used to keep the path in Lemma D.4 bounded.Source: D Theoretical Analyses: Details of… · pp. 26–27
• Lemma D.4 constructs a path on [0, 1] whose gradient decays linearly to zero and whose endpoint is the minimizer.Along the path, the loss does not increase, which helps establish that the path exists through t = 1.Source: D Theoretical Analyses: Details of… · p. 27
• After the loss falls below µρ²/8, the iterates satisfy the condition for clipping not to occur.The proof invokes Lemma D.10 and then Lemma D.7 to show the relevant quantity is at most ρ for all subsequent iterations.Source: D Theoretical Analyses: Details of… · p. 32

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

• The loss Lµ,β is defined as a quadratic function on R2 with parameters µ and β.The supplied passage does not include the complete quadratic formula.Source: §D.1 · p. 32
• SignGD’s convergence rate on the two-dimensional quadratic loss must depend on the condition number β/µ.The theorem states a lower bound under which iterates from every initialization with loss at most Δ reach loss at most ϵ at both steps T−1 and T.Source: §D.1 · p. 32