M+框架來了，增加LLM隱空間記憶，不再受上下文窗口限制

本文的第一作者Yu Wang來自加州大學圣地亞哥分校，主要研究方向為Memory for LLMs and Memory for LLM Agents. 該論文于2024年在MIT-IBM Waston Lab實習過程中完成，通訊作者Zexue He也畢業(yè)于加州大學圣地亞哥分校。

M+是在 MemoryLLM 之上提出的長期隱空間記憶擴展框架：通過把「過期」隱藏向量寫入CPU - 側(cè)長期記憶池，再用協(xié)同檢索器拉回最相關(guān)記憶，它將 8 B 級模型的有效記憶跨度從原本不到 20 k tokens 提升到 160 k tokens 以上，同時顯存占用保持不變。

• 論文標題：M+: Extending MemoryLLM with Scalable Long-Term Memory
• 論文鏈接：https://arxiv.org/abs/2502.00592
• 代碼倉庫：https://github.com/wangyu-ustc/MemoryLLM
• 開源模型：https://huggingface.co/YuWangX/mplus-8b

背景：上下文 ≠ 記憶 & 現(xiàn)有記憶模型的缺陷

上下文窗口并不能直接等價于記憶。GPT-4.1 之類的模型即便支持 100 萬 token，也會隨窗口線性升高顯存與延遲，難以落地本地部署。

業(yè)界主流做法是 “Token-Level Memory”：把歷史內(nèi)容或三元組存在數(shù)據(jù)庫 / 向量庫，檢索后再拼接回 prompt；MemGPT等系統(tǒng)即屬此類。該類做法不需要重復訓練，直接結(jié)合 GPT-4 這樣的大模型便可以獲得很不錯的性能，但是，它也會有一些隨之而來的問題：(1)冗余：原始文本并非最緊湊表達，重復率高。(2)沖突難解：遇到相互矛盾或不斷更新的信息時，文本級沖突消解復雜。(3)多模態(tài)能力弱：由于數(shù)據(jù)庫格式為文本，處理音頻或者圖片，視頻數(shù)據(jù)將相對困難。

因此，我們希望探索隱空間 (Latent-Space) 的 Memory -- 既壓縮又可端到端訓練，更接近人類在神經(jīng)激活中存儲信息的方式。

M + 的關(guān)鍵改進：Long-Term Memory

在 MemoryLLM 中，我們?yōu)?8B 的 Llama3 模型引入了約 1.67B 的 Memory。Llama3-8B 的 Transformer 共包含 32 層。當?shù)谝粚咏邮盏皆~輸入后，會通過 Embedding 層將詞轉(zhuǎn)化為一系列 4096 維的向量。基于這一特點，我們設(shè)計了 MemoryLLM，在每一層都加入 N 個 Memory Tokens（實驗中 N=12800）。在生成過程中，這些 Memory Tokens 會作為每一層的 Prefix，通過 Cross-Attention 將信息注入后續(xù)層，使模型能 “看到” 保存在 Memory Pool 中的歷史信息。

在更新階段，我們會將每層 Memory Token 中最后 K 個（實驗中 K=256）與需要寫入的信息一同送入 Transformer，再次經(jīng)過 Cross-Attention，將信息壓入新的 Memory Tokens 中（如下圖所示）。與此同時，我們在原有 Memory 中隨機丟棄 K 個舊 Token，并將新生成的 K 個 Token 放到 Memory 尾部，完成更新。

基于這樣的設(shè)計，借助每層 12800 個 Memory Vectors，我們在 50k tokens 內(nèi)都能保持良好的信息留存（最早 MemoryLLM-7B 版本只做到 20k，后續(xù)在 GitHub 提供的新版模型 https://github.com/wangyu-ustc/MemoryLLM 可達 50k）。然而，這樣的記憶容量仍無法滿足我們對更長序列的期待。要進一步擴展 Memory，單靠原有的 1.67B 容量已遠遠不夠，因此我們提出了Long-Term Memory。

如何高效實現(xiàn) Long-Term Memory？考慮到 MemoryLLM 中每一個 Memory Token 本質(zhì)上都來自 Hidden States，我們將那些在更新過程中被 “丟棄” 的 Memory Token 并非直接舍棄，而是將其保存在長期記憶池中（如下圖）。

僅僅保存是不夠的，我們還需要具備強大的提取能力。最初我們嘗試用 Attention 來從長期記憶中檢索 Hidden States，但實驗表明 Attention 在提取 Hidden States 時效果有限（在論文的消融實驗中做了詳盡對比）。因此我們提出協(xié)同提取器（Co-trained Retriever），并與全模型進行聯(lián)合訓練（如下圖）。

通過這一結(jié)構(gòu)，我們將模型的有效記憶跨度從 50k 一舉提升到 160k，且由于 Memory 主要駐留在 CPU，不會顯著增加 GPU 負擔。

M + 的實驗結(jié)果

顯著性能提升及更少的 GPU 使用：在 Longbook-QA 和 Longbook-Event-QA 兩個數(shù)據(jù)集上，我們都在更少 GPU 的使用下（單卡 18GB 左右）獲得了更強大的性能。

更強的信息留存能力：在 SQuAD 數(shù)據(jù)集上表現(xiàn)出遠超 MemoryLLM-7B 以及相關(guān) ablation baseline 的信息留存能力，可以達到 160k 依舊不完全遺忘過去的信息。

結(jié)語

M+ 展示了我們在探索隱空間長期記憶領(lǐng)域的重要進展，也為下一代具備持續(xù)記憶能力的語言模型提供了堅實的技術(shù)支撐。未來，我們將繼續(xù)研究更高效的存儲機制、更智能的檢索策略，以及與多模態(tài)輸入更自然融合的隱空間記憶架構(gòu)。在此方向上，M+ 不僅是對 MemoryLLM 的一次擴展，也是我們對 “讓模型擁有接近人類記憶能力” 這一愿景的又一次有力實踐。