技術的サマリー：インタードメイン・アテンション

問題定義

トランスフォーマーと深層状態空間モデル（SSM）は、シーケンスモデリングにおける対照的な設計パラダイムを表しています。標準的なアテンション機構（Vaswani et al., 2017）は、成長するキー・バリュー（KV）キャッシュを利用し、コンテンツベースのマッチングを通じてクエリをルーティングします。これにより鋭い想起が可能になりますが、シーケンス長に対して二次的な計算コスト（$O(N_q N)$）と線形なメモリ増大（$O(N)$）を伴います。一方、深層 SSM（例：S4、Mamba）は、全体のコンテキストを固定サイズの再帰状態に圧縮し、各ステップで $O(1)$ の時間で更新します。しかし、この圧縮は直接的なクエリ・キーマッチングの欠如という代償を伴います。固定状態は特定のクエリによって直接アドレス指定できないため、精密な検索やコンテンツベースのアテンションを必要とするタスクでは、しばしば性能が低下します。

ハイブリッド・アーキテクチャはこれらのアプローチを交互に組み合わせる試みを行ってきましたが、線形アテンションの変種はアテンションコストを直接削減しようとしてきました。本論文は異なる道筋を提案します。つまり、固定サイズ状態の効率性を維持しつつ、各クエリが圧縮された履歴にアテンションできるようにするため、SSM をアテンション機構に直接統合するというアプローチです。

手法：インタードメイン・アテンション

中核的な貢献は、カーネル・アテンションと状態空間モデルを統合するトークン・ミキサーであるインタードメイン・アテンションです。この手法は、アテンションのカーネル回帰的視点と HiPPO 様式の基底射影から導出されます。

理論的導出

1. カーネル特徴マップ: アテンション・カーネル $K(q, k)$ は、有限の特徴マップ $\xi(q)^\top \xi(k)$ によって近似されます。
2. 基底射影: すべてのキーとバリューを格納する代わりに、本手法はキー特徴 $\xi(k)$ とバリュー $v$ を、単一の SSM 再帰（具体的には HiPPO 様式）によって維持される共通の基底関数の集合 onto 射影します。これにより、シーケンス長に依存しない固定サイズの状態が得られます。
3. 圧縮係数: これらの射影の積分は、3 つの状態成分を生成します。
   • $U$: キーの基底射影係数（誘導変数）。
   • $\Gamma$: バリュー射影。
   • $\eta$: 正規化定数。
4. クエリ条件付き読み出し: 各位置において、クエリ $q$ は特徴空間 $\xi(q)$ へマッピングされます。出力は、個々のトークンではなく圧縮された係数にアテンションすることで計算されます。
   $$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$$
   行列形式では：$\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$。ここで $F_q$ はクエリ特徴行列です。

スケーラブルな実装

実用的な実装は、理論的導出の学習された緩和版であり、安定性と効率性のために設計されています。

• 特徴マップ: $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$ という、DeltaNet 様式のニューラル特徴マップを使用します。ここで SC は因果的な深度方向畳み込みです。
• 再帰コア: 複素対角 S4D 再帰を用いて HiPPO 射影を近似します。SSM への入力は、キー特徴 $\xi(k)$ とバリュー $v$ を連結したものです。
• 正規化: SSM に入る前に、キーとバリューの半分に、学習可能なヘッドごとのバイアスを適用した入力 RMSNorm を適用します。
• 分母不要な読み出し: 入力のリサイズによって引き起こされる互換性のないスケールの混合を避けるため、正規化項 $\eta$ を削除し、線形アテンション読み出し $\tilde{O} = F_q U^\top \Gamma$ を実現します。
• アーキテクチャ: 事前正規化、SwiGLU フードフォワード、ロータリー位置埋め込み（RoPE）を備えた Llama 様式のデコーダーに埋め込まれています。

主要な貢献

1. 構築: カーネル回帰と HiPPO 基底射影から導出された、クエリ条件付きの固定状態トークン・ミキサー。理想の理論的導出とスケーラブルな実装の間に明確な区別がある。
2. メカニズムの分解: 1 億 2500 万パラメータにおけるアブレーション研究により、標準的な S4D ミキサーに対する性能向上を主導する要因としてクエリ条件付き射影が特定された。キー/バリューの二重入力は、より小さく二次的な寄与をもたらす。
3. スケーリング研究: 一致した再帰状態予算の下で、FineWeb-Edu における 1 億 2500 万から 13 億パラメータまでの包括的な評価。

実験結果

本研究は、インタードメイン・アテンションを、S4D 単一のトークン・ミキサーおよび標準的な softmax アテンションのベースラインと比較します。

• S4D に対する性能: インタードメイン・アテンションは、すべてのスケール（1 億 2500 万〜13 億）において S4D トークン・ミキサーを上回り、検証パープレキシティを 13〜16% 低下させます。
• Softmax に対する性能: 13 億パラメータにおいて、インタードメイン・アテンションは、同じレシピの softmax ベースラインを、検証パープレキシティ（7.98 対 8.63）および 8 タスクの常識性スイート（平均精度で +3.03 パーセントポイント）において上回ります。1 億 2500 万パラメータでは、softmax とほぼ同程度です。
• 長さ外挿: 固定状態コアを活用し、インタードメイン・アテンションは「長さ平坦」な挙動を示します。softmax ベースラインのパープレキシティはトレーニングコンテキスト（4K）を超えると著しく劣化するのに対し、インタードメインはトレーニングコンテキストの 3.5 倍である 14K まで安定した性能を維持します。
• 想起の限界: 本論文は、インタードメイン・アテンションが他の固定状態モデルと同様に、softmax アテンションと比較して正確な文字列の連合想起（例：干し草の山の中の針タスク）に苦戦することを指摘しています。これはインタードメイン機構固有の欠陥ではなく、固定状態圧縮の構造的限界であると特定されています。

意義と主張

本論文は、インタードメイン・アテンションが、固定状態 SSM の効率性と、アテンションのコンテンツベースの検索能力との間のギャップを成功裡に埋めたと主張しています。特徴を SSM 基底 onto 射影することで、標準的なアテンションの二次コストや線形メモリ増大なしに、クエリが圧縮された履歴にアテンションすることを可能にします。

主な意義は、固定再帰状態 onto のクエリ条件付き射影が、純粋な SSM およびより大規模なスケールでは標準的な softmax アテンションを上回る性能を発揮し、かつ SSM が持つ望ましい長さ外挿特性を継承することを示した点にあります。著者らは、このアプローチが、メモリ効率と異なるコンテキスト長にわたる一貫した性能が重要な長コンテキスト言語モデリングのための実行可能な道筋を提供すると示唆しています。将来の研究方向としては、より強力な固定状態コア（例：Mamba-3）の統合、または読み出しと高速重み更新規則の組み合わせが含まれます。