[論文レビュー] Winning the Lottery Ahead of Time: Efficient Early Network Pruning

John Rachwan, Daniel Zügner+4 more|arXiv (Cornell University)|Jun 21, 2022

Advanced Neural Network Applications被引用数 4

ひとこと要約

この論文は、勾配フローの維持によってモデル性能を保ちながら構造的(pruning)なスパースなネットワークを効率的に早期に削減する手法、EarlyCroPを提案する。98.5%の重みスパarsityを達成し、SOTAの精度を実現。トレーニング時間は7倍短縮、GPUメモリ使用量は4.9倍削減、炭素排出量も70%削減され、高性能なモデルをコンsumerハードウェアでトレーニング可能にする。

ABSTRACT

Pruning, the task of sparsifying deep neural networks, received increasing attention recently. Although state-of-the-art pruning methods extract highly sparse models, they neglect two main challenges: (1) the process of finding these sparse models is often very expensive; (2) unstructured pruning does not provide benefits in terms of GPU memory, training time, or carbon emissions. We propose Early Compression via Gradient Flow Preservation (EarlyCroP), which efficiently extracts state-of-the-art sparse models before or early in training addressing challenge (1), and can be applied in a structured manner addressing challenge (2). This enables us to train sparse networks on commodity GPUs whose dense versions would be too large, thereby saving costs and reducing hardware requirements. We empirically show that EarlyCroP outperforms a rich set of baselines for many tasks (incl. classification, regression) and domains (incl. computer vision, natural language processing, and reinforcment learning). EarlyCroP leads to accuracy comparable to dense training while outperforming pruning baselines.

研究の動機と目的

• 深層ニューラルネットワークにおけるスパースなウィンニングチケットを探索するための高い計算コストを低減すること。
• モデル精度を損なわず、トレーニング時間、GPUメモリ使用量、炭素排出量を削減する構造的スプライシングを可能にすること。
• 勾配フローとニューラルタングエント・カーネル（NTK）の安定性を追跡することで、スプライシングに最適な早期トレーニング段階を同定すること。
• 勾配フローの維持によって学習ダイナミクスを保ち、大規模なスパースモデルをコンsumer GPUでトレーニング可能にする。
• 視覚、自然言語処理、強化学習のタスクにおいて、既存の非構造的および構造的スプライシングベースラインを上回ること。

提案手法

• モデルの学習ダイナミクスへの影響を最小限に抑える勾配フロー（GF）に基づくスプライシング基準を提案する。
• 勾配フローとニューラルタングエント・カーネル（NTK）の関係を活用し、早期スプライシングに適した安定したトレーニング段階を特定する。
• NTKが概ね一定である「ラジカル核（lazy kernel）」領域でスプライシングを適用し、性能低下を最小限に抑える。
• 勾配フローとモデル性能を維持しながら、完全なニューロンやフィルタを削除することで構造的スプライシングを実現。
• 複数回の密度付きトレーニングサイクルを必要とせず、1回のトレーニング実行でスパースサブネットワークを同定・スプライシングする。
• NTKの安定性とGFの維持に基づく原理的なスプライシングスケジューリングを採用し、可能な限り早いスプライシングタイミングを決定する。

実験結果

リサーチクエスチョン

• RQ1 スプライシングがモデル性能に与える影響が最小限である早期トレーニング段階を同定できるか？ これにより、効率的なサブネットワーク探索が可能か？
• RQ2 勾配フローの維持が、モデル精度と効率性を保ちながら構造的スプライシングに原理的基準として機能できるか？
• RQ3 早期構造的スプライシングは、精度、推論速度、炭素排出量の観点で、既存の構造的および非構造的ベースラインを上回るか？
• RQ4 早期スプライシングでトレーニングされたスパースモデルは、同サイズの密度付きモデルを上回る性能を達成できるか？
• RQ5 密度付きトレーニングと比較して、早期スプライシングはトレーニング時間、GPUメモリ使用量、炭素排出量をどの程度削減できるか？

主な発見

• EarlyCroP-Uは、スパース性が非常に高い状況でも密度付きモデルと同等またはそれ以上のテスト精度を達成し、LTH や GraSP を含むすべての非構造的ベースラインを上回る。
• EarlyCroP-Sは、密度付きトレーニングと比較して、トレーニング時間を7倍短縮、GPUメモリ使用量を4.9倍削減、炭素排出量を最大70%削減した。
• VGG16を用いたCIFAR-100では、EarlyCroP-Sが98.5%の重みスパarsityと89.9%のノードスパarsityを達成し、92.5%のテスト精度を記録。密度付きモデルを上回った。
• ResNeXt101_32x48dモデルでは、80エポックで93.2%の精度を達成。スパarsityは98.5%で、同サイズの密度付きモデルと比較して6.2倍速くトレーニングされ、炭素排出量は9.5倍少なかった。
• 自然言語処理タスクでは、EarlyCroP-Sは89%のスパarsityでも性能を維持し、それ以上のスパarsityではすべてのベースラインを上回った。一方、LTRはレイヤー固有の重み減衰のため失敗した。
• 強化学習タスクでは、EarlyCroP-SはLTRおよび非構造的ベースラインを上回った。非構造的モデルは、非効率な勾配計算のため、長期間トレーニングを実行した後でのみ追いついた。