[論文レビュー] Compressing Gradient Optimizers via Count-Sketches

Ryan Spring, Anastasios Kyrillidis+2 more|arXiv (Cornell University)|Feb 1, 2019

Advanced Neural Network Applications参考文献 32被引用数 5

ひとこと要約

本稿では、Adam や Momentum、AdaGrad などの一次最適化手法における補助変数を Count-Sketch データ構造を用いて圧縮することで、大規模な深層学習モデルにおけるメモリ使用量を削減する手法を提案する。埋め込み層およびソフトマックス層におけるスパarsity を活用することで、収束速度およびモデル性能を維持したまま、最大 25% のメモリ削減と、4950万クラスを有する極端な分類タスクにおいて最大 38% の高速化を達成する。

ABSTRACT

Many popular first-order optimization methods (e.g., Momentum, AdaGrad, Adam) accelerate the convergence rate of deep learning models. However, these algorithms require auxiliary parameters, which cost additional memory proportional to the number of parameters in the model. The problem is becoming more severe as deep learning models continue to grow larger in order to learn from complex, large-scale datasets. Our proposed solution is to maintain a linear sketch to compress the auxiliary variables. We demonstrate that our technique has the same performance as the full-sized baseline, while using significantly less space for the auxiliary variables. Theoretically, we prove that count-sketch optimization maintains the SGD convergence rate, while gracefully reducing memory usage for large-models. On the large-scale 1-Billion Word dataset, we save 25% of the memory used during training (8.6 GB instead of 11.7 GB) by compressing the Adam optimizer in the Embedding and Softmax layers with negligible accuracy and performance loss. For an Amazon extreme classification task with over 49.5 million classes, we also reduce the training time by 38%, by increasing the mini-batch size 3.5x using our count-sketch optimizer.

研究の動機と目的

• Adam や Momentum のような一次最適化手法における補助パラメータの増加に伴うメモリオーバーヘッドを軽減すること。
• 深層学習における最適化状態のメモリ消費を低減することで、より大きなバッチサイズとより表現力の高いモデルの実現を可能にすること。
• 埋め込み層およびソフトマックス層に内在する構造的スパarsity を活用して、性能劣化を伴わずに最適化状態を効率的に圧縮すること。
• 完全精度の最適化状態に代わる理論的裏付けのある、記憶効率の高い代替手法を提供し、SGD の収束速度を維持すること。
• 大規模な言語モデル学習および極端な分類タスクにおいて、精度損失を最小限に抑えつつ実用的利点を示すこと。

提案手法

• 埋め込み層およびソフトマックス層における二次最適化状態（例：Adam の速度およびモーメンタム）の圧縮表現を維持するために、Count-Sketch データ構造を用いる。
• フルな最適化状態の代わりに、固定サイズのスケッチ（例：[3, 266, 1024]）を格納することで、二階モーメント推定において最大 99% のメモリ削減を実現する。
• 勾配更新をビンにマッピングするためのランダムハッシュ関数を用いて、スケッチを段階的に更新することで、定数時間オペレーションを実現する。
• パラメータ更新時にスケッチをクエリしてフルな最適化状態を近似するが、理論的バウンドにより誤差を制御する。
• ハッシュ関数の数およびビンの数を調整することで、スケッチサイズと精度の間の滑らかなトレードオフを可能にする。
• スケッチベースの最適化状態を訓練パイプラインに直接統合し、標準的な最適化状態を圧縮版に置き換える。

実験結果

リサーチクエスチョン

• RQ1 Count-Sketch を用いて深層学習における補助最適化変数を圧縮しても、収束性や性能が劣化しないか？
• RQ2 大規模モデルにおいて、スケッチベースの最適化手法と標準的な Adam のメモリ使用量およびトレーニング時間はどのように比較されるか？
• RQ3 圧縮された最適化状態を用いる場合、どの程度バッチサイズを増加させられ、トレーニング速度にどのような影響を与えるか？
• RQ4 理論的分析において、スケッチベースの最適化手法は、標準的な一次最適化手法と同等の収束速度を維持するか？
• RQ5 4950万クラスを超える極めて大規模な分類タスクに対しても、本手法は効果的に適用可能か？

主な発見

• 10億語の言語モデル学習タスクにおいて、Count-Sketch 最適化手法は、11.7 GB から 8.6 GB にメモリ使用量を 25% 削減し、精度および性能に変化がなかった。
• 4950万クラスを有する極端な分類タスクにおいて、ミニバッチサイズを 3.5 倍に増加させたことで、Count-Sketch 最適化手法によりトレーニング時間を 38% 減少させた。
• Count-Sketch 最適化手法は、Recall@100 が 0.6889 であったのに対し、Adam ベースラインは 0.6881 であり、性能劣化はほとんど認められなかった。
• 1% のサイズのスケッチを用いて二階モーメント推定を圧縮することで、モデルあたりのメモリコストを 4 GB から 2.6 GB に 35% 削減した。
• 理論的分析により、Count-Sketch 最適化手法が、完全精度のベースラインと同等の SGD 収束速度を維持することが確認された。
• 本手法は、メモリ使用量と近似誤差の間で実用的なトレードオフを可能にし、計算オーバーヘッドは最小限に抑えられた。