Memento：教導 LLM 自行管理上下文 (Context)
AI 語音朗讀 · Edge TTS
Memento：教導 LLM 自行管理上下文 (Context)
我們教導模型在生成過程中自行壓縮其思維鏈 (Chain-of-Thought)。峰值 KV cache 下降了 2–3 倍，吞吐量幾乎翻倍，且被抹除的推理區塊會在 KV cache 中留下痕跡，供模型持續使用。論文、OpenMementos (228K 軌跡)、資料生成與推論程式碼皆已開源。
如果您太忙無法閱讀整篇文章，以下是我們的發現：
您可以教導模型將其思維鏈分割成多個區塊，將每個區塊壓縮成密集的「memento」，並從該處繼續進行推理。對約 30K 個範例進行標準 SFT (監督式微調) 足以讓模型學會此技能。
這能將峰值 KV cache 減少 2–3 倍，並使服務吞吐量幾乎翻倍，且僅有微小的準確度差距，這些差距會隨著規模擴大而縮小，並透過 RL (強化學習) 消除。
被抹除的區塊並未完全消失：其資訊透過 KV cache 的表示形式向前洩漏，形成了一個隱含的第二通道，若沒有此通道，準確度會顯著下降。
我們發布了 OpenMementos (基於 OpenThoughts-v3 建構的 228K 標註軌跡)、資料生成管道，以及具備原生區塊遮罩 (block masking) 功能的 vLLM 修補程式。
此項研究由以下人員完成：Vasilis Kontonis (@vkontonis), Yuchen Zeng (@yzeng58), Shivam Garg (@ShivamGarg91462), Lingjiao Chen (@ChenLingjiao), Hao Tang (@tanghao95), Ziyan Wang (@ZiyanWang98), Ahmed Awadallah (@AhmedHAwadallah), Eric Horvitz (@erichorvitz), John Langford (@JohnCLangford), Dimitris Papailiopoulos
這是一次與優秀合作夥伴團隊共同完成的驚人合作。Vasilis、Yuchen、Shivam 和 Lingjiao 在研究與工程方面承擔了大部分繁重的工作，並做出了我獨自一人無法企及的貢獻。感謝 Hao 和 Ziyan 在專案早期的努力。感謝 Eric 在整個過程中的指導，John 深入鑽研 vLLM 的工作，以及 Ahmed 提供的支援、資源與追求此研究的自由。
我的角色主要是提出意見，且大多時候是錯的 😊
這是 Microsoft Research (@MSFTResearch) 內部精品研究實驗室 AI Frontiers 的研究成果。
問題：LLM 不知道如何管理其上下文
目前已證實推理模型可以透過生成大量 token 來解決難題。測試時計算 (Test-time compute) 非常有效，並在競賽級數學與程式撰寫方面帶來了顯著進展，但它也可能導致單次推論呼叫產生數十萬個 token。這大約是一本書的長度。所有這些 token 都會留在記憶體中，無論它們是否有用，都以相同的成本被關注。模型沒有內建機制來壓縮已推導出的內容、保留結論並繼續前進。
目前有外部管理此問題的方法，例如執行獨立的摘要器、以壓縮後的上下文重新啟動 API 呼叫、或圍繞模型建立編排邏輯。然而，這些都是附加在模型周圍的系統，而非模型本身學到的技能。我們認為，決定記住什麼以及忘記什麼，可以也應該是模型在訓練過程中學會的一項技能。
Memento 正是教導語言模型這項技能。經過 Memento 訓練（即 mementified）的模型會將其推理過程分割成語意連貫的區塊。當一個區塊完成時，模型會產生一個 memento：這是對該區塊結論、關鍵中間值、公式和策略決策的簡潔、資訊密集型壓縮。將 memento 視為一個引理 (lemma)：它是未來推理步驟需要繼續進行的最小記錄。
一旦產生 memento，先前的思考區塊就會從注意力機制中被遮罩，其 KV cache 條目也會被清除。從那一點開始，模型只會看到過去的 memento 以及它目前正在處理的區塊。這意味著上下文會在模型處理區塊時增加，但在產生 memento 並移除區塊後會急劇下降。這產生了一種鋸齒狀模式，使得峰值記憶體需求僅為標準平坦 CoT 軌跡所需的一小部分。
以下是實際的 memento 生成過程：
重要的是，這一切都發生在單次生成呼叫中，無需重新啟動、獨立的摘要器或編排層。模型自行分割、壓縮並遮罩其推理過程。
我們將 Memento 應用於五個模型：Qwen2.5-7B、Qwen3 8B 和 32B、Phi-4-Reasoning (14B) 以及 Olmo3-7B-Think。它在所有模型上都有效。峰值 KV cache 下降了 2-3 倍，且準確度差距很小，這些差距會隨著規模擴大而縮小，並透過 RL 進一步消除。
我們還發現了一些未預料到的情況：被抹除的區塊雖然在物理上從 KV cache 中移除，但並沒有完全從模型的表示形式中消失。稍後會詳細說明！
如何教導上下文管理？將其加入訓練資料中！
教導這種行為需要一種不太常見的訓練資料：大規模、高品質的推理軌跡，這些軌跡被分割成多個區塊，每個區塊都配有一個 memento，該 memento 以模型可以從中繼續推理的方式捕捉區塊的結論。直覺很簡單：如果您獲取推理軌跡、進行分割、加入適當的摘要，並對結果進行 SFT，也許模型就能學會自行管理上下文。
聽起來很簡單，但與許多事情一樣，過程中有多個組件損壞並需要修復。
首先，我們決定基於 OpenThoughts 進行建構：由 QwQ-32B 生成的推理軌跡，這些軌跡已經具有相當高的品質且被社群廣泛使用，這讓我們省去了從零開始生成所有內容的麻煩。現在的問題是：我們如何從原始軌跡轉變為分割後的、帶有每個區塊邊界 memento 的註解軌跡？挑戰在於推理軌跡沒有自然的區塊邊界，即想法相互混合、計算跨越多個句子，而 CoT 的「切割」位置更多取決於語意，而非格式或其他明顯指標。
我們首先嘗試了顯而易見的方法：將軌跡貼到前沿模型中，要求它直接進行分割和摘要。這行不通！即使您先將軌跡切成碎片也不行，因為您並不完全知道在哪裡切割。尋找良好的分區需要同時考慮區塊連貫性、大小平衡和語意邊界，這是一種複雜的組合優化問題，LLM（至少我們嘗試過的那些）很難一次完成。
因此，我們將問題拆解。首先，我們將每個軌跡分割成原子單元（句子、程式碼區塊、數學方程式），這些單元無法再有意義地拆分。然後，LLM 對每個句子間的邊界進行評分，從 0（思維中斷，會破壞流程，即不好）到 3（重大轉換，自然的停止點，即好）。這是一個局部問題，而 LLM 非常擅長處理局部問題。給定這些分數後，實際邊界位置的全局優化則由動態規劃處理，它在最大化邊界品質的同時，懲罰不均勻的區塊大小。這正是（根據我們的經驗）LLM 很難進行 zero-shot 的地方，但卻是傳統動態規劃能完美解決的地方。
一旦有了分割後的軌跡，我們現在需要壓縮每個區塊。壓縮器 LLM 為每個區塊產生一個 memento，我們在 Prompt 中明確解釋，任務不是普通的摘要，而是狀態壓縮：產生足夠緊湊的內容，使模型僅憑 memento 就能繼續推理，而無需看到原始區塊。於是，memento 就誕生了！
然後，一個獨立的評判 LLM 會根據各種維度（提取的公式、保留的值、命名的方法、包含的驗證、無幻覺、結果優先結構）評估每個 memento，如果分數未達標，評判者會提供具體、可執行的回饋（不是「需要更多細節」，而是「缺少公式：K² − 3K + 3」等），壓縮器隨後重試。
這種迭代優化被證明至關重要。單次壓縮在我們的評分標準下幾乎達不到 28% 的通過率，因為初始 memento 通常會遺漏後續區塊依賴的精確公式或中間值。兩輪評判回饋將通過率提高到了 92%。
注意：對於管道中的所有 LLM 呼叫，我們都使用了 GPT-5.x，但任何能力足夠的模型都應該有效。完整的管道已開源，我們希望人們使用它、改進它，並建立比我們更好的資料集。
資料生成管道如下所示：
最終資料集 OpenMementos 包含 228K 個標註軌跡，其中 54% 為數學、19% 為程式碼、27% 為科學問題。我們測量得出，memento 實現了約 6 倍的軌跡級壓縮：每個軌跡約 11k 個推理 token 被壓縮至 2k 個 token 以下的 memento。
以下是 OpenMementos 的一些酷炫壓縮統計數據。
訓練：我們如何對模型施加壓力？
我們已經標註了具有區塊結構和 memento 的軌跡。顯而易見的下一步是：對它們進行 SFT；但具體該怎麼做？目標是讓模型最終僅從 memento 進行推理，並將原始區塊遮罩，移除其 KV cache 條目。
一種選擇是使用標準因果注意力 (causal attention) 對標註軌跡進行訓練，然後在推論時遮罩區塊。這在一定程度上有效，但這意味著訓練過程從未對模型施加任何壓力，要求它真正將資訊打包到 memento 中，因為它在訓練期間總是可以退回到關注完整區塊，而到了推論時，它基本上是孤軍奮戰。
我們希望訓練與推論相匹配：如果區塊在模型部署時會被遮罩，那麼它們在訓練時也應該被遮罩。
但從一開始就直接使用區塊遮罩進行訓練效果也不好。模型試圖同時學習三件事：區塊-memento 格式、如何在嚴格限制下進行壓縮、如何僅依賴 memento 進行下一個區塊的生成。它很難同時處理這三者。
我們發現課程學習 (curriculum) 非常重要。
階段 1 使用標準因果注意力，並對所有 token 進行損失計算。模型學習格式：何時結束區塊、如何撰寫 memento、結構看起來是什麼樣子。它仍然可以看到所有內容，因此還沒有壓縮壓力。階段 2 隨後引入硬性限制：在每個 memento 之後，先前的思考區塊會被完全遮罩，不參與後續的注意力計算。現在模型必須產生包含未來推理所需一切內容的 memento，因為原始區塊已經消失。這似乎是真正學習發生的階段：模型被迫將更多資訊打包到 memento 中，就像一種 RL 風格的壓力訊號，推動其進行自包含的壓縮。非常酷！
我們還發現，與其他關於技能教學的研究一致，在 OpenMementos 資料的一小部分上進行訓練就足夠了。即使是從 228K 池中抽取約 30K 個樣本，每個階段訓練 5 個 epoch，序列長度為 32K，也足以讓模型掌握這項技能。
對於已經具備良好推理能力的模型 (Qwen3、Olmo3、Phi-4-reasoning)，兩個階段就足夠了；非推理基礎模型（如 Qwen 2.5 7B）則需要先進行一輪標準推理 SFT。Memento 不需要比標準推理 SFT 更多的資料，它只是需要不同類型的資料。
壓縮如何影響準確度？
Memento 的第一個明顯擔憂是，關注更少的 token 應該會損害準確度。當我們第一次查看數據時，確實出現了下降。這來自哪裡？
我們最初的反應是，這一定是壓縮和稀疏性造成的：模型看到的上下文少得多，所以它當然會變差。但隨後我們進行了對照研究，情況變得更有趣。
關鍵的洞察是，我們在 QwQ-32B 生成的 OpenThoughts 軌跡上進行訓練，這與我們正在微調的模型不同。我們的幾個目標模型是在 QwQ 之後發布的，有時更強大。因此，我們進行了對照：獲取每個基礎模型，在相同的原始 OpenThoughts 軌跡（沒有區塊結構，沒有 memento）上進行 SFT，並僅測量由此產生的準確度下降。事實證明，僅僅對其他模型的推理軌跡進行 SFT 就已經會付出代價。當我們將 Memento 與該對照組而非未經修改的基準進行比較時，壓縮帶來的額外下降很小，在某些情況下甚至可以忽略不計。
但我們仍然對剩餘的任何準確度差距感到好奇。所以我們問：模型還能解決相同的問題嗎？
為了測試這一點，我們在 AIME 2024/25/26 上為所有三個模型家族的每個問題生成了 64 個完成結果，答案是絕大多數情況下可以。基礎模型與 Memento 解決問題的重疊率平均為 96.4%，在某些設定下達到 100%。模型保留了解決這些問題的能力，下降的是在任何單次嘗試中解決它們的一致性。
這是一個重要的區別，因為它很可能意味著差距是可以縮小的。例如，我們發現即使是 k=3 的多數投票也足以讓 Memento 模型不僅匹配對照組，還能匹配原始基準。這證實了該能力仍然存在於分佈中。
自然的下一步是 RL。不出所料，它有效：使用 CISPO 微調 Qwen3-8B Memento 檢查點可以恢復 AIME'26 和 GPQA-Diamond 的分數（有時甚至超過原始基準），同時在 RL 之後 KV 的節省仍然相當可觀。
即使沒有 RL，規模擴大也有獨立的幫助。從 Qwen3-8B 到 32B，儘管兩個模型都是在相同的 QwQ-32B 軌跡上訓練的，但差距顯著縮小：較大的模型能更優雅地處理分佈不匹配和壓縮的影響。
因此，結論是：壓縮保留了能力，任何一致性的損失主要歸因於訓練資料的不匹配，而非根本性的限制，且 RL 和規模擴大都能進一步縮小差距。
這是我們 mementified Qwen3-8B 的一條不錯的 RL 曲線：
雙重資訊流
在專案早期，關於推論應該如何運作有很多討論。最簡單的方法，也是讓我們生活更輕鬆的方法，是重新啟動：每次產生 memento 時，終止 KV cache 並以僅包含累積的 memento 文字開始新的 API 呼叫。無需在 vLLM 內部實作非因果稀疏注意力，這被證明是非常痛苦的。只需重新啟動呼叫即可。
但我們一直回到一個擔憂：當產生 memento 時，模型仍然可以看到完整的思考區塊，因此 memento token 在其自身生成過程中會關注區塊 token。區塊僅在 memento 完成後才被遮罩。這意味著 memento 的 KV cache 條目是作為區塊內容的函數計算出來的。因此，即使在區塊文字消失後，來自區塊的某些內容仍會在 memento 的 KV 表示形式中存活。如果下一個區塊關注 memento，它就是以間接方式關注這種先前內容的隱含、軟性表示。在重新啟動設定中，您會丟棄所有這些內容。
也就是說，關於被遮罩區塊的非平凡資訊在 KV cache 表示形式中存活了下來，超出了實際 memento token 所捕捉到的範圍。因此，一個問題一直困擾著我們：這種隱含的資訊通道對準確度真的重要嗎？還是我們想太多了？
於是我們進行了消融實驗。採用相同的 Qwen3-8B 檢查點，比較正常的 Memento 推論（遮罩區塊但保持 memento KV 狀態完整）與重新啟動模式（在每個 memento 邊界從零開始重新計算整個 KV cache，因此 memento 本身從未關注過它們的區塊）。
重新啟動模式使 AIME'24 的分數從 66.1% 下降到 50.8%。15 個百分點，對於任何理性的觀察者來說，這都不會被視為雜訊。事實上，這強烈表明流經 KV 表示形式的側通道資訊非常重要。為了確保這一點，我們希望進一步測試這個假設。
因此，我們設計了一個簡單的實驗：獲取一個模型，將一個隨機的 5 位數密碼注入目標區塊，遮罩該區塊，並對從未直接關注被遮罩區塊的下游 memento 的 KV 狀態訓練線性探針。您能恢復僅存在於隱含 KV 通道中，而不存在於任何 memento 文字中的資訊嗎？
哦，是的，您可以！探針可以很好地重建密碼，遠高於隨機機率，這正是因為 KV 狀態中發生的資訊洩漏。這種洩漏集中在較深的層中，隨著與目標區塊距離的增加而衰減，但即使在七個區塊之外仍然可以檢測到，並且隨著模型容量的增加而擴大。
我們還在一個小型受控玩具 Transformer（4 層，810K 參數）上驗證了這一點，其中洩漏在訓練檢查點中是恆定的，即使任務準確度從 77% 提高到 95%。這是殘差連接、因果注意力和就地遮罩 (in-place masking) 的架構後果。
我們認為這將 Memento 與 InftyThink 和 Accordion-Thinking 等方法區分開來，後者丟棄原始 token 並僅從摘要文字重建上下文，從而完全丟失了這個隱含通道。
上述發現使我們確信，在 vLLM 內部實作正確的區塊遮罩是值得的基礎設施工作，而不是採取基礎設施上更容易的重新啟動路徑。
讓 Memento 在 vLLM 中運作
Memento 的區塊遮罩是資料依賴的，並且在生成過程中不斷變化，因為遮罩哪些 token 取決於模型產生什麼。遺憾的是，沒有生產級推論框架開箱即用地支援這一點。我們從 HuggingFace 後端開始，這足以驗證區塊遮罩並將所有內容保留在單次推論呼叫中確實有幫助，但一旦我們確信，顯然我們需要在 vLLM 內部正確地建構它。
這被證明是痛苦的，但最終是可行的。關鍵的設計選擇是物理 KV cache 壓縮，而不是邏輯遮罩：當一個區塊完成時，其 KV 條目會被物理移除，釋放的插槽會歸還給 KV 池。這意味著標準的 FlashAttention 和 paged-attention 核心完全無需修改即可運作，因為它們永遠看不到被移除的 token。該實作純粹在 vLLM Python 層級運作，可以作為現有 vLLM 安裝的修補程式進行安裝。
在單個 B200 GPU 上，並發請求數為 240（Qwen3-8B，32K 最大 token），Memento 維持 4,290 tok/s，而標準版為 2,447（1.75 倍吞吐量），並在 693 秒內完成批次處理，而標準版為 1,096 秒。這些收益來自於在區塊完成時移除 KV 條目，使引擎在標準 vLLM 受到 KV-cache 限制的情況下能夠維持更大的批次大小。
這種基礎設施也被證明對 RL 至關重要：生成 32K-token 的訓練 rollout 需要在生成過程中進行區塊遮罩，每個 rollout 都在運行中產生並壓縮區塊。
如果沒有 vLLM 分支，在此規模下進行 RL 是不可行的。
下一步是什麼？
從這裡開始，兩件事似乎很自然。首先，擴大 RL 配方：我們使用 Qwen3-8B 的結果還很初步，pass@64 分析清楚地表明還有很大的改進空間。具有更多 RL 計算能力的更大模型應該能帶我們去到有趣的地方。
其次，對我們來說更重要的是：Agent。Memento 是作為測試案例為數學、程式撰寫和科學推理而建構的，並不是因為我們認為單輪數學和程式撰寫是最有趣的應用。區塊與壓縮模式映射到任何模型累積長軌跡中間狀態且有限上下文視窗成為瓶頸的場景。終端和 CLI Agent 自然是多輪的，其中每個動作-觀察循環都被佈置為一個自然的區塊，而選擇性地記住和遺忘的能力正是似乎缺失的（至少在 OSS 模型/Agent 中是如此）。最近關於 Agent 場景中上下文壓縮的研究（例如來自 Anthropic 和 OpenAI 的研究）指向了相同的方向，我們認為這裡有很大的探索空間。
結語
Memento 最初是為了教導模型壓縮其自身推理過程的嘗試。這確實有效：2-3 倍的 KV 減少，在準確度大致保留的同時，吞吐量幾乎翻倍。但我們從這個專案中獲得了兩個感覺比效率提升更重要的洞察。
首先是上下文管理可以透過對正確資料進行標準訓練來教導。一個對區塊或摘要沒有概念的模型，在對約 30K 個範例進行 SFT 後，可以學會分割其自身的推理、壓縮每個區塊，並從壓縮版本繼續。這是一項非平凡的、非因果的技能，涉及稀疏注意力、選擇性遺忘、狀態壓縮，這些都是透過完全傳統的訓練獲得的。我們認為，實際上還有更廣闊的非常規能力空間可以透過這種方式教導。
其次是由硬 token 及其 KV 表示形式支援的雙重資訊流。當您在單次前向傳遞中遮罩一個區塊時，區塊的資訊並沒有完全消失：它存在於區塊仍然可見時計算出的 memento 的 KV 表示形式中。這既是有用的，也是架構上不可避免的，我們還不知道這個隱含通道可以推到多遠，特別是在 RL 的情況下。
這兩部分指向相同的方向：記憶體管理應該是一項習得的能力，而模型學習所需的努力比我們預期的要少。
我們認為 Memento 只是第一步，還有很長的路要走，包括更好的訓練資料、更強大的 RL 和 Agent 應用。我們正在所有這些方面繼續努力，同時我們正在發布 OpenMementos (228K 標註推理軌跡)、我們的完整資料生成管道，以及具有原生區塊遮罩功能的 vLLM 分支。
在此期間，請停止清除您的 KV cache。您的模型記得比您想像的要多。
論文 · OpenMementos · 資料管道程式碼 · 具備區塊遮罩功能的 vLLM 修補程式