[論文レビュー] Memory-Efficient Adaptive Optimization

Rohan Anil, Vineet Gupta+2 more|arXiv (Cornell University)|Jan 30, 2019

Advanced Neural Network Applications被引用数 19

ひとこと要約

本稿では、Adam や Adagrad などの適応的最適化手法のメモリオーバーヘッドを削減するために、2次統計の圧縮・低ランク近似を維持することで、記憶効率の高い適応的最適化手法 SM3 を提案する。この手法により、同等またはより優れた収束性能を維持しながら、大幅に大きなミニバッチサイズやモデルサイズを可能にし、大規模言語モデルおよび画像分類タスクでトレーニング時間を最大2倍速くする。収束保証も維持する。

ABSTRACT

Adaptive gradient-based optimizers such as Adagrad and Adam are crucial for achieving state-of-the-art performance in machine translation and language modeling. However, these methods maintain second-order statistics for each parameter, thus introducing significant memory overheads that restrict the size of the model being used as well as the number of examples in a mini-batch. We describe an effective and flexible adaptive optimization method with greatly reduced memory overhead. Our method retains the benefits of per-parameter adaptivity while allowing significantly larger models and batch sizes. We give convergence guarantees for our method, and demonstrate its effectiveness in training very large translation and language models with up to 2-fold speedups compared to the state-of-the-art.

研究の動機と目的

• 大規模トレーニングにおけるモデルサイズやミニバッチサイズの制限要因となっている、Adam や Adagrad などの適応的最適化手法の高いメモリオーバーヘッドに対処すること。
• パラメータごとの適応性の利点を保持しつつ、記憶消費を著しく削減する手法を開発すること。
• 特に自然言語処理（NLP）およびビジョン分野における非常に大きなモデルのトレーニングを可能にするために、より大きなミニバッチサイズに向けたメモリ解放を実現すること。
• 凸オンライン最適化設定において理論的な収束保証を提供すること。
• 実験的に、記憶使用量を削減しつつ、同等またはより優れた収束性能と、より短いウォールタイムを達成できることを示すこと。

提案手法

• SM3 は、各パラメータごとの2次モーメントを完全に保持する代わりに、2次統計の低ランク近似を維持する。
• ランク1の近似を用いて、記憶効率の高い確率的更新ルールにより、適応的統計を圧縮する。
• パラメータごとの適応性を維持しつつ、ストレージを最小限に抑えるために、対角スケーリング因子を適用する。
• 標準的なディープラーニングフレームワークと互換性があり、最小限のコード変更で実装可能である。
• フィッシャー情報行列の構造的近似を通じて、ネイチャラ・グレディエントに類似した更新を活用する。
• 活性化構造に関する事前知識を必要とせず、勾配パターンに動的に適応する。

実験結果

リサーチクエスチョン

• RQ1 記憶効率の高い適応的最適化手法としての SM3 は、Adam や Adagrad と同等の収束性能を維持しながら、メモリ使用量を削減できるか？
• RQ2 SM3 が得るメモリ節約効果を、分散トレーニングにおけるミニバッチサイズの拡大にどの程度活用できるか？
• RQ3 2次統計の低ランク近似が、大規模モデルにおける収束速度に影響を与えるか？
• RQ4 Adafactor や Shampoo などの既存の記憶効率の高い手法と比較して、SM3 はメモリ使用量、速度、精度の面でどのように差をつけるか？
• RQ5 固定された計算リソースの制約下で、SM3 は標準的な適応的最適化手法よりもウォールタイムでの収束が速いか？

主な発見

• WMT’14 en→fr 翻訳タスクでは、SM3 は Adam より 40% のメモリ使用量を削減し、BLEU スコア 40.50 を達成した。
• BERT-Large の言語モデリングでは、バッチサイズを2倍にした状態で、SM3 は Adam より 35% のウォールタイム短縮で 70% のマスクLM精度に到達した。
• バッチサイズ 2048 で、SM3 はコアあたり 6.02 GiB のメモリ使用量に抑えられ、これはバッチサイズ 1024 の Adam と同等の水準であったが、2^16 までスケーリング可能であった。
• AmoebaNet-D を用いた ImageNet では、SM3 はトップ-1精度 78.71%、トップ-5精度 94.31% を達成し、最先端の性能を再現した。
• 大規模モデルにおいて、SM3 は Adam と比較して 40% のメモリ使用量削減を達成しながら、同程度の収束速度と精度を維持した。
• SM3 の1ステップあたりの処理時間は Adam より 3% 速く、バッチサイズを2倍にした場合、収束までのウォールタイムは最大 50% 減少した。