Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

技術要約：因果的属性プルーニング (Causal Attribution Pruning: CAP)

問題提起
大規模言語モデル（LLM）は強力な多段階推論能力を示しますが、推論コストが膨大であり、実世界への導入を妨げています。モデルのプルーニング（枝刈り）は、再学習なしに冗長なパラメータを削除する解決策を提供しますが、既存の手法は複雑な推論タスクにおける性能を低下させることがよくあります。現在のアプローチは、重みの大きさ（例：Wanda）や局所的な勾配統計などの相関的なシグナルに依存しています。これらの手法は「小さな重みは重要ではない」と仮定していますが、低マグニチュードのパラメータであっても、多段階推論の回路において極めて重要な役割を果たすことがあります。さらに、純粋に相関的な基準では、思考の連鎖（CoT）推論に不可欠な層間の相互作用を捉えることができず、頻繁にモジュール全体が機能を失う「レイヤー崩壊（layer collapse）」を引き起こします。

手法：因果的属性プルーニング (CAP)
著者らは、推論性能を維持するように設計された、トレーニングフリーかつワンショットのプルーニング手法である**Causal Attribution Pruning (CAP)*を提案しています。静的な特性に基づいて重みをランク付けする手法とは異なり、CAPは特定のモデルコンポーネントがタスクの性能に与える因果的影響*を測定します。このフレームワークは、以下の3つのステージで動作します。

1. キャリブレーションと因果的スコアリング (ステージ1):

   • 推論中心のベンチマーク（GSM8K（数学）、StrategyQA（マルチホップ常識）、ARC-Challenge（科学））からサンプリングすることで、小さなキャリブレーション・データセット（$D_{cal}$）を構築します。
   • モデル内の各アテンションヘッド $h_\ell$ について、CAPは介入的マスキング実験を行います。フォワードパス中にそのヘッドを一時的にマスクし、無傷のモデルと比較した際のトークンレベルのクロスエントロピー損失の増加量（$\Delta(h_\ell)$）を測定します。
   • 小規模なサンプルサイズによる分散に対する堅牢性を確保するため、データセットを $K$ 個のサブセットに分割し、損失増加の中央値を算出します。
   • このプロセスにより、単に数値的に大きいだけで冗長なヘッドではなく、機能的に極めて重要なヘッドを特定します。

2. 重みレベルの重要度変換 (ステージ2):

   • ヘッドレベルの因果的スコアは、正規化された重要度因子 $I(h_\ell)$ に変換され、$[1.0, 10.0]$ の範囲にマッピングされます。影響力の高いヘッドには高い因子が割り当てられます。
   • これらの因子は、対応する投影行列（$W_q, W_k, W_v, W_o$）内の個々の重みへと伝播されます。特定の重み $W_{ij}$ の重要度スコアは、そのマグニチュードと、関連するヘッド領域の重要度因子の積として計算されます：$S_{ij} = |W_{ij}| \cdot I(h_{region})$。
   • 自然なヘッド分割構造を持たないMLP層については、共存するアテンションヘッドから導出された層平均の重要度を割り当てます。これは近似であると注記されています。

3. グローバル・プルーニング (ステージ3):

   • モデル全体のすべての重みが、重要度加重マグニチュードスコア（$S_{ij}$）によってランク付けされます。
   • 目標とするスパース性（例：10%、20%、50%）を達成するために、二分探索を通じてグローバルな閾値が決定されます。閾値を下回る重みはゼロに設定され、標準的な推論インフラストラクチャと互換性のある非構造化スパースパターンが生成されます。

主な貢献

• プルーニングのための因果的属性の定式化: 著者らは、個々のアテンションヘッドをマスクした際の期待されるタスク特異的な損失の増加量を測定する方法を定式化し、相関的なプロキシを超えて機能的な重要性を定量化しました。
• 予算を考慮した重みレベルのプルーニング: ヘッドレベルの因果的スコアを重みレベルの重要度値に変換する手順を導入しました。これにより、重要な推論回路の構造を維持しながら、グローバルなスパース制約の下での微細なマグニチュード・プルーニングが可能になります。
• 実証的検証: 本手法は、Llama-3-8B-InstructおよびMistral-7B-Instructを用い、GSM8K、StrategyQA、ARC-Challengeのベンチマークにおいて評価されています。

結果

• 中程度のスパース性 (10–20%): CAPは、推論性能を維持する上で一般的にWandaのベースラインを上回ります。
  • Llama-3-8B-Instructの20%スパース性において、CAPはARC-ChallengeにおいてWandaに対して61%の相対的改善（70.8%対43.9%の精度）を達成しました。
  • GSM8Kの10%スパース性において、CAPはWandaに対して20%の相対的改善（71.5%対59.5%）を示しました。
  • 著者らは、CAPが特定のタスク（例：Llama-3のGSM8K 20%）においてWandaに劣る場合があるものの、Wandaの相関的基準が失敗する多段階推論タスクにおいて、一貫して性能を維持していると述べています。
• 高スパース性 (50%): 両手法とも性能が著しく低下します。CAPはLlama-3において壊滅的な失敗（パープレキシティが428.2まで跳ね上がる）を起こしており、これはMLP層に対する粗い処理に起因しています。このレベルでは、アテンションヘッドの近似に基づき低い重要度が割り当てられた必須のMLP重みが、グローバル閾値によって除去されてしまいます。
• パープレキシティ vs. 推論: 中程度のスパース性において、CAPはWandaと比較してパープレキシティがわずかに高いにもかかわらず、しばしばより優れた推論精度を達成します。これは、一般的な言語モデリングの質と、タスク特異的な推論の維持との間にデカップリング（分離）が存在することを示唆しています。
• アーキテクチャの限界: Mixture-of-Experts (MoE) モデルに関する予備的な実験では、動的なルーティングによりヘッドの重要性が非定常となるため、CAPの効果は限定的であることが示唆されています。

意義と主張
本論文は、アテンションヘッドレベルの因果的属性が、同等のスパース性レベルにおいて、相関的なプルーニング手法よりも推論性能を維持するための優れた基準であることを主張しています。推論タスクに対するヘッドの因果的寄与を明示的に測定することで、CAPは「数値的には小さいが機能的に極めて重要な」パラメータを、マグニチュードに基づく手法が破棄してしまう前に特定し、保護することができます。

著者らは、自らの研究の範囲が限定的であることを強調しています：

• これはトレーニングフリーの手法であり、重みの更新（SparseGPTなど）を必要とするアプローチと比較して計算効率が高いです。
• これは**タスクに整合（task-aligned）**しており、一般的な言語モデリングではなく、推論データに特化してキャリブレーションされています。
• MLPニューロンに対する微細な因果的属性の欠如や、動的にルーティングされるアーキテクチャ（MoE）への適用における困難さなど、限界も認めています。

核心的な貢献は、介入的かつタスクに整合したプルーニングが、一般的なパープレキシティや静的な重み統計に最適化された手法よりも、特定の能力（推論）をより良く維持できることを示した点にあります。