PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hong Liu, Zhiyuan Li, David Hall, Percy Liang, Tengyu Ma

arXiv:2305.14342v4cs.LGcs.CLmath.OC

TL;DR

LLM pre-training is costly, while existing second-order optimizers often add too much per-step overhead. Sophia uses clipped stochastic diagonal-Hessian preconditioning with infrequent estimates, and achieves a 2x speed-up over AdamW in language-model training while retaining nearly the same per-step cost.

  • Problem

    LLM pre-training has high time and monetary cost, while sophisticated Hessian-based optimizers often incur excessive per-step overhead.

  • Method

    Sophia uses an inexpensive stochastic diagonal-Hessian estimate as a pre-conditioner, divides gradient EMAs by it, clips updates, and estimates curvature every k = 10 steps.

  • Results

    2x speed-up over AdamW in number of steps, total compute, and wall-clock time is achieved on GPT language modeling.

  • Takeaways & Limitations

    Sophia adapts to heterogeneous parameter curvatures, with a simplified runtime bound independent of the local condition number and worst-case curvature.

  • Takeaways & Limitations

    Vanilla second-order information can be unreliable when curvature changes rapidly or becomes negative in non-convex landscapes, motivating positive-curvature preconditioning and clipping.

Abstract

from arXiv · show

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, Second-order Clipped Stochastic Optimization, 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. Theoretically, we show that Sophia, in a much simplified setting, adapts to the heterogeneous curvatures in different parameter dimensions, and thus has a run-time bound that does not depend on the condition number of the loss.

1 Introduction

Sophia targets the time and cost bottleneck of LLM pre-training with a lightweight second-order optimizer. It combines diagonal-Hessian preconditioning with clipping and infrequent estimation to reduce optimization steps while controlling update risks.

  • Motivation: LLM pre-training is a major efficiency bottleneck because massive datasets and model sizes require hundreds of thousands of parameter updates.PaLM, for example, was trained for two months on 6144 TPUs at a cost of 10 million dollars.
  • Method: Sophia uses a stochastic diagonal-Hessian estimate as a pre-conditioner, dividing an exponential moving average of gradients by the Hessian estimate and clipping the result element-wise.The optimizer can use Hutchinson or Gauss-Newton-Bartlett estimation, with Hessian estimates updated intermittently.
  • Results: 50% fewer steps, total compute, and wall-clock time are required to reach the same validation pre-training loss as Adam on GPT models.The reported comparison covers GPT-2 and GPT NeoX models across multiple sizes; the figure caption describes the baseline as AdamW.
  • Method: Sophia estimates the Hessian every k = 10 steps, reducing average overhead while its clipping limits worst-case updates under inaccurate or rapidly changing curvature.The safeguard also addresses the non-convexity that can cause vanilla Newton’s method to converge to maxima or saddle points.
  • Theory: Sophia’s simplified theoretical runtime bound does not depend on the local condition number or worst-case curvature.This reflects adaptation to heterogeneous curvatures across parameter dimensions.

2 Method

Sophia addresses heterogeneous curvature by using a lightweight diagonal-Hessian preconditioner with per-coordinate clipping, while reducing the risks and overhead of second-order optimization. Its Hessian estimates are updated intermittently, and its design combines curvature adaptation with safeguards for non-convexity and changing curvature.

  • Motivation: Modern loss landscapes can have dispersed curvature across parameter dimensions, limiting the effectiveness of shared-step methods such as GD and Adam.For a 125M-parameter GPT-2, the positive diagonal Hessian entries are dispersed.
  • Motivation: GD must use a learning rate suited to the sharpest direction, slowing convergence in flatter directions, while SignGD can bounce in sharp directions.In the toy example, θ[2] converges slowly and θ[1] reaches the valley before bouncing.
  • Motivation: Vanilla Newton’s method can fail under negative curvature or rapidly changing Hessians, converging to maxima or saddle points instead of minima.Sophia therefore uses positive-curvature preconditioners and clipping to mitigate unreliable second-order information.
  • Sophia: Sophia uses an exponential moving average of gradients divided by an EMA of diagonal Hessian estimates, with the Hessian EMA updated every k steps.The optimizer updates the gradient EMA each step and refreshes the Hessian estimate intermittently.
  • Sophia: Per-coordinate clipping limits update magnitudes, using stochastic momentum SignSGD as a backup when curvature estimates are negative or very small.The worst-case update size is bounded by ηρ, improving stability when second-order information is misleading.
  • Diagonal Hessian Estimators: Sophia offers two lightweight diagonal-Hessian estimators: Hutchinson’s estimator and the Gauss-Newton-Bartlett estimator.Hutchinson uses a Hessian-vector product without forming the full Hessian, while GNB exploits cross-entropy structure and produces a positive semi-definite Gauss-Newton estimate.

3 Experiments

Experiments evaluate Sophia against AdamW, AdaHessian, and Lion on language modeling, using step-matched validation loss, compute, wall-clock time, stability, and ablations. Across model sizes, Sophia achieves faster convergence with low overhead and improved stability.

  • 3.3 Results: 2x faster than AdamW in number of steps, total compute, and wall-clock time across all evaluated model sizes.The paper reports that the scaling law is also in favor of Sophia over AdamW.
  • 3.1 Experimental Setup: Sophia is evaluated on GPT-2 models from 125M to 770M parameters and GPT NeoX models at 1.5B and 6.6B parameters.Experiments use OpenWebText and the Pile, comparing primarily against AdamW, with AdaHessian and Lion as additional baselines.
  • 3.3 Results: 0.04 smaller validation loss is achieved by Sophia-H and Sophia-G on the 355M model with the same number of steps.On the 770M model, Sophia-H and Sophia-G achieve a 0.05 smaller validation loss with the same 100K steps.
  • 3.4 Analysis: Less than 5% overall wall-clock overhead is incurred compared with AdamW, while Hessian computation accounts for 6% of total compute.The diagonal Hessian estimate uses a reduced batch size every 10 steps, and Sophia has the same memory cost as AdamW.
  • 3.4 Analysis: Sophia-H seldom triggers gradient clipping, whereas AdamW and Lion trigger it in more than 10% of steps under the same threshold.Sophia-H also does not require the attention re-parameterization trick used to stabilize AdamW and Lion.
  • 3.4 Analysis: Sophia's Hessian update frequency, diagonal Hessian pre-conditioner, and element-wise clipping are examined in ablation studies.The Hessian-frequency study compares k = 1, 10, and 100 on a 30M GPT-2 model over 100K, 200K, and 400K steps.

4 Theoretical Analysis

The theoretical analysis studies a deterministic version of Sophia under convexity and Hessian-continuity assumptions. It derives runtime behavior that adapts to heterogeneous curvatures without depending on the local condition number.

  • Runtime bound: Sophia’s runtime bound does not depend on the local condition number or worst-case curvature, reflecting adaptation to heterogeneous parameter curvatures.
  • Assumptions: Sophia’s deterministic analysis assumes a twice continuously differentiable, strictly convex loss with a unique minimizer.
  • Assumptions: The Hessian-continuity assumption bounds the ratio between nearby Hessians within a neighborhood of radius R.
  • Deterministic analysis: The analysis uses the full Hessian in eigenspace because diagonal preconditioning may not align with the natural coordinate system of general functions.
  • Runtime bound: The runtime consists of burn-in before entering a local region followed by exponential error decay, yielding logarithmic dependence on 1/ϵ.
  • Comparison: On simple quadratic functions, simplified Adam, or SignGD, has a convergence rate that depends on the condition number.

5 Related work

The related-work discussion contrasts adaptive first-order optimizers with second-order methods in deep learning. Adam dominates Transformer optimization, while second-order approaches offer curvature information but traditionally incur greater complexity.

  • Stochastic Adaptive First-order Optimizers in Deep Learning: Adaptive first-order optimization progressed from RProp and AdaGrad to RMSProp and Adam, which became dominant for Transformer training.
  • Stochastic Adaptive First-order Optimizers in Deep Learning: Lion is a searched-for adaptive first-order method that is substantially faster than Adam on vision Transformers and diffusion models but offers limited speed-up on LLMs.
  • Stochastic Adaptive First-order Optimizers in Deep Learning: Adam outperforms SGD by a large margin on language modeling tasks with Transformers, unlike the pattern reported for some vision tasks.
  • Stochastic Adaptive First-order Optimizers in Deep Learning: The related work situates AdaFactor among methods used to train large Transformer models.

6 Conclusion

The paper introduces Sophia as a scalable second-order optimizer for language-model pre-training. Sophia converges in fewer steps at nearly unchanged per-step cost and achieves a 2x speed-up over AdamW across several efficiency measures.

  • 6 Conclusion: Sophia is introduced as a scalable second-order optimizer for language model pre-training.
  • 6 Conclusion: Sophia converges in fewer steps than first-order adaptive methods while maintaining almost the same per-step cost.
  • 6 Conclusion: 2x speed-up compared with AdamW is achieved in number of steps, total compute, and wall-clock time on GPT language modeling.

A Additional Experiment Results

Additional experiments examine clipping dynamics and performance across training durations. Sophia’s clipping proportion approaches 60%, and it outperforms AdamW and Lion across different total step counts.

  • Training dynamics: The proportion of Sophia parameters undergoing clipping approaches 60% after the initial training stage.
  • Training dynamics: Figure 9 visualizes the proportion of clipped parameter updates and the ℓ2 norm of the EMA of the diagonal Hessian.
  • Results with different number of steps: Across different choices of total training steps, Sophia outperforms AdamW and Lion by a large margin.

B Additional Experiment Details

Experiments tune optimizer hyperparameters through grid searches on smaller models before transferring selected settings to larger models. Sophia uses k=10 Hessian-estimation intervals as a balance between convergence speed and computation overhead.

  • Hyperparameter tuning: Baseline hyperparameters include peak learning rate, β1, β2, and weight decay, with most tuned by grid search on a 30M GPT-2 model.Peak learning rate is tuned separately across model sizes.
  • Sophia settings: Sophia transfers β1 and β2 tuned on a 30M model to larger models, then uses β1 = 0.96, β2 = 0.99, and k = 10.The reported tuning procedure fixes γ = 0.01 and k = 10 initially, then evaluates k choices.
  • Sophia settings: k = 10 offers the best balance between convergence speed and computation overhead compared with k = 1 or k = 100.

B.2 Model and Implementation Details

The experiments use GPT-family models trained with distributed methods, gradient accumulation, and mixed precision across GPUs and TPUs. The setup also includes a scaling-related architectural adjustment for some models, though Sophia does not require that trick.

  • Models: The study uses GPT-2 models spanning small, medium, and large configurations, plus 30M and 540M models.The 30M model supports efficient hyperparameter grid search.
  • Data and architecture: GPT-2 models are trained on OpenWebText with the GPT-2 tokenizer, GELU activations, and disabled bias and dropout.The training set contains 9B tokens.
  • Prompt examples: The evaluation context includes binary and three-way reasoning prompts, including entailment-style True/False/Neither decisions.
  • Training infrastructure: Distributed data parallelism, gradient accumulation, and bfloat16 enable large batch training across A5000, A100, and TPU systems.The 125M and 355M models use 10 A5000 GPUs, while the 770M model uses 8 A100 GPUs.
  • Stability: Scaling attention by the inverse of layer index addresses instability for larger learning rates, but Sophia does not need this trick.

B.3 Downtream Evaluation

Downstream evaluation uses few-shot SuperGLUE testing with fixed prompting and greedy decoding. Accuracy is measured on validation examples averaged across multiple exemplar selections, with prompt formats illustrated separately.

  • Evaluation protocol: The models are evaluated on four SuperGLUE subtasks using two-shot prompting.
  • Evaluation protocol: Greedy decoding is used for downstream evaluation, with the subtask prompt designs shown in Figure 11.
  • Evaluation protocol: Prompts contain an instruction followed by two training examples, and validation accuracy is averaged over five exemplar selections.Examples are sampled from the training split.
  • Learning-rate tuning: Figure 12 presents results from peak learning-rate tuning.

C Limitations

The paper’s evaluation remains limited in model scale, downstream coverage, and domain breadth because of computational-resource constraints. Broader comparisons and applications are left for future work.

  • Scaling scope: Sophia scales to 6.6B-parameter models, but comparisons with AdamW and Lion at larger sizes are omitted because of limited resources.
  • Evaluation scope: Downstream evaluation covers pre-training validation losses and only four SuperGLUE subtasks.The authors cite limited model capabilities at this scale as part of the reason for restricted downstream evaluation.
  • Domain scope: Applications beyond large language modeling, including computer vision, reinforcement learning, and multimodal tasks, are left for future work.

D Theoretical Analyses: Details of Section 4

The analysis develops a two-phase convergence argument under strict-convexity and Hessian regularity assumptions, using auxiliary lemmas about bounded trajectories, descent, and clipped updates. It also notes that the assumptions do not require a Lipschitz Hessian.

  • Convergence analysis: Theorem 4.3 analyzes convergence in two phases: an initial loss decrease followed by exponential error decay.The first phase decreases loss to a threshold involving µ, ρ, and η.
  • Auxiliary lemmas: Under the assumptions, the loss diverges as ∥θ∥2 approaches infinity, supporting boundedness arguments for the analyzed trajectories.This property is stated in Lemma D.1 and used in the ODE-based proofs.
  • Assumptions: The theoretical assumptions require only multiplicative Hessian control within a constant-radius neighborhood, not Hessian Lipschitzness.This is stated explicitly as a scope of Assumption 4.2.
  • Auxiliary lemmas: Strict convexity and Hessian lower bounds are used in contradiction arguments to establish loss and distance inequalities around θ∗.These arguments show that moving toward the minimizer cannot produce the contradictory loss relationships assumed in the proofs.
  • Auxiliary lemmas: The ODE construction yields a trajectory whose gradient follows ∇L(θ(t)) = (1 − t)∇L(θ) and reaches θ∗ at t = 1.The trajectory remains within the relevant bounded region, and the loss does not increase along it.
  • Descent analysis: The descent analysis studies clipped preconditioned updates and combines the Descent Lemma with other lemmas to establish the convergence claims.The clipping operation is explicitly applied to the preconditioned gradient direction, while the proof uses local Hessian control.

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

This section establishes a lower bound for signGD on a two-dimensional quadratic loss. The bound shows that its convergence rate must depend on the loss condition number.

  • Lower-bound construction: The quadratic loss Lµ,β is used to show that signGD’s convergence rate depends on the condition number β/µ.The construction provides the setting for the lower-bound theorem.
  • Lower-bound theorem: Theorem D.12 assumes signGD reaches loss at most ϵ at consecutive steps from all initializations within loss level ∆.Under these assumptions, the theorem imposes a lower bound on the required number of steps.
  • Proof: The proof uses two initializations and the fact that each coordinate changes by η between consecutive signGD iterates.This coordinate-wise movement yields the condition-number-dependent lower bound.
Loading 2305.14342v4…