🤖 machine learning

Raven: High-Recall Sequence Modeling with Sparse Memory Routing

Ravenは、固定されたメモリスロットのサブセットのみを更新するために疎で入力依存的なメモリルーティングを採用することで、高密度な状態空間モデルの干渉問題とスライディングウィンドウ・アテンションの硬い退避制限との間のバランスを効果的にとり、長文脈の想起能力を向上させた線形時間シーケンスモデルである。

原著者： Arshia Afzal, Aviv Bick, Eric P. Xing, Volkan Cevher, Albert Gu

公開日 2026-07-29

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Arshia Afzal, Aviv Bick, Eric P. Xing, Volkan Cevher, Albert Gu

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

ある物語を聞いた直後のことを思い出そうとしていると想像してください。しかし、あなたの脳には奇妙なルールがあります。新しい単語を聞くたびに、これまでに書いたすべての文章にその新しい単語を混ぜ込みながら、物語全体を最初から書き直さなければならないというルールです。これは、「状態空間モデル（State-Space Models）」と呼ばれる一部のコンピュータプログラムがどのように機能しているかを示しています。これらのモデルは、長い間物語の全体的な雰囲気（バイブス）を維持することには長けていますが、すべてを徹底的に混ぜ合わせてしまうため、その中に埋もれた名前や日付といった特定の詳細を見つけ出すことは不可能です。一方で、別の種類の記憶を想像してみてください。それは、決まったページ数のノートのようなものです。最後のページがいっぱいになると、そのページを破り捨てて新しいものと入れ替えます。これは「スライディング・ウィンドウ（Sliding Window）」モデルの仕組みです。これらは、直近の数ページを完璧に覚えていることができますが、情報が端に押し出された瞬間に、それは永遠に消えてしまいます。

これは、人工知能の分野の科学者たちが解決しようとしている大きな問題です。詳細を混乱させることなく、かつ古いからといって最初の方を忘れてしまうこともなく、非常に長い間物語を記憶できるコンピュータの脳を、どうすれば構築できるでしょうか？私たちは、情報を効率的に処理しながら、特定の重要な事実を長期間保持できるシステムを必要としています。これは、AIがより賢くなるにつれ、より長い本を読み、巨大な文書を分析し、会話の最初に行われた指示を、数千語が経過した後でも覚えておく必要があるため、極めて重要です。

そこで登場するのが、まさにこのメモリのパズルを解くために研究者によって設計された、新しい種類のAIモデル「Raven」です。Ravenのメモリを、ぐちゃぐちゃな書き直しや、単純なゴミ箱型のノートではなく、何百ものロッカーがあるハイテクなロッカールームとして考えてみてください。旧来のシステムでは、新しいアイテムが到着するたびに、それは一度にすべてのロッカーへと強制的に押し込まれるか、あるいは重要かどうかに関わらず、到着順に基づいて最も古いアイテムを追い出す形でロッカーに押し込まれていました。Ravenはこのルールを変えました。Ravenは、ボディーガードのような役割を果たすスマートな「ルーター」を使用します。新しい情報が到着すると、ボディーガードはその内容を確認し、こう判断します。「これは退屈な事実か？それなら頻繁に掃除される共有ロッカーに入れろ。これは超重要な秘密か？それなら誰も触れることのできない特別な専用ロッカーに入れろ」。

この論文は、Ravenの設計図である「ルーティング・スロット・メモリー（Routing Slot Memories: RSMs）」というフレームワークを紹介しています。鍵となる革新は、Ravenが「情報の書き込み場所」と「その情報がどれくらい長く留まるか」を切り離している点にあります。以前のモデルでは、これら二つの要素は絡み合っていました。Ravenは「スパース（疎）」なルーティング・システムを使用しており、これは各新しいテキストに対して、選択された少数のグループ（メモリスロット）のみを更新し、残りの部分は一切手を触れずにそのままにしておくことを意味します。これは非常に重要なことで、なぜなら、一度にすべてを更新しようとすると発生する「干渉」を防ぐことができるからです。もし特定のロッカーが極めて重要なパスワードを保持している場合、Ravenのルーターはそのロッカーを無視することを選択できるため、他のメモリがその役割を果たしている間も、パスワードを安全にそこに置いたままにすることができるのです。

研究者たちは、いくつかの非常に困難なメモリゲームでRavenをテストしました。その一つが、巨大な文書の中に隠された特定の文章を見つけ出す「針の穴探し（Needle in a Haystack）」テストです。これらのテストにおいて、Ravenは、学習時の長さの16倍にも及ぶ文書の中でも、その「針」を見つけ出すことができました。Mamba-2やスライディング・ウィンドウ・アテンションのような他のモデルが、テキストが長くなるにつれて失敗したり忘れたりし始める中で、Ravenは完璧に近い精度を維持しました。例えば、32,000トークンのテストにおいて、Ravenは91%以上の精度を維持しましたが、他のモデルは大幅に低下しました。

本当に素晴らしいのは、Ravenがこれを行うために特別な追加のトリックを必要としないことです。他のモデルで使用される複雑な畳み込み（コンボリューション：短期記憶フィルターのようなもの）や、特別な位置マーカーにも依存していません。ただ、このスマートなルーティング・システムを使用しているだけなのです。論文は、モデルが単なる「位置（文の中のどこにあるか）」ではなく、「コンテンツ（その単語が何か）」に基づいてメモリの割り当てを学習させることで、自然な「専門化」が生み出されることを示唆しています。一部のメモリスロットは、検索に不可欠な詳細を保持するエキスパートとなり、他のスロットは物語の一般的な流れを扱うようになります。

著者らは、このアプローチが単独で機能するだけでなく、他の種類のAIアーキテクチャと組み合わせた場合にも機能することを発見しました。Ravenを標準的なアテンション・メカニズム（現在最も人気のあるAIモデルで使用されているもの）と組み合わせたハイブリッド・システムは、スライディング・ウィンドウ・アテンションや他の線形モデルを使用したモデルよりも、ロングコンテキスト・タスクにおいて優れた性能を発揮しました。論文は、メモリの割り当てを意図的で学習可能な選択肢として扱うことが、よりスマートで効率的なAIを構築するための強力な方法であることを示唆しており、Ravenが長期的な持続性と精密な想起の両方に優れたモデルとの間の溝をうまく埋めたことを結論付けています。

技術要約: Raven — 疎なメモリ・ルーティングによる高リコール・シーケンスモデリング

1. 問題提起

SSM（状態空間モデル）や線形Transformerといった線形時間シーケンスモデルは、長文脈のリコール（想起）におけるメモリ管理において、根本的なトレードオフに直面している。既存のアーキテクチャは一般に、以下の2つの極端な形態のいずれかに分類され、それぞれ特有の限界を抱えている：

高密度な状態更新 (SSMs): Mamba-2、GLA、GDNなどのモデルは、固定サイズの状態を維持し、毎ステップ**高密度（dense）**に更新する。すべてのメモリ・スロットが新しいトークンごとに更新され、学習された減衰メカニズムに頼って忘却を行う。これは原理的には無期限の持続を可能にするが、高密度な更新は干渉を引き起こし、後続の情報によって特定の個々のトークンが上書きされたり希釈されたりすることを防ぎ、特定のトークンを保持することを困難にする。
疎な位置ベースの書き込み (Sliding-Window Attention - SWA): SWAは、明示的なトークンの表現を格納することで**疎（sparse）**に書き込むが、固定された直近のウィンドウ内に制限される。これはFIFO（先入れ先出し）戦略を用いており、位置に基づいて古いトークンを強制的に排除（エビクション）する。これにより、ウィンドウ内での信頼性の高い想起が可能になるが、関連するトークンが排除されると、その重要性に関わらずリコール率はほぼゼロに低下する。

核心となる問題は、既存のモデルが「情報の書き込み場所」と「情報の持続期間」を混同していることである。SSMはあらゆる場所に書き込み、減衰に頼って忘却する。SWAは（現在のウィンドウ以外には）どこにも書き込まず、位置に頼って忘却する。どちらのモデルも、疎でコンテンツ依存の書き込みと、明示的かつ制御可能な減衰を効果的に組み合わせることはできていない。

2. 手法: ルーティング・スロット・メモリ (RSMs)

著者らは、情報の書き込み場所と持続の期間を切り離す統一フレームワークとして、ルーティング・スロット・メモリ (RSMs) を提案する。RSMsは固定された $M$ 個のメモリ・スロットを維持し、各タイムステップ $t$ において、どのサブセットのスロットに新しい情報を書き込み、減衰を適用するかを決定するルーターを使用する。

2.1 RSM フレームワーク

RSM の一般的な漸化式は以下のように定義される：
$S_t = (1 - r_t) \odot S_{t-1} + r_t \odot (D_t S_{t-1} A_t + U_t)$
ここで：

$S_t \in \mathbb{R}^{M \times d}$ は、 $M$ 個の独立したスロットを持つメモリ状態である。
$r_t \in \mathbb{R}^M$ は、各スロットの書き込み強度を決定するルーティング・ベクトルである。
$D_t$ と $A_t$ は、減衰と特徴量の混合を制御する。
$U_t$ は、入力 $x_t$ から導出される書き込み内容である。

このフレームワークは、既存のモデルを以下のスペクトラム上に位置づける：

SSMs: 高密度なルーター ( $r_t = \mathbf{1}_M$ ) を使用し、緩やかな減衰を伴ってすべてのスロットを更新する。
SWA: 時間依存の一熱（one-hot）ルーター ( $r_t = e_t$ ) を使用し、正確に1つのスロットを更新し、強制的な排除（減衰なし）を行う。
GSA/ABC: 入力依存のルーターを使用するが、通常はすべてのスロットを更新する（高密度な書き込み）。これは、メモリを隔離する能力を制限する。

2.2 Raven アーキテクチャ

Raven は、疎な入力依存ルーティングと明示的な減衰を組み合わせることで、「中間領域」を占めるように設計された RSMs の具体的な実装である。

デュアル・ステート構造: SWAと同様に、Raven は分離された Key ( $S^k_t$ ) および Value ( $S^v_t$ ) 状態を維持する。
疎な入力依存ルーティング: 固定ウィンドウや高密度な更新の代わりに、Raven は入力 $x_t$ $x_{t}$ から導出されるルーター $g_t$ $g_{t}$ を使用する。DeepSeek MoE のルーティングから適応した Top-K 選択戦略を採用し、更新する少数のスロットのサブセット $K$ $K$ を選択する。
- スコアはシグモイド関数 $m_t = \sigma(W x_t)$ を通じて計算される。
- 上位 $K$ 個のエントリのみが保持され、それ以外はゼロに設定される。
- 選択されたスロットが適切な書き込み強度を受け取るよう、ルーターは正規化される。
選択的減衰: 減衰は選択されたスロットに対してのみ適用される。漸化式は以下の通りである：
$S^k_t = \exp(a_t r_t) \odot S^k_{t-1} + (1 - \exp(a_t r_t)) \odot k^\top_t$
ここで、 $a_t$ は入力依存のスカラー減衰（Mamba-2 に類似）である。スロットが選択されない場合（ $r_t[i] \approx 0$ ）、 $\exp(a_t r_t[i]) \approx 1$ となり、そのスロットは**凍結（フリーズ）**され、再び選択されるまでその内容を完全に保持する。
簡素化されたブロック設計: 多くの線形モデル（例：Mamba, GLA）とは異なり、Raven は短距離畳み込みを省略している。Raven は、局所および長距離のコンテキスト統合を、ルーティングと減衰のメカニズムのみに依存して行うため、よりシンプルなブロック構造を実現している。

3. 主な貢献

ルーティング・スロット・メモリ (RSMs): SSMs と SWA を統一し、「書き込み場所」（ルーティング）と「忘却メカニズム」（減衰）を明示的に分離する形式的なフレームワーク。
Raven モデル: 入力依存の疎なルーティングを用いてメモリ・スロットのサブセットを更新しつつ、選択されたスロットにのみ減衰を適用することで、RSMs を実装した線形時間シーケンスモデル。
長距離リコールのためのメカニズム: 検索に重要なトークン（例：パスキー）を専用のスロットにルーティングして凍結させることで、Raven は、高密度な SSM で共通の破壊的な干渉や、SWA のハード・エビクションを防ぐ。
長さの汎化: 明示的な長さのカリキュラムや位置補間（RoPE など）を用いずに、訓練長の最大 16倍 のコンテキストにおいて強力な性能を達成した。これは、スロット間の可変的な実効シーケンス長が、モデルを多様な時間的ダイナミクスに自然にさらしているためである。

4. 実験結果

著者らは、Mamba-2、GLA、GDN、SWA、および Transformer 変種（FoX）を含む強力なベースラインと比較して、リコール集中的なベンチマークおよび一般的な言語モデリング・タスクで Raven を評価した。

Needle-in-a-Haystack (NIAH):
- 400M スケール: Raven は 16K トークンまでほぼ完璧な精度（ $\ge 99\%$ ）を維持し、32K トークン（訓練長の 16 倍）でも高い精度（ $> 91\%$ ）を保った。対照的に、Mamba-2 と GDN は 8K を超えると著しく低下した。
- 800M スケール: Raven は NIAH-1 において 64K トークンで 91.0% の精度を達成し、次に優れたベースライン（GDN の 45.2%）を大幅に上回った。
実世界の想起タスク:
- SWDE（ウェブ抽出）、FDA（文書抽出）、および SQuAD において、Raven は一貫して線形ベースラインを上回った。例えば FDA では、線形モデルの中で唯一 22% の精度を超え、最高の SSM ベースラインよりも約 8 ポイント向上した。
言語モデリング:
- 標準的なゼロショット・ベンチマーク（Lambada, PIQA, HellaSwag など）において、Raven は高密度な SSM（Mamba-2, GLA）と同等またはそれ以上の性能を示し、疎なルーティングが一般的な言語モデリング能力を損なわないことを証明した。
ハイブリッド・アーキテクチャ:
- フル・アテンション層と組み合わせたハイブリッド・モデルにおいて、Raven は、長文脈のリコール・タスク（特に 16K–64K の範囲）において、Mamba-2、GDN、または SWA-RoPE を使用したハイブリッドモデルを大幅に上回った。

5. 重要性と主張

本論文は、Raven が線形シーケンスモデリングにおける根本的なアーキテクチャ上のギャップに対処していると主張している。どこに情報が書き込まれるかと、それがどのくらい長く持続するかを明示的に分離することで、Raven は、二次的なアテンションや複雑なハイブリッド構成に頼ることなく、線形モデルではこれまで到達不可能であったレベルの長距離リコールを実現した。

主な主張は以下の通りである：

SWA に対する優位性: 位置ベースの排除を、コンテンツ依存のルーティングと選択的減衰に置き換えることで、SWA が完全に失敗するタスクにおいても完璧なリコールを可能にした。
高密度 SSM に対する優位性: 選択的に書き込むことで高密度な状態更新による干渉を回避し、特定のトークンを専用のスロットに無期限に保持することを可能にした。
簡潔さ: 短距離畳み込みを用いることなくこれらの利点を得ており、従来の線形モデルよりもシンプルなアーキテクチャを実現している。
創発的な専門化: 分析によれば、Raven は検索に重要なトークン（例：パスワード）のために特定のメモリ・スロットを割り当て、他のスロットを一般的なコンテンツに使用するように自然に学習しており、明示的な監督なしに、構造化されたコンテンツ認識型のメモリ割り当てを構築している。

著者らは、RSM フレームワークは、ルーティングと忘却は、減衰メカニズムの副産物としてではなく、第一級のデザイン選択肢として扱うべき、直交する設計軸であることを示唆していると結論付けている。