[論文レビュー] Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

Vamshi Sunku Mohan, Kaustubh Gupta+2 more|arXiv (Cornell University)|Feb 26, 2026

Machine Learning in Healthcare被引用数 0

ひとこと要約

この論文は Mamba 状態空間モデルにおける Activation Subspace Bottlenecks を特定し、事後 steering と Stable-Mamba アーキテクチャを導入する。複数の SSM とベンチマークでタスク固有の調整なしに一貫した性能向上を実現。

ABSTRACT

State-space models (SSMs) have emerged as an efficient strategy for building powerful language models, avoiding the quadratic complexity of computing attention in transformers. Despite their promise, the interpretability and steerability of modern SSMs remain relatively underexplored. We take a major step in this direction by identifying activation subspace bottlenecks in the Mamba family of SSM models using tools from mechanistic interpretability. We then introduce a test-time steering intervention that simply multiplies the activations of the identified bottlenecks by a scalar. Across 5 SSMs and 6 diverse benchmarks, this intervention improves performance by an average of 8.27%, without requiring any task-specific tuning. Finally, we validate that the identified bottlenecks are indeed hindering performance by modifying them to yield an architecture we call Stable-Mamba, which achieves long-context performance gains when retrained from scratch.

研究の動機と目的

• transformer の代替として SSM の解釈性と steerability の動機付け。
• mechanistic interpretability ツールを用いて Mamba の Activation Subspace Bottleneck を特定。
• タスク固有の調整なしに性能を向上させる事後 steering の実証。
• ボトルネックを緩和し長文脈性能を改善する Stable-Mamba アーキテクチャの提案。
• 複数の SSM および多様なベンチマークでアプローチを検証。

提案手法

• SSM の隠れ状態から attention のような影響行列を構築し、活性化サブスペースを定義。
• Sparse Autoencoders (SAE) を用いて活性化の低次元要約を取得し、辞書学習で解釈可能な成分を抽出。
• Stochastic Parameter Decomposition (SPD) を適用して層レベルおよびパラメータレベルの指標（エントロピー、分散、ランク等）を取得。
• Delta-Sensitive サブスペースを、状態更新に結びつく高い活性化分散を持つものとして特定。
• Delta-Sensitive サブスペースをスカラー因子で増幅する事後 steering を実施（試験値は 2 および 5 を含む）。
• Stable-Mamba を、アーキテクチャ的変更（マルチタイムスケール動作、アンサンブル出力、疎なグローバルコンテキスト、学習されたゲーティング、勾配スケーリングなど）で構築し、ゼロから再訓練。

実験結果

リサーチクエスチョン

• RQ1 Mamba SSM で情報ルーティングをボトルネックにする活性化サブスペースは何か。
• RQ2 特定したボトルネックを事後 steering で、タスク固有の調整なしに性能を改善できるか。
• RQ3 解釈性に基づくアーキテクチャ変更（Stable-Mamba）はボトルネックを緩和し、タスク間で長文脈性能を改善するか。
• RQ4 特定されたボトルネックは複数の SSM バリアントやベンチマーク間で移植可能か。
• RQ5 ボトルネックの因果的役割を示す証拠は何か。

主な発見

• 事後 steering は 6 つの多様なベンチマークで 7 SSM の性能を平均 8.27% 向上させた。
• steering の対象はボトルネック層の delta-sensitive サブスペース（例：Vanilla Mamba の Layer 20）で、 steering 因子は 5（一部サブスペースでは 2）を用いた。
• Activation Subspace Bottlenecks は Layer 20 での情報圧縮と相関があり、緩和されると情報の流れが滑らかになりエントロピーが低下。
• Stable-Mamba はゼロから再訓練され、長文脈の利得と情報ボトルネックの低減を示す（例：KL 発散の低下と勾配活動の増加）。
• Steered Mamba および Stable-Mamba は RULER、Long Range Arena、LongBench v2 でベースラインより長文脈性能が優れる。
• steering 効果は 7 SSM にわたり一般化し、IFEval や QA などのタスクで顕著な改善を示す。