[論文レビュー] Efficient Neural Network Training via Forward and Backward Propagation Sparsification

Xiao Zhou, Weizhong Zhang+3 more|arXiv (Cornell University)|Nov 10, 2021

Domain Adaptation and Few-Shot Learning参考文献 51被引用数 15

ひとこと要約

この論文は、重みと構造パラメータの更新を分離することで、効率的な順方向および逆方向伝搬を実現する完全にスパースな深層ニューラルネットワークの訓練手法を提案する。2回の順方向伝搬のみを必要とする分散低減型方策勾配推定器を導入し、完全にスパースな逆誤差伝搬を可能にし、従来の手法と比較して最大10倍の高速化を達成する。

ABSTRACT

Sparse training is a natural idea to accelerate the training speed of deep neural networks and save the memory usage, especially since large modern neural networks are significantly over-parameterized. However, most of the existing methods cannot achieve this goal in practice because the chain rule based gradient (w.r.t. structure parameters) estimators adopted by previous methods require dense computation at least in the backward propagation step. This paper solves this problem by proposing an efficient sparse training method with completely sparse forward and backward passes. We first formulate the training process as a continuous minimization problem under global sparsity constraint. We then separate the optimization process into two steps, corresponding to weight update and structure parameter update. For the former step, we use the conventional chain rule, which can be sparse via exploiting the sparse structure. For the latter step, instead of using the chain rule based gradient estimators as in existing methods, we propose a variance reduced policy gradient estimator, which only requires two forward passes without backward propagation, thus achieving completely sparse training. We prove that the variance of our gradient estimator is bounded. Extensive experimental results on real-world datasets demonstrate that compared to previous methods, our algorithm is much more effective in accelerating the training process, up to an order of magnitude faster.

研究の動機と目的

• スパースな順方向伝搬にかかわらず、逆方向伝搬が密行列を必要とする既存のスパース訓練手法の非効率性を解消すること。
• 構造パラメータの勾配推定を再考することで、完全にスパースな順方向および逆方向伝搬を可能にすること。
• スケーラブルでメモリ効率の良い手法を構築し、メモリオーバーヘッドを増加させることなく、より大きな過パラメータ化されたネットワークの訓練を可能にすること。
• PyTorch や TensorFlow などの標準的なディープラーニングフレームワークにおいて、実際の計算高速化を達成すること。

提案手法

• グローバルスパースネスの制約下で、スパースニューラルネットワークの訓練を制約付き最小化問題として定式化し、重みと構造パラメータ（ベルヌーイマスク確率）を同時に最適化する。
• 訓練を2段階に分離する：重みの更新は標準的なチェーンルールに基づく逆誤差伝搬（マスク構造によるスパース化を伴い）、構造パラメータの更新は新規の分散低減型方策勾配推定器（VR-PGE）を用いる。
• VR-PGE を用いて、2回の順方向伝搬のみで構造パラメータの勾配を推定し、逆方向伝搬の必要性を排除することで、完全にスパースな計算を実現する。
• ベルヌーイパラメータの合計を制約することでグローバルスパースネスを制御し、訓練全体を通して一貫したネットワークサイズを保つ。
• 順方向および逆方向伝搬の両方でスパースサブネットワークを用い、FLOPs を密度ベースラインの約 ρ² にまで削減する。ここで ρ はチャンネル残存比である。
• 新しいサブネットワークのサンプリング頻度を低くする戦略（例：50イテレーションごと）を導入し、CPUに全モデルを格納し、GPUにサブネットワークを保持することで、メモリ圧力を軽減する。

実験結果

リサーチクエスチョン

• RQ1 密な逆誤差伝搬に依存せずに、深層ニューラルネットワークにおける完全なスパース訓練（順方向および逆方向伝搬の両方）を達成できるか？
• RQ2 分散低減型方策勾配推定器が、チェーンルールに基づく勾配推定に代わって構造パラメータの勾配推定に用いられても、訓練の安定性および収束性を維持できるか？
• RQ3 完全にスパースな訓練が、標準的なディープラーニングフレームワークにおいて実際の計算高速化をもたらすか？
• RQ4 提案手法は、過大なメモリまたは計算コストを伴わずに、より大きなネットワークにスケーリング可能か？

主な発見

• VGG-19 において、0.8% のパラメータしか残存しない状態（ρ = 0.8）で、訓練の計算時間が最大 7.41× 速くなった。
• ResNet32 を用いた CIFAR-10 では、3.06% のパラメータが保持された状態で 4.50× の高速化が達成され、バリデーション精度は 90.82% を維持した。
• 提案手法の VR-PGE 推定器の分散は、証明可能に有界であり、安定的かつ信頼性の高い勾配更新を保証する。
• すべてのスパースネスレベルにおいて高い精度を維持し、極めて低いチャンネル残存比でもわずかな低下（例：VGG-19 では 93.84% から 90.82%）にとどまる。
• 実験的結果から、本手法は従来のスパース訓練ベースラインと比較して著しく高速であり、最大で10倍の高速化が確認された。
• サブネットワークのサンプリング頻度を 50 イテレーションごとに低下させても、最終的な精度が保持され、効率的なメモリ管理と大規模モデルへのスケーラビリティが可能となった。