SOAP: Improving and Stabilizing Shampoo using Adam

要約

深層学習の最適化タスクにおいて、高次のプレコンディショニング手法である Shampoo が Adam よりも有効であるという証拠が増えています。
ただし、Shampoo には、第一モーメント量と第二モーメント量の移動平均を更新するだけの Adam と比較した場合、追加のハイパーパラメーターと計算オーバーヘッドが含まれるという欠点があります。
この研究は、Shampoo (1/2 乗で実装) と Adafactor (Adam のメモリ効率の高い近似) との間の正式な接続を確立し、Shampoo が Shampoo のプリコンディショナーの固有基準で Adafactor を実行することと同等であることを示します。
この洞察は、より単純で計算効率の高いアルゴリズムの設計につながります。 $\textbf{S}$hampo$\textbf{O}$ ( $\textbf{P}$reconditioner の固有基底に $\textbf{A}$dam を使用)
石鹸）。
Shampoo の計算効率の向上に関して、最も簡単なアプローチは、Shampoo の固有分解の計算頻度を下げることです。
残念ながら、私たちの経験的結果が示すように、これはパフォーマンスの低下につながり、この頻度でさらに悪化します。
SOAP は、Adam と同じように、現在の (ゆっくりと変化する) 座標ベースで 2 番目の瞬間の移動平均を継続的に更新することで、この劣化を軽減します。
さらに、SOAP は回転された空間で Adam を実行するのと同等であるため、Adam と比較して追加のハイパーパラメーター (事前調整周波数) が 1 つだけ導入されます。
私たちは、360 メートルと 660 メートルのサイズのモデルを使用した言語モデルの事前トレーニングで SOAP を経験的に評価します。
大規模なバッチ方式では、SOAP は AdamW と比較して反復回数を 40% 以上削減し、実測時間を 35% 以上削減し、Shampoo と比較して両方のメトリクスで約 20% 改善しました。
SOAP の実装は https://github.com/nikhilvyas/SOAP で入手できます。

要約(オリジナル)

There is growing evidence of the effectiveness of Shampoo, a higher-order preconditioning method, over Adam in deep learning optimization tasks. However, Shampoo’s drawbacks include additional hyperparameters and computational overhead when compared to Adam, which only updates running averages of first- and second-moment quantities. This work establishes a formal connection between Shampoo (implemented with the 1/2 power) and Adafactor — a memory-efficient approximation of Adam — showing that Shampoo is equivalent to running Adafactor in the eigenbasis of Shampoo’s preconditioner. This insight leads to the design of a simpler and computationally efficient algorithm: $\textbf{S}$hampo$\textbf{O}$ with $\textbf{A}$dam in the $\textbf{P}$reconditioner’s eigenbasis (SOAP). With regards to improving Shampoo’s computational efficiency, the most straightforward approach would be to simply compute Shampoo’s eigendecomposition less frequently. Unfortunately, as our empirical results show, this leads to performance degradation that worsens with this frequency. SOAP mitigates this degradation by continually updating the running average of the second moment, just as Adam does, but in the current (slowly changing) coordinate basis. Furthermore, since SOAP is equivalent to running Adam in a rotated space, it introduces only one additional hyperparameter (the preconditioning frequency) compared to Adam. We empirically evaluate SOAP on language model pre-training with 360m and 660m sized models. In the large batch regime, SOAP reduces the number of iterations by over 40% and wall clock time by over 35% compared to AdamW, with approximately 20% improvements in both metrics compared to Shampoo. An implementation of SOAP is available at https://github.com/nikhilvyas/SOAP.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google