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

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

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

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

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

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

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

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

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

在更新階段，我們會(huì)將每層 Memory Token 中最后 K 個(gè)（實(shí)驗(yàn)中 K=256）與需要寫(xiě)入的信息一同送入 Transformer，再次經(jīng)過(guò) Cross-Attention，將信息壓入新的 Memory Tokens 中（如下圖所示）。與此同時(shí)，我們?cè)谠?Memory 中隨機(jī)丟棄 K 個(gè)舊 Token，并將新生成的 K 個(gè) Token 放到 Memory 尾部，完成更新。

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

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

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

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

M + 的實(shí)驗(yàn)結(jié)果

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

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

結(jié)語(yǔ)

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