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

要約

言語モデルの事前トレーニングに莫大なコストがかかることを考慮すると、最適化アルゴリズムの重要な改善は、トレーニングの時間とコストの大幅な削減につながるでしょう。
Adam とその亜種は何年もの間最先端であり、より洗練された 2 次 (ヘッセ行列ベース) オプティマイザーでは、ステップごとのオーバーヘッドが多すぎることがよくあります。
この論文では、事前調整器として対角ヘッセ行列の軽量推定を使用する、シンプルでスケーラブルな 2 次オプティマイザである Sophia、2 次クリップ確率的最適化を提案します。
更新は、勾配の移動平均を推定ヘッセ行列の移動平均で割った後、要素ごとにクリッピングされます。
クリッピングは最悪の場合の更新サイズを制御し、軌道に沿ったヘッセ行列の非凸性と急速な変化による悪影響を抑制します。
Sophia は、数回の反復ごとに対角ヘッセ行列を推定するだけであり、ステップごとの平均時間とメモリ オーバーヘッドは無視できます。
125M から 1.5B までのサイズの GPT モデルを使用した言語モデリングでは、Sophia は、ステップ数、総計算時間、実時間の点で Adam と比較して 2 倍の高速化を達成し、50% 少ないステップで同じ複雑さを実現します。
総計算量が減り、実測時間が短縮されます。
理論的には、非常に単純化された設定で、Sophia がさまざまなパラメーター次元の不均一な曲率に適応するため、損失の条件数に依存しない実行時の限界があることが示されます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google