MSA 程式碼準時開源!突破 1 億 Token 記憶極限,效能超越 58 倍大模型 RAG

幾週前,@EverMind 發表了關於「記憶稀疏注意力」(Memory Sparse Attention, MSA)的論文,主要建構了一套系統,能讓模型記住多達 1 億個詞元(相當於一千本書的長度),並且效能損失不到 9%,依然能夠找到正確答案。在多項基準測試中,這個擁有 40 億參數的模型甚至超越了基於其 58 倍規模模型所建構的 RAG 系統。

其核心思路在於:與其在單獨的資料庫中搜尋並寄望於找到正確資訊(RAG 演算法的運作方式),還不如直接將記憶功能建構到模型的思維模式中。它能夠端到端地學習哪些資訊需要記住、哪些需要忽略,無需單獨的檢索流程。

論文發表後引起業界極大迴響,各界殷切期盼其開原始碼以供研究、共同精進。昨天,EverMind 兌現承諾,完整開源整個專案程式碼,不到一天時間就收穫 2.5K star,可見大家對這一技術的關注度。

專案簡介

長期記憶是通用智慧的基礎,但全注意力的計算瓶頸將大多數大型語言模型的有效上下文長度限制在 128K–1M。現有方案——混合線性注意力、固定大小的狀態記憶(如 RNN)、以及RAG/智慧體等外部儲存——要麼在極端長度下精度快速下降、延遲持續增長,要麼缺乏端到端可微性或動態記憶管理能力,要麼需要複雜的工程流水線。我們提出了記憶稀疏注意力(MSA):一個端到端可訓練、可擴展的稀疏隱式狀態記憶框架。核心思路包括:

  • 可擴展的稀疏注意力 + 文件級 RoPE(並行/全域),在訓練和推論中均實現近線性複雜度
  • KV 快取壓縮搭配記憶並行推論引擎,在 2×A800 GPU 上實現1 億 token吞吐;
  • 記憶交織機制,支援跨分散記憶片段的多輪、多跳推論。

在長上下文問答和 NIAH(大海撈針)基準測試中,MSA超越了同骨幹 RAG、頂尖 RAG 方案和領先的長上下文模型。在前所未有的16K→1 億 token範圍內,MSA 僅有不到 9%的效能下降,為將記憶容量與推論能力解耦提供了一條可行路徑。

從 16K 擴展到 1 億 token:MSA 將 top-k 選擇與稀疏注意力融合,在推論時允許文件解耦的同時保持端到端可微。在 MS MARCO 上,MSA 效能下降不到 9%,並展現出強大的外推能力。部分基線曲線因上下文長度限制而提前結束。

圖 1:16K→1 億 token 擴展曲線

圖 1:MSA 在超長上下文下的可擴展性


核心貢獻

  • 記憶稀疏注意力(MSA):一個端到端可訓練可擴展稀疏注意力層,結合文件級 RoPE,實現 O(L) 複雜度,從16K→1 億 token效能下降不到 9%
  • KV 快取壓縮 + 記憶並行:分層儲存(路由鍵駐留 GPU、內容 K/V 存放 CPU),分散式打分,按需傳輸,在2×A800上實現1 億 token推論。
  • 記憶交織:自適應交替執行「生成式檢索 → 上下文擴展 → 生成」,大幅提升跨文件多跳推論能力。
  • 全面評測:MSA 在長上下文問答和 NIAH 基準上超越同骨幹 RAG、頂尖 RAG 方案和領先長上下文模型,在大規模下展現出更優的穩定性準確率

整體設計

架構

MSA 將檢索與生成融合到一個可微的閉環中。文件隱式狀態(K/V/Kᵣ)通過分塊均值池化進行壓縮。路由投影器利用餘弦相似度計算相關性(先對注意力頭取均值,再對 token 取最大值),選出Top‑k 文件,然後將其壓縮後的 K/V與查詢的本地K/V拼接,用於自回歸解碼。路由僅作用於上層;下層保持獨立文件處理,實現層級對齊。

  • 並行(文件級)RoPE:每個文件的位置從 0 開始重置,避免短訓練、長推論之間的位置漂移,使 64k 訓練能外推到 1 億。
  • 全域 RoPE(活躍上下文):查詢的起始索引偏移 k(Top‑k 檢索塊數),保持因果順序:背景 → 查詢 → 生成

圖 2:MSA 層

圖 2:MSA 層(稀疏注意力 + 文件級 RoPE)

圖 2:記憶稀疏注意力層與並行/全域 RoPE


推論流水線

MSA 採用三階段流水線(圖 3):

  1. 全域記憶編碼(離線):對語料進行前向計算,快取分塊池化後的 (K̄, V̄, K̄ᵣ)。
  2. 線上路由與上下文組裝:將查詢投影為 Qᵣ,與 K̄ᵣ 匹配選出Top‑k,然後僅載入選中的 K̄/V̄ 與本地上下文拼接。
  3. 稀疏生成:在稀疏上下文上進行自回歸生成。

記憶並行將 K̄ᵣ 分片到多個 GPU(廣播查詢 → 本地打分 → 全域歸約)。內容 K̄/V̄ 存放在主機記憶體中,被選中時非同步拉取——在1 億 token部署下兼顧顯存吞吐

圖 3:三階段推論與記憶交織

圖 3:推論流程圖 3:離線編碼 → 線上路由 → 稀疏生成;可選多輪交織用於多跳推論


實驗結果

實驗設定問答:9 個資料集(MS MARCO v1、NQ、DuReader、TriviaQA(10M)、NarrativeQA、PopQA、2WikiMultiHopQA、HotpotQA、MuSiQue),記憶庫規模277K→1000 萬 token,評測指標:LLM 評分(0–5)NIAH (RULER):8 個子任務,32K→100 萬 token,報告平均準確率。骨幹模型:Qwen3‑4B‑Instruct‑2507。對比同骨幹 RAG 和頂尖 RAG 方案(KaLMv2 + 大模型生成器,可選重排序)。

表 2:MSA 對比同骨幹 RAG(Qwen3‑4B)

總結:平均得分3.760,相比標準 RAG 提升+16.0%,相比 RAG+重排序提升+11.5%,相比 HippoRAG2 提升+14.8%(均取各方案最佳 @k);在同骨幹組中,MSA 在除 NarrativeQA 外的所有資料集上領先。

資料集Token 數Qwen3-4B R@1R@5R@10Qwen3-4B (RR) R@1R@5R@10HippoRAG2 R@1R@5R@10MSA (自適應)
MS MARCO v17.34M2.8933.0113.0052.9343.0323.0172.6763.0053.0194.141
Natural Questions1.47M3.4523.3743.2973.4943.4083.3853.3383.3893.3743.545
DuReader277K3.7263.5793.5943.8483.6183.6072.9413.4853.4154.155
TriviaQA (10M)10M4.1334.4144.2734.3134.3754.3914.1884.4304.3674.621
NarrativeQA538K1.6112.5672.8603.6383.4923.5361.9592.6282.6553.395
PopQA1.18M2.9593.2733.2993.3153.2643.2663.1113.2493.2493.433
2WikiMultiHopQA722K1.0653.0553.1361.1873.0573.1591.0453.1803.3304.280
HotpotQA1.35M2.2523.5823.7872.6423.9904.0223.2303.7703.9704.061
MuSiQue1.41M0.9361.7521.9281.1441.9601.9651.0201.9072.0952.211
平均2.5593.1793.2422.9463.3553.3722.6123.2273.2753.760

表 2:同骨幹 RAG 對比 MSA(@1/@5/@10 vs MSA @自適應)


表 3:MSA 對比頂尖 RAG(大骨幹模型)

總結:對比KaLMv2+Qwen3‑235BKaLMv2+Llama‑3.3‑70B(有/無重排序),MSA 在4/9個資料集上取得最高分,平均得分3.760,相對於各最強配置分別提升+7.2%+5.0%+10.7%+5.4%。在少數資料集(如 MuSiQue)上的差距主要源於參數量和固有推論能力的差異。

資料集KaLMv2 + Qwen3‑235B R@1R@5R@10Qwen3‑235B (RR) R@1R@5R@10KaLMv2 + Llama‑3.3 R@1R@5R@10Llama‑3.3 (RR) R@1R@5R@10MSA (自適應)
MS MARCO v12.8463.0283.0272.8863.0202.9952.6492.9042.9192.8812.9552.9524.141
Natural Questions3.7113.6703.6943.6213.6103.6453.6753.6743.6623.7563.6653.6473.545
DuReader4.0443.9913.9783.9733.9323.8914.0513.8463.7423.9673.7763.7804.155
TriviaQA (10M)4.3674.6564.5784.4924.3204.5554.2734.7404.7194.5474.7034.6954.621
NarrativeQA1.4132.1302.4273.2123.4273.3751.2902.1232.3823.1503.2633.3173.395
PopQA2.8103.3473.3963.2683.3803.3762.7873.2983.3053.3373.3843.3623.433
2WikiMultiHopQA2.6463.5793.5821.8553.3813.5831.3393.2633.4451.6513.3323.5414.280
HotpotQA3.4974.0904.2253.3414.1414.1943.0703.8964.1273.4284.1454.2034.061
MuSiQue1.9882.4622.6471.8012.5222.6051.7042.3172.2581.8952.4622.6142.211
平均3.0363.4393.5063.1613.5263.5802.7603.3403.3963.1793.5213.5683.760

表 3:頂尖 RAG 方案(強檢索器 + 大生成器 + 可選重排序)對比 MSA


圖 4:RULER NIAH 穩定性(32K→100 萬)

總結:MSA 在100 萬 token下仍保持94.84%的準確率。未經改動的骨幹模型在超過128K後急劇下降(100 萬時僅 24.69%)。混合線性注意力長上下文模型在≥128K/256K時明顯退化。外部記憶智慧體(如 RL‑MemoryAgent‑14B)雖然較為穩定,但絕對準確率更低,衰減幅度也大於 MSA。

圖 4:RULER NIAH 32K→100 萬圖 4:準確率隨上下文長度的變化(越高越好)


實作說明

  • 訓練:1589.5 億 token 持續預訓練,使用輔助路由損失,隨後進行兩階段 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-isolation

2. 下載模型

mkdir ckpt
huggingface-cli download --resume-download EverMind-AI/MSA-4B --local-dir ckpt/MSA-4B

3. 下載基準資料

基準資料託管在 EverMind-AI/MSA-RAG-BENCHMARKS,首次運行時會自動下載到data/目錄。

4. 運行

# 在基準測試上運行推論
bash scripts/run_benchmarks.sh eval_benchmark
# 計算 LLM 評分
bash scripts/calculate_llm_score.sh eval_benchmark

專案位址

https://github.com/EverMind-AI/MSA

相關文章推薦

分享網址
AINews·AI 新聞聚合平台
© 2026 AINews. All rights reserved.