Google が「シャンプー」という 2 次最適化アルゴリズムを提案、Transformer のトレーニング時間を 40% 削減

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

機械学習の最適化手順は現在、一次手法によって支配されています。

SGD であれ Adam であれ、このタイプの最適化アルゴリズムは損失関数の 1 次導関数 (勾配) を計算し、特定の規定された方法で勾配降下方向に重みを反復します。

実際、2 次勾配は勾配の微分を計算し、最も適切な降下方向と速度をより速く見つけることができるため、より優れた特性を持っています。

ただし、計算コストとストレージコストを考慮すると、2 次最適化アルゴリズムはほとんど使用されません。

最近、Google Brain は新しい 2 次前処理方法を提案しました。これにより大きな改善がもたらされ、SGD、Adam、AdaGrad などの 1 次アルゴリズムを上回り、ニューラルネットワークのトレーニング時間が短縮されました。

これは、Transformer トレーニングタスクのどの一次方法よりもはるかに高速であり、同等またはそれ以上の精度を達成できます。ジェフ・ディーンもTwitterで「いいね！」せずにはいられなかった。

シャンプーアルゴリズム

この記事は、以前の 2 次手法であるShampoo アルゴリズムを実際的に改良したものです。

なぜ「シャンプーアルゴリズム」と呼ばれるのでしょうか?これは、実はこのタイプのアルゴリズムのユーモラスな名前です。シャンプーの広告のスローガンは通常、「こする、すすぐ、繰り返す」というもので、これは単純な繰り返しの無限サイクルを示しており、最終的にはボトルのシャンプーがなくなることになります。

この機械学習最適化アルゴリズムは、本論文の責任著者である Yoram Singer 氏の論文「Shampoo: Preconditioning Stochastic Tensor Optimization」で初めて使用されました。この論文は 2018 年に ICML に掲載されました。

シャンプーアルゴリズムは、2 つのプリコンディショナーLt と Rt の統計値を追跡する必要があります。

次に、これら 2 つの前処理演算子の 4 乗根を計算し、それらを逆数にします。 2 つの行列にそれぞれ左側と右側の勾配ベクトルを掛け、t+1 ステップの勾配を反復して、次の式を使用して勾配を取得します。

上記のプロセスは単純な繰り返しのように見えるため、著者はこれを「シャンプー」と呼んでいます。

2018年の論文は理論的な説明に重点を置いていましたが、このような単純な「髪を洗う」ステップでさえ、実際の適用には多くの困難が伴います。

このステップで最も計算量が多いのは、Lt-1/4 と Rt-1/4 です。これら 2 つの数値を計算するには、コストのかかる特異値分解が必要です。

実際、4 次逆根は SVD 法だけでなく、Schur-Newton アルゴリズムでも計算できます。さらに、行列の次元が大きくなるにつれて、後者によって節約される時間はますます大きくなります。

シュア・ニュートン法は、GPU や TPU などのニューラルネットワークアクセラレータの計算リソースを消費することなく、通常の CPU で計算できます。

しかし、それでも行列の根の逆数を計算するのにはまだかなりの時間がかかります。この問題に対処しなければ、トレーニング速度を上げることは不可能になります。

そこで著者は非同期コンピューティング手法を使用し、TensorFlow で Lingvo を使用してトレーニングループを改善しました。

CPU は、トレーニングデータの収集と処理のほか、チェックポイントやトレーニングステータスの要約などの補助的なアクティビティも担当します。 GPU や TPU などのアクセラレータがトレーニングループを実行する場合、通常はアイドル状態または十分に活用されていない状態になり、自動的に倍精度コンピューティングが提供されます。

これにより、トレーニングで消費されるリソースを増やすことなく、前処理演算子を計算するのに最適です。