Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (2305.14342v4)

Abstract: Given the massive cost of LLM 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.

• The paper introduces Sophia, a scalable second-order optimizer that uses diagonal Hessian estimates to reduce computational costs in language model pre-training.
• Sophia utilizes both Hutchinson’s and Gauss-Newton-Bartlett estimators to efficiently approximate curvature, lowering training iterations and runtime.
• Experimental results demonstrate a 2x speedup over Adam, highlighting improved scalability and robustness across various model sizes.

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

This essay explores the mechanics and implications of "Sophia: A Scalable Stochastic Second-order Optimizer for LLM Pre-training" (2305.14342). The paper presents Sophia, a second-order optimization algorithm designed to enhance the efficiency of LLM pre-training. The optimizer leverages the advantages of second-order methods while maintaining computational efficiency akin to first-order methods.

Introduction

Sophia addresses the burgeoning challenge of high computational costs in LLM pre-training. Traditional optimizers like Adam have dominated the landscape due to their balance between computational demand and performance. However, the intrinsic limitations of first-order methods often lead to inefficiencies as model sizes and datasets grow. Sophia introduces a second-order approach that utilizes a lightweight estimate of the diagonal Hessian, thus offering a significant reduction in both the number of iterations and total computational resources required.

Methodology

Sophia employs a diagonal Hessian as a pre-conditioner, updated every few iterations to control computational overheads. This approach minimizes the update variances and effectively addresses sharp curvature issues across different parameter dimensions. The optimizer's update rule is defined by the equation:

$$\theta_{t+1} = \theta_t - \eta_t \cdot \text{clip}(m_t / \max\{\gamma \cdot h_t, \epsilon\}, 1)$$

with $\theta_t$ being the model parameters at step $t$, $m_t$ the moving average of gradients, $h_t$ the estimated Hessian diagonal, and $\gamma$, $\epsilon$ tuning parameters to prevent extreme updates.

Figure 1: Comparison of numbers of steps to reach the same validation loss. Across all model sizes, Sophia achieves significant speedup.

The algorithm supports two Hessian estimators: Hutchinson's unbiased estimator, which provides an unbiased diagonal approximation using Hessian-vector products, and the Gauss-Newton-Bartlett (GNB) estimator, which leverages structure in neural loss functions for a more deterministic diagonal approximation.

Experimental Results

The paper demonstrates that Sophia achieves a 2x speed-up over Adam in the number of steps, total computation, and wall-clock time, maintaining performance across model sizes from 125M to 1.5B parameters.

Figure 2: Validation loss on OpenWebText. Compared to AdamW and Lion, Sophia achieves lower loss across all model sizes.

One key finding is the scaling law: as model sizes increase, the gap in efficiency between Sophia and traditional first-order optimizers widens, demonstrating enhanced scalability. The optimizer's robustness is further highlighted in its ability to maintain stability against hyperparameter variations, a typical challenge in large-scale training settings.

Theoretical Insights

Sophia's effectiveness in adapting to heterogeneous curvatures across parameter dimensions offers a sustainable path for optimizing training processes without the extensive overhead typically associated with second-order methods. Through rigorous theoretical analysis, the authors provide runtime bounds that do not depend heavily on local condition numbers or curvature extrema, ensuring convergence efficiency.

Practical Implications

Practically, Sophia can be integrated into existing training pipelines with minimal changes to the architecture or computational framework. It leverages auto-differentiation frameworks like PyTorch and JAX to efficiently compute Hessian-vector products and other necessary operations, making it accessible for a wide array of model configurations and training infrastructures.

Conclusion

Sophia represents a noteworthy advancement for large-scale LLM training, marrying the precision of second-order methods with the efficiency required for practical deployment. Its design demonstrates that sophisticated optimization does not necessitate prohibitive computational costs, paving the way for more efficient training of increasingly larger models.

Sophia's contribution is significant for those seeking to enhance the scalability of LLM training with practical deployment under realistic computational constraints. As the field progresses, incorporating such advanced optimization strategies is likely to be crucial in addressing future challenges in AI scalability and efficiency.