技術要約: MuonSSM – シーケンスモデリングのための直交化状態空間モデル

問題提起

状態空間モデル（SSM）は、長シーケンスモデリングにおいて効率的な線形時間のアテンションメカニズムに代わる手法として台頭している。しかし、既存のSSM（Mamba、DeltaNetなど）は、長いホライゾンにおいて不安定性やメモリの劣化に悩まされることが多い。本論文は、現在のSSMがメモリ状態に対して**一次のランク1更新（first-order, rank-one updates）**に依存していることを指摘している。これらの更新は、主に以下の3つの問題を引き起こす：

1. スペクトル異方性（Spectral Anisotropy）: メモリ行列の特異値が極めて不均一になり、条件数の悪い最適化ランドスケープを引き起こす。
2. 勾配の劣化（Gradient Degradation）: 入力依存の遷移行列を通じた繰り返しの収縮により、長シーケンスにおける勾配消失が生じる。
3. メモリ干渉（Memory Interference）: 時間的な慣性（temporal inertia）とバランスの取れた更新幾何学が欠如しているため、新しい更新が以前の情報を上書きしてしまう。

近年の派生モデルはゲーティングや正規化を通じてこれらを緩和しようと試みているが、蓄積される更新の幾何学を明示的に制御したり、長距離の信号伝搬を改善するための時間的慣性を導入したりするものではない。

手法: MuonSSM

著者らは、効率的な並列スキャンのために必要なアフィン構造を維持しつつ、再帰的な遷移行列ではなく、メモリ更新の幾何学を明示的に制御することでSSMの学習を安定化させるフレームワークであるMuonSSMを提案する。このアプローチは、以下の2つの主要なコンポーネントを用いて標準的なSSMのダイナミクスを拡張する。

1. モメンタム拡張ダイナミクス（Momentum-Augmented Dynamics）

MuonSSMは、タイムステップを横断して更新方向を蓄積する補助的なモメンタム行列 $M_t$ を導入する。これはモメンタムベースの最適化を模倣しており、状態更新に「時間的慣性」を与える。

• モメンタム状態は次のように更新される: $M_t = \gamma M_{t-1} + \text{NS}(\tau \beta_t v_t k_t^\top)$。
• メモリ状態 $S_t$ は、以前の状態とモメンタム項を組み合わせることで更新される: $S_t = S_{t-1}(\alpha_t(I - \beta_t \eta k_t k_t^\top)) + M_t$。
• これにより、勾配フローのための追加のチャネルが作成され、長距離のクレジット割り当てにおける指数関数的な減衰が緩和される。

2. 単一イテレーション・ニュートン・シュルツ（NS）正規化

スペクトルの不均衡に対処するため、入力注入（ランク1行列）は、モメンタム状態に加算される前に、軽量な単一ステップのニュートン・シュルツ（Newton–Schulz, NS）イテレーションを用いて正規化される。

• 演算子: $\text{NS}(X) = (a + b \tilde{X}\tilde{X}^\top + c (\tilde{X}\tilde{X}^\top)^2)\tilde{X}$。ここで、$\tilde{X}$ はフロベニウスノルムによってスケールされた入力である。
• 機能: この演算子は特異値を（1の近くに）拘束し、明示的なSVD（特異値分解）を必要とせずに入力のランクを保持する。
• 幾何学的効果: 標準的なフロベニウス正規化とは異なり、NS演算子のバックプロパゲーションにおけるヤコビアンは、現在のランク1の書き込みに対して直交する方向を増幅する。これにより「ランクの豊穣化（rank enrichment）」が促進され、メモリ状態が低次元の部分空間へと崩壊することを防ぐ。

並列化可能性

メモリとモメンタムの結合されたダイナミクスは、**ブロックアフィン再帰（block-affine recurrence）**として記述できる。これにより、モデルは並列結合スキャン（parallel associative scans）を利用することができ、再帰の深さを $O(L)$ から $O(\log L)$ に削減し、計算複雑性がシーケンス長 $L$ に対して線形であることを保証する。

主な貢献

1. フレームワーク設計: スキャンベースの再帰にモメンタムベースの二次のダイナミクスと暗黙的なスペクトル条件付けを統合したMuonSSMの導入。
2. 理論的分析:
   • 拡張されたダイナミクスが並列化可能性を維持することを証明。
   • モメンタム経路が非消失の勾配成分を導入し、長距離伝搬を安定化させることを実証。
   • NS演算子が特異値を拘束し、直交する方向へバックワード信号をバイアスすることでランクの豊穣化を駆動することを示す分析。
3. 実証的検証: 言語、視覚、時系列の各モダリティにおけるタスクにおいて、一貫した改善を示す包括的な実験。

実験結果

著者らは、言語、視覚、時系列の3つのモダリティにわたって、MuonSSMをベースライン（Mamba, LongHorn, Gated DeltaNet）と比較評価した。

• 言語モデリング: FineWeb-Eduデータセットにおいて、MuonSSMのバリアントはより低いパープレキシティと、常識推論ベンチマークにおける優れた性能を達成した。特に、「Needle-in-a-Haystack（ haystackの中の針）」検索タスクにおいて、MuonSSMは最大8Kのコンテキスト長まで高い精度を維持し、訓練ホライゾン（2Kトークン）を超えると急速に性能が低下するベースラインを大幅に上回った。
• 視覚: MambaVisionフレームワークにおいてMambaミキサーをMuonSSMに置き換えることで、ImageNet-1K分類、物体検出（COCO）、セマンティックセグメンテーション（ADE20K）において一貫した利得が得られた。決定的なことに、MuonSSMは、破損（IN-C）、レンダリング（IN-R）、および敵対的攻撃（IN-A）に対する堅牢性が大幅に向上しており、テクスチャ・バイアスから形状不変な特徴学習へのシフトを示唆している。
• 時系列: ヒューマン・アクティビティ・レコグニション（Human Activity Recognition）タスク（MuWiGes, UESTC-MMEA-CL, MMAct）において、MuonSSMのバリアントはベースラインを凌駕し、データセットの複雑さが増すにつれてその性能差が拡大した。

アブレーションと分析:

• スペクトル条件付け: 可視化により、標準的なSSMはスペクトルの崩壊（高い条件数）に苦しむ一方で、MuonSSMはほぼ等長（isometric）なスペクトルを維持することが確認された。
• 容量 vs 幾何学: ベースラインの状態次元を2倍にしても、MuonSSMによる大幅な改善と比較すると、わずかな利得しか得られなかった。これは、恩恵が単純な容量拡張ではなく、幾何学的条件付けに由来することを確認している。
• 効率性: ステップあたりの定数倍のオーバーヘッドはあるものの、MuonSSMは収束を加速させ（約1.3倍）、目標性能に達するまでの総訓練ステップ数を削減する。

重要性と主張

本論文は、更新の幾何学的条件付けが、安定してスケーラブルなシーケンスモデリングへの原則的な経路であることを提唱している。モメンタムとニュートン・シュルツ正規化を通じてメモリ更新の幾何学を明示的に制御することにより、MuonSSMは、並列スキャンの線形時間の効率性を損なうことなく、一次のSSMが持つ根本的な不安定性に対処している。

著者らは、このアプローチが以下を実現すると主張している：

• 勾配伝搬を改善し、スペクトル増幅を抑制する。
• ランクの崩壊を防ぐことで、長いホライゾンにわたってメモリ表現を豊かにする。
• 長コンテキストのタスクにおいてアテンションに代わる堅牢な選択肢を提供し、安定性は複雑な遷移行列のパラメータ化だけでなく、更新ダイナミクスを通じて達成できることを示す。

本研究は、将来のシーケンスモデリングにおいて、更新メカニズムの安定性を再帰構造から切り離すことが、多様なモダリティにわたって安定した学習を実現するためのシンプルかつ効果的なメカニズムを提供することを示唆している。