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

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.

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

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

Rating

8: accept, good paper

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

2 fair

Contribution

4 excellent

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.

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. [1] Identifying and attacking the saddle point problem in high-dimensional non-convex optimization (https://arxiv.org/abs/1406.2572)

Questions

Adam is essentially RMSProp + momentum, where RMSProp is a method for adjusting the learning rate. If my understanding is correct, Sophia is an entrywise learning rate scaling method, so a natural question is: have you considered combining Sophia with momentum? How well does Sophia work for (potentially parameter-efficient) fine-tuning tasks instead of training a foundation model from scratch? The 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?

Rating

8: accept, good paper

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

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

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. The 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. The section 2.3 on GNB was not so clear, perhaps a more complete explanation in the appendix would help. There is only weak theoretical justification for the optimizer, so it is justifiably relegated to the appendix.

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.

Rating

6: marginally above the acceptance threshold

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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.

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

Questions

Thank you for an overall interesting and nicely written paper. I'd be interested in understanding the following. **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. **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. **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.) **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.

Rating

8: accept, good paper

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Ethics concerns

The paper is concerned with LLMs, particularly their fast pre-training. I don't have any specific concerns, but given that LLMs constitute an emerging technology with several nascent ethical issues, I recommend this paper for an ethics review.