テクニカルサマリー：AMPによる効率的なLLM：アテンションヘッドとMLPのプルーニング

問題提起

大規模言語モデル（LLM）は、認知タスクにおいて人間レベルの性能を達成しているが、膨大なパラメータ数に起因する高い計算コストと推論速度の低下という課題を抱えている。これらの要因は、リソースが制限された環境へのデプロイを妨げている。プルーニング（枝刈り）はモデルサイズを削減するための有望な戦略であるが、既存の手法には以下のような重大なトレードオフが存在する：

• 構造化プルーニング（Structured Pruning）： コンポーネント全体（例：レイヤーやヘッド）を削除するが、多くの場合、計算コストの高いファインチューニングを必要とするか、重要なコンポーネントを削除した際に性能低下を招く。
• 半構造化プルーニング（Semi-Structured Pruning）： スピードアップを実現するために、特定のハードウェア（例：2:4スパース性パターンを持つ特定のGPU）を必要とし、消費者向けハードウェアでの展開を制限する。
• 非構造化プルーニング（Unstructured Pruning）： 不規則な疎行列を作成するため、特殊なソフトウェアの最適化なしには、汎用ハードウェア上で推論を加速させることができない。
• 一般的な限界： 既存の技術の多くは、損失ベースの指標、テイラー展開、または複雑なヒューリスティックに依存しており、これらは計算コストが高く、柔軟性に欠け、あるいは専用のハードウェアを必要とする。したがって、性能低下を最小限に抑えつつ、ハードウェアに依存せず、実効的な推論スピードアップを実現できる効率的な手法が求められている。

手法：AMP（Attention Heads and MLP Pruning）

著者らは、マルチヘッド・アテンション（MHA）内の重要度の低いアテンションヘッドと、マルチレイヤー・パーセプトロン（MLP）内のニューロンを削除する構造化プルーニング手法であるAMPを提案している。重みの大きさや損失の変化のみに依存する手法とは異なり、AMPは、各コンポーネントが残差ストリーム（residual stream）への寄与度を評価するために、入力データを重みに投影することで、構造的な重要性を評価する。

コアメカニズム

この手法は、情報が残差ストリームを通じて流れるという原理に基づいている。AMPは、各コンポーネントの重要性を、そのレイヤーの最終出力に対する活性化の寄与の大きさに基づいて算出する。

1. MHAコンポーネント：

   • MHAの出力は、個々のアテンションヘッドからの寄与の総和として分解される（$O = \sum h_n W_n$）。
   • $n$番目のヘッドの重要度（$I_n$）は、ヘッドの出力（$h_n$）とその対応する出力投影行列（$W_n$）の積の$\ell_1$ノルムとして定義される：
     $$I_n = \|h_n W_n\|_1$$
   • このアプローチは、単純な重みノルムによるプルーニングとは異なり、重みと入力の間の非線形な相互作用を考慮している。

2. MLPコンポーネント：

   • 本手法は、LLaMAなどのモデルで使用されているSwiGLUメカニズム、具体的にはUp投影とGate投影を対象とする。
   • SwiGLUは要素ごとの乗算を使用するため、アーキテクチャの整合性を維持するために、ニューロンは両方の投影からペアでプルーニングされる必要がある。
   • $m$番目のニューロンのペアの重要度（$I_m$）は、サンプル入力トークンにおけるGate投影とUp投影の要素ごとの積の$\ell_1$ノルムの平均として計算される：
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. プルーニング・プロセス：

   • 小規模なキャリブレーション・セット（例：Alpacaからランダムに抽出した50サンプル）を使用して、平均活性化値を抽出する。
   • 全レイヤーにわたって、コンポーネント（ヘッドおよびニューロン）を重要度順にランク付けする。
   • 目標とする圧縮率 $c$ を達成するために、すべてのレイヤーにおいて下位 $c\%$ のアテンションヘッドと $c\%$ のMLPニューロンを一様に削除する。
   • 学習後（Post-Training）： プルーニングされたモデルは、性能を回復させるために、低ランク適応（LoRA）を用いた軽量なファインチューニングを受ける。

主な貢献

• 新規の重要度指標： AMPは、入力活性化を重みに投影してコンポーネントの重要性を評価するという、データ駆動型の指標を導入しており、これは従来の重みのみ、あるいは損失ベースのヒューリスティックとは異なる。
• ハードウェアに依存しないスピードアップ： 本手法は、特殊なハードウェアやスパース・テンソル・コアを必要とせず、実用的な推論スピードアップ（最大1.25倍）を実現し、消費者向けGPUへの適用を可能にする。
• 効率性： プルーニングの決定プロセスは消費者向けGPUでわずか数分で完了し、その後のファインチューニングもフル再学習と比較して極めて少ないリソースしか必要としない。
• コヒーレンス検証： 著者らは「コヒーレンス・チェック」実験を提案・実行し、AMPによって特定された「最も重要な」コンポーネントを削除した場合に壊滅的な性能低下が発生することを示し、指標が重要な構造と非重要な構造を区別できる能力を検証した。
• 柔軟性： 本手法は異なるLLMファミリー（LLaMA, Phi）に適用可能であり、固定されたスパース性パターン（例：2:4）に縛られることなく、柔軟な圧縮率を選択できる。

実験結果

著者らは、LLaMA 7B、LLaMA-2 7B、Phi-1.5、およびPhi-2モデルに対して、ゼロショットベンチマーク（WinoGrande, HellaSwag, ARC-e/c, PIQA）およびパープレキシティ（WikiText2）を用いてAMPを評価した。

• 性能：
  • LLaMA 7B： 20%のプルーニング率において、AMPは平均精度**64.52%**を達成し、LLM-Pruner (59.71%) や Shortened LLaMA (63.46%) などの最先端手法を上回った。
  • LLaMA-2 7B： 30%のプルーニング率において、AMPは平均精度**61.02%**を達成し、SliceGPT (51.50%) や LLM Surgeon (59.03%) を凌駕した。
  • 改善度： AMPは、既存の構造化プルーニング手法に対し、平均タスク精度で最大1.49ポイント上回った。
• 推論速度：
  • AMPは、LLaMA 7Bで1.25倍、LLaMA-2 7Bで1.19倍のスピードアップを達成した。
  • これはShortened LLaMA (1.13倍) を上回り、特殊なハードウェアを必要とせずにWanda (1.24倍) のような半構造化手法に匹敵する。
• パープレキシティ： ファインチューニングされたプルーニング済みモデルは、非ファインチューニング版と比較してWikiText2において著しく低いパープレキシティを示し、その差は高いプルーニング率において拡大した。
• アブレーション研究：
  • コヒーレンス・チェック： プルーニングの順序を逆にし（最も重要なコンポーネントを削除）、精度が25ポイント以上低下したことは、指標の妥当性を裏付けている。
  • コンポーネントの分離： MHAヘッドのみ、あるいはMLPニューロンのみをプルーニングした場合、両方を同時にプルーニングする場合よりも大幅に性能が悪化し、バランスの取れたアプローチの必要性が示された。

重要性と主張

本論文は、AMPが以下の特性を備えることで、モデル圧縮における重要なトレードオフに対処していると主張している：

1. 効果的： 高い圧縮率を維持しながら、精度において現在の最先端の構造化プルーニング手法を上回る。
2. 効率的： プルーニングの決定（数分）および性能回復（軽量なLoRAファインチューニング）の両方において、最小限の計算リソースしか必要としない。
3. 実用的： 標準的なハードウェア上で実効的な推論スピードアップを可能にし、リソースが制限された環境へのデプロイを容易にする。
4. 持続可能： 推論レイテンシとパラメータ数の減少により、運用コストを約17.39%、炭素排出量を約19.97%削減できることから、Green AIの原則に合致している。

著者らは、本手法が非必須の構造を正確に特定して削除できることを強調しており、これにより、特殊なハードウェア・インフラストラクチャを必要とせずに、高度に最適化されたLLMのデプロイが可能になるとしている。また、プルーニング後の性能を回復させる上で、事後学習（ファインチューニング）が極めて重要な役割を果たすことも強調している。