技術要約: PRISM (Parallel Residual Iterative Sequence Model)

1. 問題提起

生成的なシーケンスモデリングは、Transformerベースのアーキテクチャが持つ**表現力（expressivity）と、線形シーケンスモデルの効率性（efficiency）**との間の根本的な緊張関係に直面している。

• 効率性のボトルネック: 既存の効率的なアーキテクチャ（例：Linear Attention、MambaやDeltaNetのようなSSM）は、オンライン凸最適化（OCP）目的関数に対する線形化されたワンショット解に依存している。これらの手法はハードウェア効率の高い並列スキャンをサポートしているが、理論的にはランク1のボトルネックによって制約されている。各トークンのメモリ状態への寄与は、構造的に単一の外積（$v_t k_t^\top$）に限定されており、これが高次の特徴量相互作用を捉えたり、単一ステップ内で再構成誤差を反復的に最小化したりする能力を制限している。
• 表現力のボトルネック: TTT（Test-Time Training）のような強力な反復手法は、多段階の非線形勾配降下を実行することでランク1の制限を解消し、**ランク-L更新（Rank-L updates）を実現する。しかし、これは直列依存のボトルネック（Serial Dependency Bottleneck）**を導入する。すなわち、各勾配ステップは前のステップから更新された状態に依存するため、ハードウェアの並列性が損なわれ、$O(NL)$ の逐次計算が発生する。これにより、TTTは線形回帰よりも大幅に遅くなる（最大174倍）。

解決すべき中心的な問いは、**「いかにして並列スキャンの互換性（高い効率性）を維持しながら、ランク-L更新（高い表現力）を実現できるか？」**である。

2. 手法: PRISM アーキテクチャ

PRISMは、TTTの「ゲート $\times$ 残差 $\times$ 方向」の反復パターンを、並列化可能な形式へと明示的に再構築することで、この緊張関係を解消する。本アーキテクチャは、以下の2つのコアメカニズムに基づいている。

2.1. 書き込みと忘却のデカップリング（Write-Forget Decoupling）

PRISMは、スペクトル摂動解析に基づき、状態更新 $S_t = S_{t-1} A_t + B_t$ を2つの異なるパスに分解する：

• 忘却（Forgetting, $A_t$）: 忘却演算子は近似誤差に対して頑健である。PRISMはこれを、構造化された効率的なランク1の形式（Gated DeltaNetと同一）に固定し、複雑な非線形性を必要とせずに長期的なメモリの安定性を確保する。
• 書き込み（Writing, $B_t$）: 注入演算子（injection operator）はランクの忠実度に対して極めて敏感である。PRISMはここに表現力を割り当て、反復ソルバーを通じて**ランク-L注入（Rank-L injection）**を構築する。

2.2. 入力アンカーによるループ展開（Input-Anchored Loop Unrolling）

明示的なソルバー（残差や勾配の方向が進化する隠れ状態 $S_{t-1}$ に依存するもの）の直列依存性を回避するため、PRISMは多段階の解を単一の並列パスへと「償却（amortize）」する2段階のプロキシメカニズムを導入する：

1. ステージ1：アンカー構築: 入力シーケンスに対する短縮畳み込み（ShortConv）を用いて、状態に依存しないプロキシ $u_t \approx S_{t-1} k_t$ を計算する。これは局所的な履歴エネルギーを捉え、その後のすべての計算のアンカーとして機能するため、実際の隠れ状態への依存性を排除できる。
2. ステージ2：並列精緻化ループ: アンカー $u_t$ を使用して、$L$ 回の精緻化ステップを並列に展開する。各ステップ $l \in \{1, \dots, L\}$ において：
   • 方向（Direction）: 学習された射影 $k_t^{(l)} = W_k^{(l)} u_t$ が、ランク-$l$ 更新の基底を提供する。
   • ゲイン（Gain）: 学習された予測器 $p_t^{(l)} = W_p^{(l)} u_t$ が、文脈的ゲイン $\sigma'(S_{t-1} k_t)$ を近似する。
   • 残差更新（Residual Update）: 残差は貪欲に減算される：$r_t^{(l+1)} = r_t^{(l)} - \delta_t^{(l)}$。ここで $\delta_t^{(l)} = \text{GELU}(p_t^{(l)} \odot r_t^{(l)})$ である。
   • 信頼度（Confidence）: ステップサイズゲート $\beta_t^{(l)}$ が更新を調整する。

最終的な注入行列は $B_t = \sum_{l=1}^L \beta_t^{(l)} \delta_t^{(l)} (k_t^{(l)})^\top$ として累積される。すべてのコンポーネントは（アンカーを介して）入力のみに依存するため、全更新は並列プレフィックススキャンに必要な「状態独立条件」を満たしている。

3. 主な貢献

1. 統一されたOCPの視点と分類学: 本論文は、既存の手法を共通の最適化フレームワークの中に位置づけ、ランク1/並列性のトレードオフを中央のボトルネックとして特定した。そして、PRISMがランク-Lの表現力と並列スキャンの互換性の両方を達成した最初の手法であることを確立した。
2. PRISM アーキテクチャ: 書き込みと忘却のデカップリングおよび入力アンカーによるループ展開を利用した、ハードウェア認識型の実装。この設計により、明示的なソルバーの逐次計算を回避しながら、その表現力を保持している。
3. 理論的な表現力の保証: 著者らは、入力アンカー形式が**ランク蓄積（Rank Accumulation）**をもたらすことを正式に証明した。PRISMの注入演算子は、トークンごとにランク-Lの状態変更を生み出し、理想的な非線形デルタ則の解の品質に漸近的に近づく。さらに、$L=0$ のとき、PRISMは構造的にGated DeltaNetへと縮退し、後方互換性を確立している。

4. 実験結果

4つのレコメンデーション・ベンチマーク（Amazon Books, Movies, Electronics, Yelp）および言語モデリング・タスク（SlimPajama）において実験が行われた。

• 性能（ランク-Lの有効性）: PRISMは、TTTやTitansのような明示的な逐次ソルバーと同等の性能を達成し、いくつかのケースではそれらを上回った。Amazon Moviesにおいて、PRISMはすべての線形モデルの中で最高のAUCを記録し、Titansを上回った。また、フルアテンションTransformer（例：SASRec）との差も縮めており、「線形化による表現力の損失（expressivity tax）」がランク-Lの追跡によって無視できるレベルになることを示唆している。
• 効率性（並列スキャン）: PRISMは、様々なシーケンス長にわたって安定したトレーニングスループット（~57K–61K tokens/s）を維持し、ランク1のベースライン（Gated DeltaNet, Mamba-2）と同等の性能を発揮する。対照的に、TTTはその逐次的な性質ゆえに174倍遅く、Transformer++はシーケンス長に対して二次関数的に劣化する。
• アブレーション研究:
  • 反復の深さ: 反復的な精緻化を取り除くと（$L=0$）、最大の性能低下が発生し、多段階の精緻化が極めて重要であることが検証された。
  • 非線形性: 精緻化ループにおけるGELU活性化関数を取り除くと性能が低下し、最適化品質の高い更新には非線形な成形が必要であることが確認された。
  • メカニズムのプロービング: ランク1のモデルが失敗する論理タスク（XOR, Parity）において、PRISMは100%の精度を達成し、その恩恵が単なる空間的ランクの拡大ではなく、反復的な非線形計算から来ていることを証明した。

5. 意義と主張

本論文は、PRISMが最適化の忠実度とハードウェア効率の高い並列トレーニングの間の溝を埋めることに成功したと主張している。反復的な補正の構造的パターンを並列化可能なフィードフォワード演算へと蒸留することで、PRISMはランク-並列のトレードオフを解決する。

著者らは、この設計により、効率的なシーケンスモデルが、線形回帰のスループット上の利点を犠牲にすることなく、複雑でマルチモーダルなユーザーの関心や長距離の依存関係を捉えられるようになると断言している。PRISMは単なる理論的な改善ではなく、ランク1の障壁を超えるための実用的な経路として提示されており、明示的な最適化の表現力と線形モデルの効率性の両方を提供している。生成モデリングにおいて、局所的なコンテキストは最適化の軌道を予測するための非常に効率的な構造的事前分布を提供しており、正確な状態依存性からデカップリングするという戦略が、スケーリングのための実行可能な手段であることを示唆している。