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

Thanks for the feedback! We address questions below:

Q1: What is the actual computational cost of this Hessian-vector product, compared with the computation of the (stochastic) gradient?

Response: In practice, we observe that on a 1.5B model a Hessian-vector product with batchsize=64 costs approximately the same time as a stochastic gradient with batchsize=480.

Q2: Does the 2x speedup still exist when considering larger models typically with billions of parameters?

Response: So far 1.5B and 6.6B GPT NeoX are the largest models which we have resources to run. We still observe 2x speedup of Sophia over Adam on the 1.5B model, but we still need more time and compute to complete the 6.6B run with more steps.

Q3: Can authors add more explanation for Figure 10?

Response: 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.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 is close to the largest peak learning rate for AdamW in that setup. Note that with the unstable 1.5e-4 learning rate, AdamW is slower than Sophia-G before blowing up.

Q4: I am also curious about the performance of Sophia on other models like LLaMA and BERT.

Response: We provide results of mask language modeling with BERT in the table below. In terms of validation pre-training loss, Sophia is better than AdamW and Lion with less than 70% compute spent. We will provide results on T5 and contrastive pre-training in the next version.

Thanks for noting that "The contributions of this paper are significant" and "It is an interesting idea to use clipping to regularize the update size".

Q1: 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?

Response: We'd like to clarify that Sophia uses exponential moving average (EMA) of stochastic gradients in the update rule as in the second paragraph of Section 2.2 in page 4 of the manuscript. We will make this more prominent in the revision.

Q2: How do you explain the success of Sophia? Is it due to regularization via gradient clipping?

Response: As demonstrated in the ablation study (Section 3.5), element-wise update clipping and the Hessian estimators are both reasons for Sophia's success. Compared to AdaHessian, the Hutchinson estimator in Sophia-H has a better denoising effect. Compared with empirical Fisher, label resampling (equation 8) makes the GNB estimator of Sophia-G an unbiased estimator of the Gauss-Newton term. Element-wise clipping allows infrequent updates of the diagonal Hessian, which reduces the overhead of computing second-order derivatives.

Thanks for the feedback!

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

Response: We provide results of mask language modeling with BERT in the table below. In terms of validation pre-training loss, Sophia is better than AdamW and Lion with less than 70% compute spent.

Q2: "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.

Response: Thanks for raising this point! Note that experiments of 125M, 355M, and 770M in Figure 4 (a)(b)(c) are carried out on OpenWebText with GPT-2, while those of 1.5B and 6.6B are on the Pile with GPT NeoX. The difference in datasets and architecture means that we cannot build a scale law with all the 5 scales. ,as a result, don't have 1.5B and 6.6B data points in the scaling law figure (Figure 1 (d)). We will modify the sentence to "On OpenWebText, the gap between Sophia and Adam with 100K targeted steps is larger on 355M / 770M models than 125M models" to make it more precise in the revision.

Q3: The section 2.3 on GNB was not so clear, perhaps a more complete explanation in the appendix would help.

Response: We added an expanded version in the appendix.

Thanks for noting the "novelty and creativity" of the paper and that "the authors make second-order methods computationally tractable for pre-training LLMs, which is a big achievement". We answer the questions as follows:

Q1: The theoretical analysis seems 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. 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?

Response: We agree that Assumption E.2 requires the curvature to change at a reasonable rate. However, this assumption contains natural cases which do not satisfy the self-concordance assumptions [1] or bounded third-order derivatives assumptions [2] in Newton's method analysis. For example, $f(x) = \log(1+e^x)-x/2$ is not self-concordant (when x goes to infinity) but satisfies Assumption E.2. When x is very large, the hessian of $f$ is very small, but the gradient is close to 1. Newton's method will overshoot, but with clipping Sophia is fine with this case. For decomposable nonconvex functions, we provide analysis in Appendix G of the revision. We leave analysis for more general non-convex functions to future work.

[1] Boyd, S. P. and Vandenberghe, L. Convex optimization.
[2] Nesterov, Y. and Polyak, B. T. Cubic regularization of newton method and its global performance

Q2: 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?

Response: We did not observe a significant trend of the trace of Hessian of the pre-training loss of models trained with Sophia-G and AdamW and Lion as shown in the table below on 125M and 355M models pre-trained for 100K steps. In language model pre-training, we typically have strict data deduplication and pass one data point only once. Therefore the training loss and validation loss are almost the same and we are not overfitting training data. Performance on few-shot evaluation in Figure 5 also indicates the models trained with Sophia are transferable to downstream tasks.