[論文レビュー] Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients

Milad Alizadeh, Shyam A. Tailor+5 more|arXiv (Cornell University)|Feb 16, 2022

Advanced Neural Network Applications被引用数 11

ひとこと要約

Prospect Pruning (ProsPr) は、初期化時に重みを pruning する新しい手法であり、最初の数ステップの最適化においてメタ勾配を用いて、pruning 後の訓練可能性を明示的に考慮することで、訓練可能で安定したサブネットワークを特定する。単一のステップで pruning を行い、データ要件を低減しながら、視覚分類タスクで最先端の精度を達成しており、損失関数の最適化ダイナミクスを重要度基準に組み込むことで、従来手法を上回っている。

ABSTRACT

Pruning neural networks at initialization would enable us to find sparse models that retain the accuracy of the original network while consuming fewer computational resources for training and inference. However, current methods are insufficient to enable this optimization and lead to a large degradation in model performance. In this paper, we identify a fundamental limitation in the formulation of current methods, namely that their saliency criteria look at a single step at the start of training without taking into account the trainability of the network. While pruning iteratively and gradually has been shown to improve pruning performance, explicit consideration of the training stage that will immediately follow pruning has so far been absent from the computation of the saliency criterion. To overcome the short-sightedness of existing methods, we propose Prospect Pruning (ProsPr), which uses meta-gradients through the first few steps of optimization to determine which weights to prune. ProsPr combines an estimate of the higher-order effects of pruning on the loss and the optimization trajectory to identify the trainable sub-network. Our method achieves state-of-the-art pruning performance on a variety of vision classification tasks, with less data and in a single shot compared to existing pruning-at-initialization methods.

研究の動機と目的

• 既存の初期化時 pruning 手法が、pruning 後の訓練可能性を十分に考慮していないという問題を解決すること。
• 初期化直後の単一ステップでの評価にとどまる現在の重要度基準の短絡的特性を克服すること。
• pruning が最初の数ステップの最適化軌道および損失関数に与える影響を明示的にモデル化する手法を開発すること。
• 反復的再訓練や段階的 pruning を繰り返さずに、初期化段階で高精度な single-shot pruning を可能にすること。
• サブネットワークのスパarsity と初期段階からの訓練可能性の両方を満たすものを効率的に特定することで、ニューラルネットワーク pruning の実用性と効率性を向上させること。

提案手法

• ProsPr は、初期重みに対する損失の勾配（すなわちメタ勾配）に基づいて重要度スコアを計算する。このメタ勾配は、初期化後の数ステップの最適化を経た後、初期重みに関する損失の勾配として得られる。
• 確率的勾配降下法の最初の数ステップを逆伝播することで、特定の重みを pruning した場合に初期学習段階での損失に与える影響を推定する。
• 計算効率と勾配安定性の向上を図るため、メタ勾配の一次近似を用いる。
• すべての重みについて重要度スコアを計算し、そのスコアが低い重みを pruning する。
• pruning マスクを初期重みに直接適用し、残りのサブネットワークを再訓練せずにそのまま学習を開始する。
• 本手法は single-shot を想定しており、従来手法で一般的な繰り返しの pruning と再訓練のループを回避する。

実験結果

リサーチクエスチョン

• RQ1 pruning 後の最適化軌道を考慮する初期化時 pruning 手法は、従来手法を上回る性能を達成できるか？
• RQ2 初期学習段階のメタ勾配を組み込むことで、静的である重要度指標に比べて、訓練可能なサブネットワークの特定がどの程度向上するか？
• RQ3 ProsPr は、単一の pruning ステップと少ないデータ量で、どの程度最先端の精度を達成できるか？
• RQ4 メタ勾配による訓練可能性のモデリングによって、反復的または段階的 pruning 策略の必要性がどの程度低減できるか？
• RQ5 メタ勾配は、損失関数および最適化ダイナミクスに対する pruning の高次効果を効果的に捉えることができるか？

主な発見

• ProsPr は、複数の視覚分類ベンチマークで、既存の初期化時 pruning 手法を上回る最先端の精度を達成した。
• 本手法は、反復的再訓練や段階的 pruning を必要とせず、単一のステップでの pruning で高い精度を達成した。
• 特に高いスパarsity 水準でも、従来手法と比較して性能の低下が小さく抑えられた。
• メタ勾配の使用により、pruning が初期学習ダイナミクスに与える影響をモデル化することで、訓練可能なサブネットワークの特定がより正確になった。
• 従来手法よりも少ないデータ量で効果的に動作し、サブネットワーク発見におけるデータの効率性が向上した。
• 3 ステップを超えてメタ勾配を計算すると、勾配安定性の問題が生じ、本手法は初期段階の少数ステップに限定される。