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 などのアクセラレータがトレーニング ループを実行する場合、通常はアイドル状態または十分に活用されていない状態になり、自動的に倍精度コンピューティングが提供されます。

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

非同期計算の使用

各ステップですべてのテンソルの前処理演算子を計算しますが、前処理された勾配は N ステップごとに 1 回計算され、処理のために CPU に渡されます。

この期間中、GPU または TPU は引き続き計算を行っており、更新された事前トレーニング済み演算子が取得されるまで、トレーニング プロセス中に過去の事前条件付け済み演算子が使用されます。

計算はパイプライン化され、トレーニング ループをブロックせずに非同期で実行されます。その結果、Shampoo アルゴリズムの最も計算が難しいステップでは、全体的なトレーニング時間にほとんど追加コストがかかりません。

これだけじゃ十分ではありません。著者はシャンプー アルゴリズムにいくつかの改良を加え、大規模モデルのトレーニングに適したものにしました。これには、ステップ サイズと方向の分離、大きなテンソルの前処理、大きなテンソルの複数のチャンクへの分割が含まれます。

最大速度67%増加

WMT'14 Transformer の英語からフランス語への翻訳トレーニング タスクでは、アルゴリズムによって 1.67 倍の高速化が達成され、時間が 40% 短縮されました。

Shampoo アルゴリズムでは、同じ精度を達成するために Adam や AdaGrad の約半分のステップ数しか必要とせず、学習率に対する耐性も AdaGrad より優れています。

これまでの非同期コンピューティングでは、N はトレーニング計算の量を決定する調整可能なパラメータでした。N が大きいほど、計算量は少なくなります。もちろん、N も結果に影響します。トレーニング プロセスのパフォーマンスと結果の品質の間でトレードオフを行う必要があります。

実験により、このアプローチでは、品質に顕著な低下が生じることなく、最大 1,200 ステップの遅延を許容できることが示されました。

Shampoo は画像分類タスクにも使用できます。

著者は、ImageNet-2012 データセットで ResNet-50 モデルもトレーニングしました。結果は、SGD よりもモメンタムを使用して速く収束しましたが、トレーニング損失は SGD と同等でしたが、テスト セットへの影響は後者ほど良くありませんでした。

一般化能力の欠点に関しては、シャンプーアルゴリズムにはさらなる改善が必要です。

論文の宛先:

出典: arxiv.org

https://arxiv.org/abs/1802.09568