An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

Mambaモデルを、巨大で非常に賢いオーケストラだと想像してください。すべてのレイヤーの中に、16個の小さな単音楽器（「モード」と呼ばれます）のバンクがあります。従来の考え方では、指揮者（モデル）は、音楽の内容がどうであれ、例えば「最高の8つ」という固定された楽器のセットを曲全体を通して選び続けると想定されてきました。

しかし、この論文はその幕を引き、衝撃的な秘密を明らかにしました。指揮者は決して固定されたセットを選んでいるわけではないのです。

代わりに、指揮者は熟練した即興演奏家です。モデルが単語（「トークン」）を読み取るたびに、メロディを奏でるために実際に必要な8つの楽器がどれであるかを、瞬時に再決定します。時にはフルートであり、時にはバイオリンであり、時にはドラムなのです。「重要な」楽器は、入力に応じて**移動（マイグレーション）**します。もしオーケストラに固定された8つのセットを使い続けるよう強制すれば、それは行進曲の楽譜を使ってジャズのソロを演奏するように強いることになります。それは機能はしますが、本物と比較するとひどい響きになります。

魔法のツール：「正確な楽器」

著者たちはどのようにしてこれを知ったのでしょうか？ 彼らは数学的な「正確な楽器」を構築しました。

オーケストラの内部構造は、楽器同士が干渉しない特殊な「対角（ダイアゴナル）」構成であるため、著者たちは出力をあらゆる楽器の貢献の完璧な総和へと分解することができました。彼らは「グラム・テンソル」（超精密なスコアカードのようなもの）を作成し、特定の楽器グループを外した場合に、どれほどの誤差が生じるかを正確に算出できるようにしました。

彼らはこのツールを実際のモデルと比較し、相対誤差が 2.3 × 10⁻⁷ であることを発見しました。これは、地球から月までの距離を測定して、誤差が髪の毛一本の幅にも満たないレベルであることを意味します。これは推定値ではなく、精密な測定値なのです。

大きな発見：「マイグレーション・ギャップ」

このツールを用いて、彼らは極小のモデル（130Mパラメータ）から、大規模なモデル（デプロイされているFalcon-Mambaのような7Bパラメータ）までを調査しました。

最も活動的なレイヤーにおいて、固定されたセットは、入力に応じて変化するセットよりも2倍多くの誤差を生むことが分かりました。

• 統計： 最も影響を受けるレイヤーにおいて、「マイグレーション・ギャップ」（固定セットと変化するセットの間の誤差比）は 0.44 から 0.57 の間でした。
• 意味： もしモデルが各特定の瞬間に対して最適な楽器を選ぶ（「入力スケジューリングされたオラクル」）ことができれば、一度固定リストを選んでそのまま使い続けるよりも、誤差を半分に減らすことができます。

これは、彼らがテストしたすべてのモデル、すなわちMamba-1ファミリー、7B Falcon-Mamba、そしてMamba-2において見られました。

何がマイグレーションを引き起こすのか？（「なぜ」か）

著者たちは問いかけました。「どの部分が、この切り替えを行っているのか？」
Mambaレイヤーには、主に3つの信号があります。

1. 書き込みマップ ($B_t$): どの楽器に入力信号を与えるかを決定します。
2. 読み出し ($C_t$): どの楽器が音として聞こえるかを決定します。
3. タイムステップ ($\Delta$): しばしば「選択性」のノブと考えられているものです。

彼らは「凍結信号（frozen-signal）」実験を行いました。各信号を平均値に固定し、一つずつ検証して、マイグレーションが停止するかどうかを確認しました。

• 結果： 書き込みマップ ($B_t$) を固定すると、マイグレーションは消失しました。モデルは楽器の切り替えをやめました。
• 驚きの事実： タイムステップ ($\Delta$) を固定しても、マイグレーションは全く変わりませんでした。

結論： 多くの人々が選択性の鍵だと考えていた「タイムステップ」信号は、マイグレーションの信号をほとんど運んでいません。真のヒーローは書き込みマップ ($B_t$) です。それは、トークンごとにどの楽器が演奏するかを決定する門番なのです。

その結果：これは利用できるのか？

著者たちは、この知識をモデルの削減（プルーニング）に利用できるか試みました。

• 静的プルーニング (Static Pruning): 平均的な活動に基づき、最高の8つの楽器を選び、それを永久に保持する。
• 入力スケジューリング・プルーニング (Input-Scheduled Pruning): 現在の文章を読み、今まさにどの8つの楽器がアクティブであるかを測定し、それらのみを保持する。

結果：
状態予算の半分（16モードのうち8モードを保持）において、入力スケジューリング方式は、生の精度において、プルーニングされていないフルモデルと同等、あるいは場合によってはそれ以上に優れた性能を発揮しました。

• 130Mモデルにおいて、スケジューリング方式のパープレキシティは 11.84 であり、未プルーニングのモデルの 12.38 を上回りました。
• 7B Falcon-Mambaにおいて、スケジューリング方式は 4.44 を記録し、未プルーニングの 4.48 を上回りました。

しかし、重要な注意点があります： 論文は、この「スケジューリング」方式が**二パス・オラクル（two-pass oracle）**であることを明記しています。つまり、まずウィンドウ全体を一度読み取ってどの楽器を保持するかを決定し、次に二度目のパスを実行して出力を生成します。これは、実世界のデプロイにおいて、計算量やメモリを節約できるわけではない（結局、データを2回読み取る必要がある）ことを意味します。

著者らは、この結果が「デプロイ可能な効率化」ではなく、「実現可能な余地（realizable headroom）」を示していることを明確にしています。これは、もし安価で高速な予測器によって、高価な最初のパスを経ることなく正しい楽器を推測できれば、極小の超効率的なモデルが巨大なモデルに匹敵するポテンシャルがあることを証明しているのです。現在の手法は、単に「到達可能な天井」を示しているに過ぎません。

否定されたもの

論文は、何がうまくいかないかについても明確に述べています。

• 静的ランキング (Static Rankings): 平均的な活動や「ハンケル・エネルギー」（GHOSTやLASTのような）に基づいて固定のモードセットを選ぶ手法は、大幅に劣ります。それらは、動く標的についていくことができません。
• タイムステップの活動量: タイムステップの活動に基づいてプルーニングを行うことは、タイムステップがマイグレーションを駆動していないため、無意味な試みです。
• 単純な予測器: 文の前半部分や「ドメイン」（コードか散文かなど）だけでマスクを予測しようとする安易なトリックを試しましたが、これらは潜在的な利得のわずか 2〜6% しか回収できませんでした。「重要な」楽器のセットは非常に速く変化するため（数百トークンの範囲内で）、スコア対象となる特定のトークンを正確に測定しなければ、完全な恩恵を得ることはできません。

まとめ

この論文は、訓練された選択的状態空間モデル（Selective State-Space Models）が、入力に基づいて内部リソースを絶えず再配分する、動的で生きている存在であることを証明しています。 「書き込みマップ」こそが、このマイグレーションの指揮者です。リアルタイムで（二パスのコストなしに）楽器を切り替えるモデルをまだデプロイすることはできませんが、この研究は、その領域の正確な地図を与えてくれました。現在の「静的」なプルーニング手法が、膨大なパフォーマンスを台無しにしていることを示しており、今後の道筋は、これらのマイグレーションをオンザフライで予測できるスケジューラーを構築することにあるのです。