[Paper Note] SOAP: Improving and Stabilizing Shampoo using Adam, Nikhil Vyas+, ICLR'25

[AkihikoWatanabe]

URL

• https://arxiv.org/abs/2409.11321

Authors

• Nikhil Vyas
• Depen Morwani
• Rosie Zhao
• Mujin Kwun
• Itai Shapira
• David Brandfonbrener
• Lucas Janson
• Sham Kakade

Abstract

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

Translation (by gpt-4o-mini)

• Shampooという高次の前処理法が、深層学習の最適化タスクにおいてAdamよりも効果的であるという証拠が増えてきている。しかし、Shampooの欠点は、Adamと比較して追加のハイパーパラメータと計算オーバーヘッドがあることである。Adamは第一および第二モーメントの量の移動平均のみを更新する。 本研究では、Shampoo（1/2乗で実装）とAdafactor（Adamのメモリ効率の良い近似）との間に正式な関係を確立し、ShampooはShampooの前処理器の固有基底でAdafactorを実行することと同等であることを示す。この洞察は、よりシンプルで計算効率の良いアルゴリズム、すなわち$\textbf{S}$hampo$\textbf{O}$と$\textbf{A}$damを前処理器の固有基底で組み合わせたSOAPの設計につながる。

Shampooの計算効率を改善するための最も簡単なアプローチは、Shampooの固有分解をより頻繁に計算することを避けることである。しかし、我々の実証結果が示すように、これによりパフォーマンスが低下し、その頻度が増すにつれて悪化する。SOAPは、Adamと同様に第二モーメントの移動平均を継続的に更新することで、この低下を軽減するが、現在の（ゆっくり変化する）座標基底において行う。さらに、SOAPは回転した空間でAdamを実行することと同等であるため、Adamと比較して追加のハイパーパラメータ（前処理の頻度）が1つだけ導入される。 我々は、360mおよび660mサイズのモデルを用いた言語モデルの事前学習においてSOAPを実証的に評価した。大規模バッチの状況下で、SOAPはAdamWと比較してイテレーション数を40%以上、経過時間を35%以上削減し、Shampooと比較して両方の指標で約20%の改善を示した。SOAPの実装は、https://github.com/nikhilvyas/SOAP で入手可能である。

Summary (by gpt-4o-mini)

• Shampooという前処理法が深層学習の最適化タスクで効果的である一方、追加のハイパーパラメータと計算オーバーヘッドが課題である。本研究では、ShampooとAdafactorの関係を明らかにし、Shampooを基にした新しいアルゴリズムSOAPを提案。SOAPは、Adamと同様に第二モーメントの移動平均を更新し、計算効率を改善。実験では、SOAPがAdamWに対して40%以上のイテレーション数削減、35%以上の経過時間短縮を達成し、Shampooに対しても約20%の改善を示した。SOAPの実装は公開されている。

[AkihikoWatanabe]

openreview: https://openreview.net/forum?id=IDxZhXrpNf