[論文レビュー] Scalable Second Order Optimization for Deep Learning

Rohan Anil, Vineet Gupta+3 more|arXiv (Cornell University)|Feb 20, 2020

Stochastic Gradient Optimization Techniques参考文献 49被引用数 29

ひとこと要約

本稿では、因子化された事前処理行列を用いて深層学習の学習を高速化する、スケーラブルでハードウェア最適化された2次順応手法（フルマトリックスAdagradの変種）の実装を提示する。CPUとアクセラレータのパイプライン処理、効率的な行列の平方根計算、アーキテクチャ拡張を活用することで、Transformer や BERT などの大規模モデルにおいて、最大47％のウォールクロック時間短縮を達成し、収束速度と学習効率の両面で、Adam や LAMB などの1次順応手法を大きく上回った。

ABSTRACT

Optimization in machine learning, both theoretical and applied, is presently dominated by first-order gradient methods such as stochastic gradient descent. Second-order optimization methods, that involve second derivatives and/or second order statistics of the data, are far less prevalent despite strong theoretical properties, due to their prohibitive computation, memory and communication costs. In an attempt to bridge this gap between theoretical and practical optimization, we present a scalable implementation of a second-order preconditioned method (concretely, a variant of full-matrix Adagrad), that along with several critical algorithmic and numerical improvements, provides significant convergence and wall-clock time improvements compared to conventional first-order methods on state-of-the-art deep models. Our novel design effectively utilizes the prevalent heterogeneous hardware architecture for training deep models, consisting of a multicore CPU coupled with multiple accelerator units. We demonstrate superior performance compared to state-of-the-art on very large learning tasks such as machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.

研究の動機と目的

• 2次順応手法が大規模な深層学習において実用的でない理由（高いメモリ、計算、通信コスト）を解決する。
• 理論的な2次順応最適化の利点と、スケールでの実装のギャップを埋める。
• 最新の異種ハードウェア（CPU + アクセラレータ）上で、効率的かつスケーラブルな2次順応学習を可能にする。
• 特に埋め込み層のような大規模な層に対して、フルマトリックス事前処理の実装における数値的およびインfraストラクチャ的課題を克服する。
• 多様な大規模タスクにおいて、Adam や LAMB などの1次順応ベースラインと比較して、顕著な収束性およびウォールクロック時間の改善を示す。

提案手法

• CPUに事前処理行列の計算をオフロードするパイプライン化された最適化パイプラインを設計し、アクセラレータはフォワード/バックワードパスを担当する。
• 高コストなスペクトル分解（例：SVD）の代わりに、行列の平方根（例：$L^{-1/4}$, $R^{-1/4}$）を効率的かつ数値的に安定した反復法で計算する。
• Shampooアルゴリズムを拡張し、埋め込み行列のような大規模かつ高ランクの層をサポートするように、事前処理行列の構造と更新ルールを変更する。
• AdaGradから導出される各層ごとの学習率を用いて訓練を安定化させ、AdaGradよりも大きな学習率を可能にする。
• TPUコア間の重複計算を削減するため、事前処理統計のクロスレプリカシャーディングを実装する。
• 非同期的なCPUオフロードを活用して、最小限の遅延で訓練ループに事前処理勾配計算を統合する。

実験結果

リサーチクエスチョン

• RQ1 現代の深層学習ワークロードにおける大規模スケールでの2次順応手法が実用的になるか？
• RQ2 フルマトリックス事前処理の計算およびメモリコストを、収束の利点を損なわず低減できるか？
• RQ3 CPU-アクセラレータシステム上で2次順応手法を効率化するために、どのようなアーキテクチャ的およびアルゴリズム的最適化が必要か？
• RQ4 Adam や LAMB などの1次順応ベースラインと比較して、2次順応手法が学習時間およびステップ数をどの程度短縮できるか？
• RQ5 数値的安定性とスケーラビリティのトレードオフが、大規模2次順応最適化器の設計にどのように影響するか？

主な発見

• WMT’14 英語→フランス語翻訳タスクにおいて、Shampooは標準Transformerで学習時間を45％短縮（12.0時間 → 6.7時間）、大規模なTransformer-Bigモデルでは37％短縮（47.0時間 → 29.5時間）した。
• 32KバッチサイズのBERT-Large言語モデル学習では、ハイパーパramータチューニングなしで、マスクed-LM精度を16％少ないステップ数で達成し、ウォールクロック時間を4％短縮（3.8時間 → 3.65時間）した。
• ResNet-50を用いたCIFAR-10では、Shampooは300ステップで93.45％の精度に到達した一方で、ベースラインは143ステップで到達し、合計の学習時間を42％短縮（1428秒 → 827秒）した。
• Criteo 1TBクリックスルーレート予測タスクでは、学習時間を13分から8.2分に短縮し、大規模スパースモデルにおいても高い効率性を示した。
• 事前処理行列の計算がまだ完全に最適化されていないモデルでも、ステップ数だけでなく実際のウォールクロック時間において顕著な高速化を達成した。
• 現在の実装では、BERT-Largeでステップ時間に14％の増加が見られたが、これは今後の最適化（例：クロスレプリカシャーディング）によりさらなる向上が期待できる。