技術要約: AIRE-Prune

問題提起

深層状態空間モデル（SSM）は、入力を内部状態に圧縮することで、シーケンシャルデータのモデリングにおいて高い有効性を示してきました。しかし、状態次元 $n$ は依然として計算およびメモリコストの主要な要因となっています。現代的なアーキテクチャは、安定性と効率性を確保するために、しばしば対角または対角化可能なパラメータ化（例：S5のようなマルチSISOまたはMIMOスタック）を利用しますが、これらのモデルは頻繁に過剰なパラメータ化に陥ります。既存のポストトレーニングによる最適化メカニズムは、$n$ の削減には限定的です。非構造的な重みプルーニングのような従来のアプローチは、状態次元を直接的に削減するものではありません。また、最近のLAST (Gwak et al., 2025) のような他の構造化された手法は、最悪ケースの周波数領域ゲイン（$H_\infty$）に依存して状態をスコアリングします。これは効果的ではあるものの、最悪ケースの指標は、タスクが稀にしか励起しないピーク増幅を強調するため、過度に保守的になる可能性があり、結果として最適な圧縮率を得られない可能性があります。

手法: AIRE-Prune

著者らは、対角（または対角化可能）なSSM向けに設計された、構造化されたポストトレーニング・レイヤー適応型プルーニング手法である AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)ing) を提案しています。最悪ケースのゲインに焦点を当てる従来の手法とは異なり、AIRE-Prune は 漸近的インパルス応答エネルギー（無限時間のホライゾンにおける各状態が寄与する総出力エネルギーを表す、典型的なケースの基準）にプルーニングを適合させています。

コアメカニズム

1. エネルギーに基づくランキング:
   $(\Lambda, B, C)$ で定義される安定な離散時間対角SSMレイヤー（ここで $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ かつ $|\lambda_i| < 1$）に対して、本手法は各状態（モード）$i$ に対して生成される単位インパルス入力による総二乗出力エネルギーを計算します。
   第 $i$ モードの閉形式スコアは以下のように導出されます：
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   このスコアは以下の3つの要素を統合しています：

   • 可制御性 (Controllability): $\|B_{i,:}\|_2^2$ によって表される（励起のしやすさ）。
   • 可観測性 (Observability): $\|C_{:,i}\|_2^2$ によって表される（測定のしやすさ）。
   • 減衰/メモリ (Damping/Memory): 極 $\lambda_i$ によって表される。単位円に近い極を持つ状態（減衰が遅い状態）は、長期的なエネルギーに大きく寄与します。

2. クロスレイヤー正規化:
   レイヤー間でエネルギーのスケールが大きく異なる可能性があるため（例：エンコーダー/デコーダーのゲインによる）、著者らはグローバルな比較を可能にするための プレフィックス正規化スコアリング 戦略を提案しています：

   • 各レイヤー $\ell$ 内で、モードはローカルエネルギー $E^{(\ell)}_i$ に基づいて降順にソートされます。
   • ソートされたエネルギーに対して、プレフィックス和 $S^{(\ell)}_{(i)}$ が計算されます。
   • グローバルなスコアは、「ハザードレート」比として定義されます：
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     この正規化により、スコアはレイヤー内で単調非増加となり、モデル全体に対して単一のグローバル閾値 $\tau$ を適用することが可能になります。

3. プルーニングの実行:
   目標とするプルーニング比率または閾値が与えられると、アルゴリズムは、正規化されたスコアが $\tau$ を超える各レイヤーの最長のプレフィックス状態を保持します。これにより、レイヤー適応型のプルーニングが可能になり、一部のレイヤーは大幅に削減されたり、あるいはその状態がグローバル閾値を下回る場合に完全に除去されたりする一方で、他のレイヤーはほぼそのまま維持されます。

主な貢献

• 閉形式の漸近的エネルギー・スコア: 本論文は、無限ホライゾンの出力エネルギーに基づく、学習されたすべてのモードに対する数学的に導出された閉形式のスコアを提供します。これにより、高価な数値シミュレーションや再学習を回避できます。
• モーダル・トランケーション（モード切断）の拡張: 本研究は、単一の線形システムに適用される古典的なモーダル・トランケーションを、効率性に不可欠な対角構造を維持したまま、深層の非線形SSMスタックへと拡張しています。
• 典型的なケース vs 最悪のケース: 重要度の指標を最悪ケースの $H_\infty$ ゲインから総インパルス応答エネルギーへとシフトさせることで、本手法は稀なピーク増幅ではなく、システムの集約的な振る舞いをターゲットにしています。
• アーキテクチャに依存しない性質: 本手法は、マルチSISO構成およびMIMOスタック（S5スタイルなど）の両方、さらにはS4DやMambaのような他のSSMファミリーにも適用可能です。

実験結果

著者らは、S5、S4D、およびMambaのバックボーンを用いて、Long Range Arena (LRA) ベンチマークおよび Speech Commands データセットでAIRE-Pruneを評価しました。

• 圧縮効率: 多様なベンチマークにおいて、AIRE-Pruneは再学習なしで、平均精度低下わずか 0.29% で平均 60.8% の状態プルーニングを達成しました。
• ベースラインとの比較:
  • LRAスイートにおいて、AIRE-PruneはRandom、Uniform Magnitude、Global Magnitude、LAMP、Uniform $H_\infty$、Global $H_\infty$、およびLASTを含むベースラインを大幅に上回りました。
  • $\le 1$ パーセントポイント (pp) の精度損失という厳格な制約下において、AIRE-Pruneは最強のベースラインであるLASTと比較して、大幅に高いプルーニング比を実現しました。例えば、Path-Xタスクでは、LASTが30%のプルーニングであったのに対し、AIRE-Pruneは70%の状態で安全にプルーニングできました。
  • Speech Commandsタスクでは、無視できるほどの精度損失（フルモデルの96.43%に対し96.40%）で、45%の状態がプルーニングされました。
• 推論効率: プルーニングは直接的にシステムレベルの利点に繋がりました。NVIDIA H100 GPU上で、本手法はタスクやプルーニング比率に応じて 1.23$\times$ から 2.86$\times$ の推論高速化と、19% から 65% のパラメータ削減をもたらしました。
• 汎用性: 本手法はS4DおよびMambaアーキテクチャにおいても堅牢な性能を示し、S4DのTextタスクで最大90%、MambaのImageタスクで80%のプルーニングを、最小限の精度低下で達成しました。

意義と主張

本論文は、AIRE-Pruneが、学習済みSSMの状態空間における 大幅な冗長性 を明らかにし、それが再学習なしで除去可能であることを主張しています。プルーニングを最悪ケースのゲインではなく漸近的応答エネルギーに合わせることで、本手法はタスク性能へのコストを最小限に抑えつつ、積極的な圧縮を実現します。

著者らは、AIRE-Pruneで観察される「ステップ状」の精度低下プロファイル（高いプルーニング閾値に達するまで精度が安定し、その後急激に低下する現象）は、重要な状態と重要でない状態の間に強い分離があることを示していると強調しています。これは、滑らかな劣化を示すベースラインとは対照的であり、ベースラインのスコアリング手法が重要な状態と重要でない状態を混在させていることを示唆しています。

実用面において、本手法は以下を可能にします：

1. グローバルな予算管理: レイヤーごとのチューニングなしに、単一のターゲット予算を設定できます。
2. ブロックレベルの削減: 状態が閾値を下回る場合にレイヤー全体をプルーニングすることを可能にし、微細なスパース性を、よりハードウェアフレンドリーなブロック構造の削減へと変換します。
3. トレーニングフリーの最適化: 追加のファインチューニングを必要としないポストトレーニングメカニズムであり、展開されたモデルに即座に適用可能です。

結論として、AIRE-Pruneは、深層SSMにおけるモデル次数削減のための、原理に基づいたエネルギーベースのアプローチを提供し、圧縮率と精度保持の両面において現在の最先端のプルーニング戦略を凌駕しています。