STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

技術要約: STARFISH

問題提起
ニューラルネットワークのプルーニング（枝刈り）は、重み、ニューロン、またはチャネルを削除することで、モデルのサイズ、メモリフットプリント、および推論レイテンシを削減するための標準的な手法である。しかし、過度なプルーニングはしばしば予測精度の著しい低下を招く。失われた性能を回復させるためにリカバリプロセス（ヒーリング）が通常用いられるが、既存の手法には重大な限界がある。従来のアプローチは、元の訓練データセットを用いたフル再学習やファインチューニングを必要とすることが多く、これは計算コストが高く、元の訓練データがプロプライエタリであったりアクセス不可能であったりする現代の設定においては、頻繁に不可能となる。さらに、最先端のリカバリ手法は、高いスパース性条件下（例：重みの75%以上を削除した場合）において、特に現代のVision Transformer (ViT) アーキテクチャに対して精度を回復させることに苦慮することが多い。

手法: STARFISH
著者らは、STARFISH (faST Accuracy Recovery in pruned networks From Internal State Healing) を提案する。これは、元の訓練データやラベルへのアクセスを必要としない、ポスト・プルーニング（枝刈り後）のリカバリ手法である。その核心となる仮説は、中間表現（内部状態）の整合性を最適化することは、最終的な出力予測のみを最適化することよりも、精度の回復のためのより効果的なプロキシ（代理指標）になるという点である。

アルゴリズムは、訓練データとは異なるがデータ分布からサンプリングされた、ラベルのない小さなキャリブレーションセット ($S_{cal}$) を用いて、以下の2段階で進行する：

1. ターゲット表現のキャッシュ: 元の密な（dense）モデルを $S_{cal}$ に対して実行し、各計算ブロック（例：正規化、非線形関数、および残差加算の後）のベクトル化された出力をキャッシュする。
2. アライメント最適化: プルンド・モデル ($f_{\theta \odot m}$) を、キャッシュされたターゲットとのアライメント損失を最小化するように訓練する。損失関数は、すべてのブロック $\ell$ における、プランド・モデルの隠れ状態 ($f^\ell_{\theta \odot m}(x)$) と密なモデルの隠れ状態 ($r^\ell_x$) の間の平均コサイン距離として定義される：
   $$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \odot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$$
   プランド・モデルの残存する重みは、キャリブレーションセットに対する経験損失を最小化するように更新される。

理論的根拠として、最終的な隠れ表現（分類ヘッドへの入力）の間の距離が、出力分布のカルバック・ライブラー（KL）ダイバージェンスの上界として機能することが示されている。これは、内部状態を整合させることが、本質的に出力誤差を制約することを示唆している。

主な貢献

1. 表現ベースのリカバリ: 本論文は、特定のプルーニングの種類、マスク、またはアーキテクチャに依存しないリカバリ手法を導入している。中間表現のアライメントが、直接的な出力ベースの目的関数（Cross-EntropyやKLダイバージェンス・マッチングなど）よりも、スパース性が高まるにつれて優れた性能を示すことを経験的に実証している。
2. 高スパース性における効率性: 本手法は、極端なスパース性レベル（例：85%のスパース性）において、元の精度を効率的に回復させる初めての手法である。これは、非常に小さなキャリブレーションセット（わずか1,000例、あるいは訓練セットサイズの0.4%）を使用し、元の訓練データへのアクセスを一切必要とせずに達成される。
3. モジュール性と速度: このアプローチは計算量的に軽量であり、小さなデータセットに対する数回のフォワードパスとバックワードパスのみを必要とする。また、構造化、半構造化、および非構造化のプルーニング・パターンすべてに適用可能である。

実験結果
著者らは、様々なVision Transformer (DeiT, ViT) および畳み込みニューラルネットワーク (MobileNet) アーキテクチャを用いて、ImageNet-1KにおけるSTARFISHを評価した。

• 高スパース性 (0.75–0.85): 75%のスパース性にプルーニングされたDeiT-Bモデルにおいて、STARFISHは元の密なモデルの精度の82%を回復したが、競合手法（CORPなど）は40%しか回復できなかった。85%のスパース性を持つDeiT3-Hモデルにおいて、STARFISHは密なモデルの精度の92.08%を維持し、プランド・モデルの性能を72.36パーセントポイント向上させた。
• 中程度のスパース性 (0.5): 様々なDeiTアーキテクチャとプルーニング範囲（MLP、QK、および結合MLP+QK）において、STARFISHは一貫して最先端の手法を上回った。例えば、MLP+QKの50%構造化プルーニングを施したDeiT-Tモデルにおいて、STARFISHは密な精度の81.95%を保持し、CORP法に対して31.54パーセントポイントの改善を示した。
• 比較: STARFISHは、特にプルーニングによってモデルが著しく劣化する領域において、SNOWS（ヘッセ行列近似を使用）やCORP（閉形式のアフィンリカバリを使用）といった手法に対して優れた性能を示した。

意義と主張
本論文は、STARFISHがリカバリのパラダイムを出力のマッチングから内部状態のヒーリングへとシフトさせることで、モデル圧縮の分野における重要な進歩であることを主張している。著者らは、このアプローチが、訓練データが入手できず、フル再学習が不可能な現代の大規模モデルにおいて特に重要であることを強調している。内部表現が、高スパース性下において最終出力よりも堅牢にモデルの知識を捉えるという原理を利用することで、STARFISHは予測性能を損なうことなく、極めて効率的なプランド・ネットワークのデプロイを可能にする。この手法は、初期のプルーニング戦略に関わらず、あらゆるプランド・ネットワークに適用できる、汎用的でプラグアンドプレイなソリューションとして提示されている。