SOAP: Improving and Stabilizing Shampoo using Adam

要約

深い学習最適化タスクにおいて、Adamをめぐる高次前処理方法であるShampooの有効性の証拠が増えています。
ただし、Shampooの欠点には、Adamと比較した場合の追加のハイパーパラメーターと計算オーバーヘッドが含まれます。これは、1階間および2番目の数量の平均を実行するだけです。
この作業は、シャンプー（アダムのメモリ効率の高い近似）であるシャンプー（1/2パワーで実装）とAdaFactorとの間に正式な接続を確立します。シャンプーは、シャンプーの前処理者の固有ベシスでAdaFactorを実行することと同等であることを示しています。
この洞察は、$ \ textbf {a} $ damを添えて、よりシンプルで計算効率の高いアルゴリズムの設計につながります：$ \ textbf {s} $ \ textbf {o} $ {a} $ dam {a} $ dam {p} $ reconitioner’s eigenbasis（
石鹸）。
Shampooの計算効率の改善に関しては、最も簡単なアプローチは、Shampooの固有カムポジションを頻繁に計算することです。
残念ながら、私たちの経験的結果が示すように、これはこの頻度で悪化するパフォーマンスの劣化につながります。
SOAPは、Adamと同様に、2番目のモーメントの実行平均を2番目のモーメントの平均を継続的に更新することにより、この劣化を軽減しますが、現在の（ゆっくりと変化する）座標です。
さらに、石鹸は回転したスペースでAdamを実行するのと同等であるため、Adamと比較して1つの追加のハイパーパラメーター（前処理頻度）のみを導入します。
360mおよび660mのサイズのモデルでトレーニング前の言語モデルの石鹸を経験的に評価します。
大規模なバッチレジームでは、SOAPはAdamWと比較して35％を超える反復回数を40％以上削減し、シャンプーと比較して両方のメトリックが約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