GHOST: Mamba2 における Phantom State の解明と構造化プルーニング手法の技術的サマリー

本論文は、大規模な Mamba2 モデルの推論ボトルネックを解決し、推論時のメモリ帯域幅を削減するための新しい構造化プルーニング手法**「GHOST (Grouped Hidden-state Output-aware Selection & Truncation)」**を提案しています。以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

1. 背景と問題定義

Mamba2 は、Mamba1 から状態空間モデル（SSM）の状態次元 $N$ を 16 から 128 に拡張することで、時間的モデリング能力を大幅に向上させました。しかし、この 8 倍の拡張は推論において重大なボトルネックを引き起こしています。

• メモリ帯域幅の飽和: 13 億パラメータモデルの場合、再帰状態（Recurrent State）のサイズが約 12MB から約 100MB に膨張し、推論中のメモリ帯域幅を飽和させ、キャッシュの局所性を低下させます。
• 既存手法の限界:
  • 非構造化プルーニング: 重みのスパース化は行えますが、入力ベクトルが密（dense）であるため、結果として得られる活性化値も密になり、メモリ帯域幅の削減には寄与しません。
  • 静的なマグニチュード基準: 重みのノルム（大きさ）のみで重要度を判断するため、システムダイナミクス（動的エネルギー）を無視します。これにより、「Phantom States（重みは小さいが実際には高頻度で活動する状態）」が誤って削除され、「Corporeal States（重みは大きいが生産性が低い状態）」が維持されるという失敗モードが発生します。
  • 勾配ベース手法: 損失感度を推定する Taylor 法などは高精度ですが、大規模モデルでは 45GB 以上の VRAM が必要となり、実用的ではありません。また、分布シフトの問題も抱えています。

2. 手法：GHOST

GHOST は、制御理論における**「バランスド・トランケーション（Balanced Truncation）」**の原理を、バックプロパゲーション（勾配計算）なしで、フォワードパスの統計情報のみを用いて近似する構造化プルーニングフレームワークです。

2.1 核心的なアプローチ

GHOST は、各状態チャネルの**「制御性（Controllability）」と「観測性（Observability）」**を同時に測定し、その積（ saliency score）に基づいて状態次元 $N$ を削減します。

1. 制御性（Controllability）の推定:
   • 入力履歴が現在の状態にどの程度影響を与えるかを評価します。
   • 隠れ状態 $H_t$ の共分散（分散）を指標として用います。入力から転送されたエネルギーの大きさを表します。
2. 観測性（Observability）の推定:
   • 現在の状態が将来の出力にどの程度影響を与えるかを評価します。
   • 出力エネルギーの Hessian 行列（曲率）を指標として用います。具体的には、即座の投影 $(C'_t)^\top C'_t$ の対角成分を近似値として利用します。
3. 重要度スコア（Saliency Score）:
   • 制御性と観測性の積を計算し、Hankel 特異値の理論的定義に合わせます。
   • $S_i \approx \sqrt{P_i \cdot Q_i}$ （$P$: 制御性、$Q$: 観測性）
   • これにより、入力から出力への情報フローとして「生きている（Living）」状態を特定し、Phantom State を特定して維持し、Corporeal State を削除します。

2.2 実装とプロセス

• グループ化処理: Mamba2 の Grouped Query Attention (GQA) 構造を利用し、グループ内のスコアをプールして層ごとに閾値を設定します。これにより、複雑なダイナミクスをモデル化するグループは多くの状態を維持し、冗長なグループは厳しくプルーニングする適応的な容量割り当てが可能になります。
• 計算コスト: 勾配や Hessian 行列の計算は不要です。校正データ（Calibration Data）に対する2 回のフォワードパスのみで完了し、メモリオーバーヘッドは $O(G \cdot N)$ と非常に低く抑えられます。
• 逐次更新: 層ごとに順次処理を行い、活性化値を更新しながら次の層を校正することで、分布シフトを軽減します。

3. 主要な貢献

1. GHOST フレームワークの提案: Mamba2 の状態次元 $N$ を、制御性と観測性の測定に基づいて構造的にプルーニングする、データ駆動型の手法を提案しました。
2. 「代理失敗（Proxy Failure）」の解決: 静的な重みのマグニチュードに依存する既存手法の盲点を克服し、Phantom State（重要だが重みが小さい）を維持し、Corporeal State（重みは大きいが活動が低い）を削除することに成功しました。
3. 高性能と低コストの両立: 勾配計算を一切行わず、Taylor プルーニングなどの高コストな手法と同等、あるいはそれ以上の性能を達成しました。

4. 実験結果

Mamba2 の 1.3B モデルおよび 130M〜2.7B の様々なスケールで評価を行いました。

• スパース性と性能のトレードオフ:
  • WikiText-2 において、状態次元を50% 削減しても、Perplexity は約 1 ポイントのみ増加（13.17 → 14.23）しました。
  • 70% の高スパース性領域でも、Taylor 法や Magnitude 法が性能崩壊（Perplexity 急増）を起こす中、GHOST は滑らかに性能を維持しました。
• シーケンス長の一般化:
  • 短い文脈（128 トークン）で校正し、長い文脈（2048 トークン）で評価した場合、GHOST は Dense モデルと同様の傾向で性能が維持されました。一方、Taylor 法は文脈が長くなるにつれて性能が急激に劣化しました。
• スケーリング法則:
  • 小規模モデル（130M, 370M）において、勾配ノイズに弱い Taylor 法が破綻する中、GHOST はすべてのスケールで安定した性能を示しました。
• ゼロショット性能と OOD 頑健性:
  • Lambada、PIQA などのゼロショットタスクや、コード生成（HumanEval）、数学（MMLU）などの分布外（OOD）タスクにおいて、GHOST は他の構造化プルーニング手法を凌駕し、Dense モデルに近い性能を維持しました。これは、GHOST が言語モデルの「次の単語予測」だけでなく、SSM としての根本的な動作に重要なチャネルを特定していることを示唆しています。
• 計算効率:
  • 推論時の VRAM 使用量は Taylor 法（45GB 超）に対し、GHOST は 15GB 程度で済み、A100 などの標準 GPU でも実行可能です。

5. 意義と結論

GHOST は、大規模 SSM モデルの推論コストを削減するための画期的な手法です。

• 理論的統合: 深層学習と制御理論（バランスド・トランケーション）を結びつけ、勾配なしでシステムダイナミクスを正確に捉えることを実証しました。
• 実用性: 大規模モデルの展開を民主化し、メモリ帯域幅のボトルネックを解消することで、エッジデバイスやリソース制約のある環境での Mamba2 の利用を可能にします。
• 将来展望: 本手法は Mamba2 特有の構造（スカラー恒等行列 A、グループ化ダイナミクス）に最適化されていますが、制御理論に基づく状態選択の原理は、他の SSM 変種やアーキテクチャへの拡張、およびシステムレベルの最適化（スパースカーネル実装など）への道を開くものです。

要約すれば、GHOST は「重みの大きさ」ではなく「実際の動的なエネルギー」に基づいて状態を選別することで、Mamba2 の推論効率を劇的に向上させつつ、モデルの知能を維持する画期的な技術です。