技術要約：Sparse Delta Memory (SDM)

問題提起

Gated DeltaNet (GDN) や Mamba2 といった線形再帰型ニューラルネットワーク（Linear RNN）は、固定されたメモリと計算量を提供するため、理論的には任意の長さのコンテキストを扱うのに適しています。しかし、それらの実用的な性能は、隠れ状態（hidden state）のサイズが小さいことによって、長文脈の想起（recall）やインコンテキスト学習において制限されています。標準的な線形RNNにおいて状態サイズを大きくすると、密な状態更新（行列とベクトルの積）によるFLOPsが線形に増加し、計算上のボトルネックとなります。一方で、ソフトマックス・アテンションを用いる標準的なTransformerアーキテクチャは、高い想起能力を維持しますが、シーケンス長が増加するにつれてKey-Value (KV) キャッシュのメモリおよび計算コストが線形に増大するという課題があります。

解決すべき核心的な課題は、1トークンあたりの計算コスト（FLOPs）を増やすことなく、いかにして線形RNNの隠れ状態容量を数桁規模でスケールさせ、長文脈性能を向上させるかという点です。

手法：Sparse Delta Memory (SDM)

著者らは、密なキー・バリューの外積を、Product-Key Memory (PKM) に基づく疎なアドレッシング・スキームに置き換えることで、Gated DeltaNetを拡張したアーキテクチャである Sparse Delta Memory (SDM) を提案しています。

コアメカニズム

1. 明示的な疎メモリ (Explicit Sparse Memory): 密な行列 $M_t \in \mathbb{R}^{d_{qk} \times d_v}$ を保持する代わりに、SDMは $N$ 個のスロットを持つ明示的なメモリテーブル $M_t \in \mathbb{R}^{N \times d_v}$ を保持します。ここで、$N$ は隠れ次元よりも数桁大きい値（例：$N \approx 10^6$）を取ることができます。
2. 疎なキー選択 (Sparse Key Selection): 各タイムステップにおいて、投影されたキー ($k'$) とクエリ ($q'$) は2つの半分に分割されます。それらの外積和（outer sum）により、メモリ・スロットに対する $N$ 個のスコアが生成されます。トップ-$W$（書き込み）およびトップ-$R$（読み取り）の選択戦略を用いることで、モデルは更新または読み取りを行う特定のメモリ・スロットを特定します。この選択は、$\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$ という性質を利用することで、フルサイズの $N$ スコア行列を実体化することなく効率的に行われます。
3. ゲート付きデルタ更新 (Gated Delta Updates): 選択された $W$ 個の書き込みスロットに対して、モデルは以下のGated DeltaNet更新規則を適用します：
   • 忘却ゲート (Forget Gate): $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • デルタ更新 (Delta Update): $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • 選択されなかったスロットは変更されません。
4. 疎な読み取り (Sparse Read): 出力は、選択された $R$ 個の読み取りスロットの重み付き和となります。

IsoFLOP デザイン

公平な比較を行うため、著者らはSDMを密なGDNベースラインのパラメータ数およびFLOPsと一致するように設計しています。

• FLOPs: 1トークンあたりのコストは $O((W+R) \times d_v)$ であり、メモリサイズ $N$ に依存しません。$W$ と $R$ をGDNベースラインのクエリ次元と等しく設定することで、計算コストは同一に保たれます。
• パラメータ: 線形投影のサイズはベースラインと同一です。メモリサイズ $N$ は、ヘッド数 $H$ と投影次元を通じて制御され、FLOPsに影響を与えることなく、状態サイズを $O(1/H^2)$ としてスケールさせることが可能です。

学習された初期状態 (Learned Initial State)

通常、小さな状態をゼロに初期化するGDNとは異なり、SDMは初期メモリ状態 $M_0$ を学習可能なパラメータとして扱います。これにより、モデルは事前学習の知識をメモリ構造の中に直接保持することができ、推論を通じて持続するパラメトリックなメモリとして機能します。

主な貢献

1. アーキテクチャ: 1トークンあたりの計算量を一定に保ったまま、線形RNNの状態容量を3桁スケール（例：キロバイトからギガバイトへ）させるSDMの提案。
2. IsoFLOP による検証: 状態容量を増やすこと（計算量を増やさずに）が、性能向上に大きく寄与することを実証。
3. パラメトリックメモリ: 線形RNNにおける学習された初期状態 $M_0$ の導入。このメカニメントが共通知識を効果的に蓄積し、推論タスクを改善することを示した。
4. スケーリング則: SDMが予測可能なスケーリング則に従うことを実証し、密なGDNを上回り、8Bパラメータ規模ではFull Attentionモデルの訓練損失をも上回ることを示した。

実験結果

著者らは、1兆トークル以上のデータを用いて、1.4Bから8Bパラメータに及ぶモデルを訓練しました。

• スケーリング則: スケーリング・ラダーにおいて、SDMはすべての計算レベルでGDNよりも一貫して低い訓練損失を達成しました。8Bスケールでは、SDMはGDNおよびFull Attentionモデルの両方よりも低い損失に到達しました。
• 長文脈想起 (RULERベンチマーク): SDMは、長文脈の検索タスクにおいてGDNを大幅に上回りました。1.4Bおよび8Bスケールにおいて、SDMはそれぞれ31.2および50.2を記録し、GDNの20.0および34.2と比較して高いスコアを達成しました。特筆すべきは、SDMが固定されたメモリ表現を用いているにもかかわらず、いくつかのRULERサブタスクにおいてFull Attentionの性能に匹敵、あるいはそれを上回ったことです。
• 短文脈および推論: SDMは、コードデータにおけるパープレキシティ（perplexity）が低く、推論ベンチマーク（MMLU、GSM8K、HumanEvalなど）における平均精度がGDNよりも高くなりました。8Bスケールでは、SDMは平均精度においてFull Attentionを上回りました。
• アブレーション研究:
  • 状態サイズ: メモリサイズを縮小すると性能が単調に低下し、大きな状態が利得の主要な要因であることを確認しました。
  • 学習された初期化: $M_0$ を学習することはSDMの性能を向上させましたが、GDNにはほとんど効果がありませんでした。これは、学習された初期化を効果的に利用するには大きな容量が必要であることを示唆しています。
  • 効率性: 8BスケールにおけるSDMの訓練スループットは、HBMへのメモリアクセスの影響によりGDNより1.49倍遅かったものの、推論時にはFull Attentionよりも6倍高速でした。

意義と主張

本論文は、SDMが線形RNNの効率性とTransformerの長文脈能力の間の溝を埋めることに成功したと主張しています。疎性を利用してメモリ容量と計算コストを切り離すことで、SDMは以下を可能にします：

• 一定の空間/計算量: KVキャッシュの無制限な増大なしに、任意の長さのシーケンス（最大100万トークンを実証）を処理できます。
• 優れた長文脈性能: 長文脈の検索およびインコンテキスト学習において、線形RNNの中で最先端の結果を出し、Full Attentionに対しても競争力のある結果を出します。
• パラメトリックな知識蓄積: 事前学習の知識をメモリ状態自体に学習・保持する能力を持ち、一般的な推論や共通知識のタスクを向上させます。

著者らは限界についても認めています。具体的には、大きなメモリ・フットプリント（モデルパラメータに匹敵する）が多大なHBMリソースを必要とすることや、現在のカーネル実装が最適化された密なGDNカーネルよりも効率が低いことです。しかし、彼らはSDMが、より優れた長期記憶を持つエージェントの開発や、長時間のビデオ、あるいは拡張された推論トレースの処理のための新たな可能性を切り開くものであると考えています。