数週間前、@EverMind 氏が「メモリスパースアテンション（MSA）」に関する論文を発表しました。この研究は、最大 1 億トークン（本 1,000 冊分に相当）までの情報を記憶し、性能低下を 9% 未満に抑えつつ正解を導き出すシステムを構築するものです。複数のベンチマークテストにおいて、このパラメータ数 40 億のモデルは、自らの 58 倍の規模を持つモデルを基盤とした RAG システムさえも上回る性能を示しました。
その中核となる発想は以下の通りです。個別のデータベースを検索し、正しい情報が見つかることを期待する（RAG アルゴリズムの動作原理）のではなく、記憶機能そのものをモデルの思考プロセスに直接組み込むというものです。これにより、何を記憶し何を無視するかを、個別の検索プロセスを介さずにエンドツーエンドで学習することが可能になります。
論文発表後、産業界に大きな反響を呼び、研究と共同進歩のためにソースコードの公開が強く待たれていました。そして昨日、EverMind 氏はその約束を果たし、プロジェクトコードの完全なオープンソース化を実現しました。公開から 24 時間以内に 2,500 のスターを獲得しており、この技術に対する関心の高さが伺えます。
プロジェクト概要
長期記憶は汎用知能の基盤ですが、フルアテンションの計算上のボトルネックにより、大半の大規模言語モデル（LLM）の実質的なコンテキスト長は128K〜1M トークンに制限されています。既存の解決策である混合線形アテンション、固定サイズの状態記憶（RNN など）、あるいはRAG/エージェントなどの外部ストレージは、極端な長さにおいて精度が急速に低下したりレイテンシが増大したりするか、エンドツーエンドの微分可能性や動的記憶管理機能を欠き、あるいは複雑なエンジニアリングパイプラインを必要とするという課題を抱えていました。私たちが提案するのはメモリスパースアテンション（MSA）です。これはエンドツーエンドで訓練可能かつ拡張可能なスパースな暗黙的状態記憶のフレームワークです。その中核的なアイデアは以下の通りです。
- 拡張可能なスパースアテンションとドキュメント単位の RoPE（並列/グローバル）を組み合わせ、訓練・推論の両方においてほぼ線形の計算量を実現します。
- KV キャッシュの圧縮とメモリの並列処理を備えた推論エンジンにより、2×A800 GPU上で1 億トークンのスループットを達成します。
- 分散した記憶断片にわたるマルチターン・マルチホップ推論を可能にするメモリの織り交ぜ（Memory Interleaving）メカニズムを搭載しています。
長いコンテキストを要する質問応答や NIAH（Needle In A Haystack：干し草の山の中の針）ベンチマークにおいて、MSAは同 backbone の RAG や最先端の RAG 手法、そして主要な長期コンテキストモデルを凌駕しました。前例のない16K から 1 億トークンの範囲において、MSA は9% 未満の性能低下に留まり、記憶容量と推論能力を分離するための実行可能な道筋を示しました。
16K から 1 億トークンへの拡張：MSA は top-k 選択とスパースアテンションを融合させ、推論時にドキュメントの非結合を可能にしながらもエンドツーエンドの微分可能性を維持します。MS MARCO において、MSA の性能低下は9% 未満であり、強力な外挿能力を示しています。一部のベースライン曲線はコンテキスト長の制限により早期に終了しています。
図 1: 超長文脈における MSA の拡張性
主な貢献
- メモリスパースアテンション（MSA）：エンドツーエンドで訓練可能な拡張可能なスパースアテンションレイヤーであり、ドキュメント単位の RoPEと組み合わせることでO(L)の計算量を実現し、16K から 1 億トークンへの拡張において性能低下を9% 未満に抑えます。
- KV キャッシュ圧縮＋メモリ並列処理：階層的ストレージ（ルートキーは GPU、コンテンツ K/V は CPU に配置）と分散スコアリング、オンデマンド転送により、2×A800上で1 億トークンの推論を実現します。
- メモリの織り交ぜ：「生成的検索→コンテキスト拡張→生成」を適応的に交互に実行し、ドキュメント横断的なマルチホップ推論能力を大幅に向上させます。
- 包括的な評価：MSA は長期コンテキスト QA や NIAH ベンチマークにおいて、同 backbone の RAG や最先端の RAG 手法、主要な長期コンテキストモデルを上回り、大規模環境下でより優れた安定性と精度を示しました。
全体設計
アーキテクチャ
MSA は検索と生成を微分可能な閉ループに統合します。ドキュメントの暗黙的状態（K/V/Kᵣ）はチャンク単位の平均プーリングによって圧縮されます。ルーティングプロジェクタはコサイン類似度を用いて関連性を計算し（アテンションヘッドの平均後、トークンごとの最大値を取得）、Top-k ドキュメントを選択します。その後、圧縮された K/VをクエリのローカルなK/Vと連結し、自己回帰的デコードに使用します。ルーティングは上位層でのみ機能し、下位層は独立したドキュメント処理を維持することで階層アライメントを実現します。
- 並列（ドキュメント単位）RoPE：各ドキュメントの位置を 0 から再設定することで、短い訓練と長い推論の間の位置のずれを防ぎ、64k での訓練から 1 億への外挿を可能にします。
- グローバル RoPE（アクティブコンテキスト）：クエリの開始インデックスをk（Top-k 検索ブロック数）だけシフトさせることで、背景→クエリ→生成という因果順序を維持します。
図 2: MSA レイヤー（スパースアテンション＋ドキュメント単位の RoPE）
図 2: メモリスパースアテンションレイヤーと並列/グローバル RoPE
推論パイプライン
MSA は3 段階のパイプライン（図 3）を採用しています。
- グローバルメモリエンコーディング（オフライン）：コーパスに対して順伝播計算を行い、チャンク単位でプーリングされた(K̄, V̄, K̄ᵣ)をキャッシュします。
- オンラインルーティングとコンテキスト構築：クエリをQᵣへ投影してK̄ᵣと照合しTop-kを選択し、選択されたK̄/V̄のみをロードしてローカルコンテキストと連結します。
- スパース生成：スパースなコンテキスト上で自己回帰的生成を行います。
メモリ並列処理によりK̄ᵣを複数 GPU にシャード化します（クエリのブロードキャスト→ローカルスコアリング→グローバルリダクション）。コンテンツのK̄/V̄はホストメモリに保存され、選択された際に非同期で取得されます。これにより1 億トークンのデプロイにおいてGPU メモリとスループットの両立を図っています。
図 3: 3 段階の推論とメモリの織り交ぜ
図 3: オフラインエンコード→オンラインルーティング→スパース生成。マルチホップ推論のためにオプションでマルチターンの織り交ぜも可能
実験結果
実験設定
質問応答：9 つのデータセット（MS MARCO v1、NQ、DuReader、TriviaQA(10M)、NarrativeQA、PopQA、2WikiMultiHopQA、HotpotQA、MuSiQue）、メモリバンク規模277K〜1000 万トークン、評価指標：LLM スコア（0〜5）。
NIAH (RULER)：8 つのサブタスク、32K〜100 万トークン、平均精度を報告。
Backbone モデル：Qwen3-4B-Instruct-2507。同 backbone の RAG や最先端の RAG 手法（KaLMv2＋大規模生成器、オプションで再ランク付けあり）と比較。
表 2: MSA と同 backbone RAG（Qwen3-4B）の比較
要約：平均スコア3.760を達成。標準 RAG 比+16.0%、RAG＋再ランク付け比+11.5%、HippoRAG2 比+14.8%の向上（各方式の最適@k を採用）。同 backbone グループ内では、NarrativeQA を除く全データセットで MSA が首位となりました。
| データセット | トークン数 | Qwen3-4B R@1 | R@5 | R@10 | Qwen3-4B (RR) R@1 | R@5 | R@10 | HippoRAG2 R@1 | R@5 | R@10 | MSA (適応型) |
|---|---|---|---|---|---|---|---|---|---|---|---|
| MS MARCO v1 | 7.34M | 2.893 | 3.011 | 3.005 | 2.934 | 3.032 | 3.017 | 2.676 | 3.005 | 3.019 | 4.141 |
| Natural Questions | 1.47M | 3.452 | 3.374 | 3.297 | 3.494 | 3.408 | 3.385 | 3.338 | 3.389 | 3.374 | 3.545 |
| DuReader | 277K | 3.726 | 3.579 | 3.594 | 3.848 | 3.618 | 3.607 | 2.941 | 3.485 | 3.415 | 4.155 |
| TriviaQA (10M) | 10M | 4.133 | 4.414 | 4.273 | 4.313 | 4.375 | 4.391 | 4.188 | 4.430 | 4.367 | 4.621 |
| NarrativeQA | 538K | 1.611 | 2.567 | 2.860 | 3.638 | 3.492 | 3.536 | 1.959 | 2.628 | 2.655 | 3.395 |
| PopQA | 1.18M | 2.959 | 3.273 | 3.299 | 3.315 | 3.264 | 3.266 | 3.111 | 3.249 | 3.249 | 3.433 |
| 2WikiMultiHopQA | 722K | 1.065 | 3.055 | 3.136 | 1.187 | 3.057 | 3.159 | 1.045 | 3.180 | 3.330 | 4.280 |
| HotpotQA | 1.35M | 2.252 | 3.582 | 3.787 | 2.642 | 3.990 | 4.022 | 3.230 | 3.770 | 3.970 | 4.061 |
| MuSiQue | 1.41M | 0.936 | 1.752 | 1.928 | 1.144 | 1.960 | 1.965 | 1.020 | 1.907 | 2.095 | 2.211 |
| 平均 | — | 2.559 | 3.179 | 3.242 | 2.946 | 3.355 | 3.372 | 2.612 | 3.227 | 3.275 | 3.760 |
表 2: 同 backbone RAG と MSA の比較（@1/@5/@10 vs MSA @適応型）
表 3: MSA と最先端 RAG（大規模 backbone モデル）の比較
要約：KaLMv2+Qwen3-235BおよびKaLMv2+Llama-3.3-70B（再ランク付けの有無を含む）との比較において、MSA は9 つ中 4 つのデータセットで最高スコアを記録しました。平均スコアは3.760であり、各最強構成に対してそれぞれ+7.2%、+5.0%、+10.7%、+5.4%の向上を示しました。一部のデータセット（MuSiQue など）での差は、主にパラメータ数と固有の推論能力の違いに起因するものです。
| データセット | KaLMv2 + Qwen3-235B R@1 | R@5 | R@10 | Qwen3-235B (RR) R@1 | R@5 | R@10 | KaLMv2 + Llama-3.3 R@1 | R@5 | R@10 | Llama-3.3 (RR) R@1 | R@5 | R@10 | MSA (適応型) |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| MS MARCO v1 | 2.846 | 3.028 | 3.027 | 2.886 | 3.020 | 2.995 | 2.649 | 2.904 | 2.919 | 2.881 | 2.955 | 2.952 | 4.141 |
| Natural Questions | 3.711 | 3.670 | 3.694 | 3.621 | 3.610 | 3.645 | 3.675 | 3.674 | 3.662 | 3.756 | 3.665 | 3.647 | 3.545 |
| DuReader | 4.044 | 3.991 | 3.978 | 3.973 | 3.932 | 3.891 | 4.051 | 3.846 | 3.742 | 3.967 | 3.776 | 3.780 | 4.155 |
| TriviaQA (10M) | 4.367 | 4.656 | 4.578 | 4.492 | 4.320 | 4.555 | 4.273 | 4.740 | 4.719 | 4.547 | 4.703 | 4.695 | 4.621 |
| NarrativeQA | 1.413 | 2.130 | 2.427 | 3.212 | 3.427 | 3.375 | 1.290 | 2.123 | 2.382 | 3.150 | 3.263 | 3.317 | 3.395 |
| PopQA | 2.810 | 3.347 | 3.396 | 3.268 | 3.380 | 3.376 | 2.787 | 3.298 | 3.305 | 3.337 | 3.384 | 3.362 | 3.433 |
| 2WikiMultiHopQA | 2.646 | 3.579 | 3.582 | 1.855 | 3.381 | 3.583 | 1.339 | 3.263 | 3.445 | 1.651 | 3.332 | 3.541 | 4.280 |
| HotpotQA | 3.497 | 4.090 | 4.225 | 3.341 | 4.141 | 4.194 | 3.070 | 3.896 | 4.127 | 3.428 | 4.145 | 4.203 | 4.061 |
| MuSiQue | 1.988 | 2.462 | 2.647 | 1.801 | 2.522 | 2.605 | 1.704 | 2.317 | 2.258 | 1.895 | 2.462 | 2.614 | 2.211 |
| 平均 | 3.036 | 3.439 | 3.506 | 3.161 | 3.526 | 3.580 | 2.760 | 3.340 | 3.396 | 3.179 | 3.521 | 3.568 | 3.760 |
表 3: 最先端 RAG 手法（強力な検索器＋大規模生成器＋オプションの再ランク付け）と MSA の比較
図 4: RULER NIAH における安定性（32K→100 万）
要約：MSA は100 万トークンの環境下でも94.84%という高精度を維持しています。一方、改修を施していない backbone モデルは128Kを超えると急激に精度が低下し、100 万時には 24.69%にまで落ち込みます。混合線形アテンションを用いた長期コンテキストモデルも128K/256K 以上で明らかに劣化します。外部メモリ型エージェント（RL-MemoryAgent-14B など）は比較的安定していますが、絶対精度は低く、MSA に比べて減衰幅も大きくなっています。
図 4: コンテキスト長に対する精度の変化（高いほど良い）
実装上の留意点
- 訓練：1589.5 億トークンを用いた継続的プレトレーニングを実施し、補助ルーティング損失を使用。その後、2 段階の SFT（8k→64kのカリキュラム学習）を行いました。
- アブレーション実験（論文 表 4）：カリキュラム拡張、メモリの織り交ぜ、継続的プレトレーニング、および生テキストの注入はいずれも顕著な貢献をしており、これらを除去すると5%〜37%の性能低下を招きます。
クイックスタート
完全な説明（プロジェクト構造、サポートされるベンチマークなど）についてはQUICK_START.mdをご参照ください。
1. インストール
conda create -n msa python=3.12 -y && conda activate msa
pip install -r requirements.txt
pip install flash-attn==2.7.4.post1 --no-build-isolation2. モデルのダウンロード
mkdir ckpt
huggingface-cli download --resume-download EverMind-AI/MSA-4B --local-dir ckpt/MSA-4B3. ベンチマーク用データのダウンロード
ベンチマーク用データは EverMind-AI/MSA-RAG-BENCHMARKS でホストされており、初回実行時に自動的にdata/ディレクトリへダウンロードされます。
4. 実行
# ベンチマークでの推論実行
bash scripts/run_benchmarks.sh eval_benchmark
# LLM スコアの計算
bash scripts/calculate_llm_score.sh eval_benchmark