[論文レビュー] Shampoo: Preconditioned Stochastic Tensor Optimization

Vineet Gupta, Tomer Koren+1 more|arXiv (Cornell University)|Feb 26, 2018

Stochastic Gradient Optimization Techniques参考文献 21被引用数 33

ひとこと要約

Shampooは、深層学習におけるテンソル構造をもつパラメータのための新しい前処理付き確率的最適化アルゴリズムであり、各テンソル次元に対して別個の前処理行列を維持することで、効率的な2次最適化を可能にする。実際の応用において、SGD、Adam、AdaGradよりも著しく高速に収束するが、構造的行列演算とTensorFlowにおける効率的な実装のおかげで、単純な最適化手法と同等の1ステップあたりの実行時間となる。

ABSTRACT

Preconditioned gradient methods are among the most general and powerful tools in optimization. However, preconditioning requires storing and manipulating prohibitively large matrices. We describe and analyze a new structure-aware preconditioning algorithm, called Shampoo, for stochastic optimization over tensor spaces. Shampoo maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions. We establish convergence guarantees in the stochastic convex setting, the proof of which builds upon matrix trace inequalities. Our experiments with state-of-the-art deep learning models show that Shampoo is capable of converging considerably faster than commonly used optimizers. Although it involves a more complex update rule, Shampoo's runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.

研究の動機と目的

• 高次元の機械学習問題におけるフル行列前処理のスケーラビリティの制限を克服すること。
• モデルパラメータのテンソル構造（例えば行列や4次元畳み込みフィルタ）を活用する実用的な前処理手法を開発すること。
• フルヘッシアン近似の代わりに、次元ごとの前処理行列を維持することで、深層学習における効率的な2次最適化を可能にすること。
• 行列トレース不等式を用いて、確率的凸最適化設定における理論的収束保証を確立すること。
• Adam や SGD といった標準的最適化手法に対して、収束速度の面で実証的に優位性を示すこと、たとえ更新ルールがより複雑であっても。

提案手法

• Shampooは、各テンソル次元ごとに前処理行列のセットを維持し、蓄積された勾配の2次統計を用いてオンラインで更新する。
• 各前処理行列に対して対称行列の平方根演算を実行し、左および右の前処理行列の幾何平均を用いて勾配をスケーリングする。
• 行列パラメータ $ W \in \mathbb{R}^{m \times n} $ に対して、更新則は $ W_{t+1} = W_t - \eta L_t^{-1/4} G_t R_t^{-1/4} $ である。ここで $ G_t $ は勾配、$ L_t, R_t $ は前処理行列である。
• 前処理行列 $ L_t $ および $ R_t $ は、$ L_t = L_{t-1} + G_t G_t^\top $ および $ R_t = R_{t-1} + G_t^\top G_t $ として更新され、勾配の外積を追跡する。
• 高次元テンソルへは、すべてのテンソルモードに同じ原則を適用することで一般化され、最適化中にテンソル構造が保持される。
• Shampooは、標準的なテンソル演算を用いてTensorFlowで効率的に実装されており、最小限のコード変更で既存の深層学習フレームワークと互換性を持つ。

実験結果

リサーチクエスチョン

• RQ1 高次元で構造的なパラメータ空間（例えば深層学習におけるテンソル）にスケーラブルに適応できる前処理付き最適化手法を設計できるか？
• RQ2 各テンソル次元ごとに別個の行列を維持する前処理戦略は、Adam や SGD といった1次最適化手法よりも高速に収束するか？
• RQ3 2次情報を利用しても、1ステップあたりの計算コストを低く抑えられるか？
• RQ4 確率的凸最適化設定において、この構造的前処理アプローチに対して理論的収束保証を確立できるか？
• RQ5 大規模な深層学習モデルにおいて、さまざまなアーキテクチャとデータセットで実際の性能はどのようになるか？

主な発見

• ResNet や Inception アーキテクチャを用いた CIFAR-10 および CIFAR-100 において、Shampoo は Adam や SGD、AdaGrad よりも著しく高速に収束し、少ないエポック数で低い訓練損失を達成した。
• LM1B 語彙モデル化ベンチマークにおいて、Shampoo は Adam や AdaGrad よりも低いテストパープレキシティを達成し、ステップ数による収束速度も速かった。
• 行列の逆平方根を含むより複雑な更新ルールであったにもかかわらず、Shampoo の1ステップあたりの実行時間は SGD や Adam と同等であり、Tesla K40 GPU でもわずかなパフォーマンスペナルティにとどまった。
• 一部のケースでは、例えばCIFAR-100におけるResNet-55では、Shampooの1秒あたりのステップ数が1.249で、Adamの1.203を上回り、実際にはAdamよりも高速だった。
• 前処理行列の遅延更新とモーメンタムの導入により、計算コストを抑えつつも精度に影響を与えることなく、平均化された実行時間の改善が達成された。
• すべてのテストモデルで、最先端の収束速度を達成し、深層学習最適化における構造に配慮した前処理の有効性を示した。