← 返回首頁

Latent Briefing:透過 KV Cache Compaction 實現多 Agent 系統的高效記憶體共享

Ramp Labs
Ramp Labs
@RampLabs
371🔁 30
𝕏 (Twitter)🔥🔥🔥🔥

AI 語音朗讀 · Edge TTS

Latent Briefing:透過 KV Cache Compaction 實現多 Agent 系統的高效記憶體共享

多 Agent 系統在協調、複雜推理與平行工作流方面展現了極大的潛力。然而,它們往往存在嚴重的 token 使用效率低下的問題。在階層式架構中,當 Orchestrator(協調者)分解任務並委派給 Worker Agent(執行者)時,可能會出現冗餘的中間推理過程。隨著 Orchestrator 的推理軌跡在多次呼叫中不斷擴展,token 的使用量會迅速累積。雖然這些方法可以提升效能,但卻付出了巨大的代價,且 Agent 之間的上下文共享效率往往不佳。

現有的上下文管理方法,例如基於 LLM 的摘要(速度慢)或透過 RAG 進行檢索(不穩定),都有各自的權衡取捨。相反地,我們利用模型的注意力模式來識別上下文中哪些部分是重要的,並在表徵層面捨棄其餘部分。這產生了一種透過直接操作模型 KV cache 來在 Agent 之間共享相關記憶體的方法。我們將此方法稱為 Latent Briefing。

在 LongBench v2 基準測試的 126 個問題上(涵蓋 0–100k token 的文件),我們的方法達成了以下成果:

  • 在各種難度和上下文長度條件下,準確度與基準線相當或有所提升
  • 在中等長度(32k–100k token)文件上,中位數 token 節省高達 49%
  • Worker 模型 token 消耗降低 65%
  • 約 1.7 秒的中位數壓縮開銷,並隨輸入長度線性擴展

遞迴 Agent 中的 Token 爆炸問題

我們採用遞迴語言模型(Recursive Language Model, RLM)框架(Zhang et al., 2025)作為我們多 Agent 系統的基礎架構。在 RLM 中,一個強大的 Orchestrator 會分解任務,並透過 REPL 環境對 Worker 模型進行重複呼叫。Orchestrator 會向 Worker 發送目標明確的查詢,要求其分析文件的特定面向、驗證假設或提取資訊。

雖然 RLM 在處理長上下文方面表現出色,但它們比傳統 LLM 的效率更低,且消耗的 token 顯著更多。此外,Worker 只能看到 Orchestrator 明確傳遞給它的內容:通常是一個目標查詢和原始文件。但 Orchestrator 在多次呼叫中已經建立了豐富的推理軌跡:測試過的假設、識別出的段落、排除的死胡同、發現的交叉引用。這些累積的上下文可以幫助 Worker 更有效地回答問題,但將所有內容作為文字傳遞會導致每次後續呼叫的輸入成本膨脹。最終,Worker 只能在狹窄的問題視角下工作,而 Orchestrator 更廣泛的理解卻未被利用。

標準解決方案都有顯著的缺點:

  • LLM 摘要:每個步驟延遲 20–60 秒,有損資訊,且摘要可能無法捕捉子任務所需的內容
  • RAG / 檢索:需要進行分塊(chunking)與嵌入(embedding),會遺漏跨塊的依賴關係
  • 全部傳遞:昂貴、緩慢,且準確度可能會因無關的上下文而下降

我們需要快速且精確的跨 Agent 記憶體機制,以嘗試減少這種 token 爆炸。


任務導向的 KV Cache Compaction

背景:AM Compaction 框架

我們的方法建立在用於 KV cache 壓縮的注意力匹配(Attention Matching, AM)框架(Zweiger et al., 2026)之上。其核心思想是:給定一個大小為 S 的 KV cache,找到一個大小為 t < S 的壓縮快取,使其產生幾乎相同的注意力輸出。

形式上,對於每個注意力頭(attention head),我們尋求壓縮組件 (C1, β, C2),使得:

其中:

  • C1 (壓縮後的 keys):從原始 key 向量中選出的子集,用於高注意力權重
  • β (偏差修正):標量調整,用於補償缺失的 keys,確保保留 keys 的 softmax 分佈近似於所有 keys 的原始分佈
  • C2 (壓縮後的 values):透過嶺回歸(ridge regression)求解出的重構值向量

原始的 AM 演算法會獨立處理每個(層, 頭)對。對於 Qwen3-14B 而言,這意味著 40 層 × 8 個 KV 頭 = 320 次序列化求解,每次運行三個步驟:

  1. Token 選擇:計算所有查詢與所有 key 位置之間的注意力分數,然後選擇總分最高的前 t 個位置。

  2. 透過 NNLS 求解 Beta:找到偏差修正 β,使得 softmax(q · C1ᵀ + β) 近似於保留 token 的 softmax(q · Kᵀ):透過帶有非負約束的投影梯度下降法求解¹。

  3. 透過嶺回歸求解 C2:求解 C2 = (XᵀX + λI)⁻¹XᵀY,其中 X 是壓縮後的 softmax 矩陣,Y 是原始注意力輸出,重構出保留注意力計算的值向量。

我們的修改:

為了適應推論環境,我們對 AM 壓縮進行了三項關鍵修改:

  1. 任務導向的查詢向量。在原始 AM 框架中,用於評分的查詢是從上下文本身採樣的。我們將其替換為從 Orchestrator 針對此特定 Worker 呼叫的任務提示(task prompt)中導出的查詢。這使得快取壓縮能夠優先處理與 Worker 任務最相關的資訊。

此處的軌跡是 Orchestrator 到目前為止的完整 context window:先前的 Worker 回應、任何 REPL 輸出以及思維鏈(chain of thought)推理。

我們將軌跡和任務提示透過 Worker Agent 進行前向傳遞。任務提示與軌跡 keys 之間的注意力分數告訴我們,Worker 認為軌跡的哪些部分與其當前任務相關。

K = 軌跡 KV cache keys 
Q = 來自 Orchestrator 針對此 Worker 呼叫的任務提示之注意力查詢

對於每個 (層, 頭):
    attn_{l,h} = softmax(Q · Kᵀ / √d)               # 任務與軌跡之間的注意力
    scores_{l,h}(pos) = RMS_q(attn_{l,h}(:, pos))    # 每個位置在任務查詢上的 RMS
  1. 透過全域評分進行共享 token 選擇。我們不再讓每個頭獨立選擇自己的前 t 個 keys,而是將所有層和頭的分數聚合為單個「每個位置的相關性分數」:
依頭部重要性加權 (來自 AM 優化後的預算分配):
    position_score(pos) = Σ_{l,h} head_weight[l,h] · scores_{l,h}(pos)

我們不是給 320 個編輯各自一份手稿讓他們獨立編輯,而是讓他們投票決定保留哪些部分。

在原始 AM 論文中,頭部重要性權重是針對特定模型(例如 Qwen3-4B)透過優化預先計算的。由於我們使用 Qwen3-14B,且目前沒有優化後的預算,我們預設使用均勻的頭部加權。儘管進行了這種簡化,共識訊號仍然有效:許多頭部一致認為值得關注的 token,對應的就是任務相關的上下文。

共享遮罩(mask)允許我們執行批次運算,在極小化效能損失的同時顯著降低開銷。

  1. 使用 MAD 標準化進行閾值處理。我們不選擇固定數量的 token(前 k 個),而是保留所有分數高於統計導出閾值的位置。MAD 標準化提供了一種穩健的離群值指標:
若滿足以下條件則保留位置 i:
    position_scores[i] > median + threshold · MAD

閾值參數控制了壓縮的積極程度:

實現即時壓縮:

原始 AM 演算法無法跨注意力頭進行批次處理,因為每個頭選擇的 token 子集不同(例如,頭 0 保留位置 {12, 45, 89, …},而頭 3 保留 {7, 45, 102, …})。結果導致對應的矩陣形狀不相容,無法堆疊成單個張量運算。這迫使執行必須序列化:對於 Qwen3-14B,需要啟動 320 次獨立的 CUDA kernel,導致 GPU 在等待每次小規模求解完成時大部分處於閒置狀態。雖然這種方法能產生高品質的壓縮,但會產生巨大的延遲(在 A100 GPU 上超過 30 秒),使其不適用於即時 Agent 工作負載。

共享全域遮罩讓我們能將所有 320 次求解堆疊為批次張量運算。自適應批次大小調整器會填滿每個批次的 GPU 記憶體,通常 40 層只需 2–3 個批次。KV 前綴快取(prefix caching)在呼叫之間重複利用了 90% 以上的表徵,因此只需對新 token 進行前向傳遞。

我們透過應用原地(in-place)softmax、順序運行階段、將 KV cache 卸載到 CPU、對預填充(prefills)進行分塊,以及在發生 OOM 錯誤時自動將批次大小減半,來優化 GPU 記憶體。

透過這些修改,AM 壓縮從 30 多秒縮短至約 1.7 秒的中位數開銷,並隨軌跡長度線性擴展,僅佔總呼叫成本的一小部分。

實驗設置

Latent Briefing 如何與 RLM 整合

在標準 RLM 設置中,Orchestrator 透過 llm_query() 向 Worker 發送目標明確的 (上下文, 查詢) 對。Worker 接收後進行處理並返回回應。每次呼叫都是獨立的,Worker 沒有先前呼叫的記憶體。

透過 Latent Briefing,Worker 在多次呼叫中維護一個 Orchestrator 軌跡的持久化 KV cache。每次呼叫時:

  1. Orchestrator 更新後的軌跡(包括新的推理和先前的 Worker 回應)透過 Worker 模型進行前向傳遞,並使用 KV 前綴快取,通常 90% 以上的 token 與前一次呼叫相同並直接重複利用。

  2. Orchestrator 針對此呼叫的任務提示透過對軌跡的注意力產生查詢向量。

  3. 軌跡的 KV cache 使用這些查詢作為相關性訊號進行壓縮。

  4. Worker Agent 使用此壓縮後的 KV cache 進行初始化並產生回應。

壓縮後的快取保留了 Worker 真正需要從 Orchestrator 記憶體中獲取的上下文資訊。

基準測試

我們在 LongBench v2 上評估了一個以 Claude Sonnet 4 作為 Orchestrator、Qwen-14B 作為 Worker 的遞迴語言模型(RLM)。這是一個涵蓋多種文件類型(包括學術論文、法律文件、小說和政府報告)的閱讀理解基準測試。我們在三個資料集和 4 個壓縮閾值下進行了評估:

壓縮閾值:基準線(無壓縮,原生 RLM)加上四個壓縮閾值(t = −1.0, 0.0, 1.0, 2.0)。每個閾值均使用相同的 Orchestrator 提示運行所有問題。


結果

準確度:壓縮與基準線相當或有所提升

在合適的閾值下,Briefing 提升了相對於基準線的準確度,在所有三個條件下均獲得了 +3 pp 的增益。我們發現最佳閾值取決於輸入資料。過度壓縮會丟棄 Worker 需要的資訊,從而降低準確度;而壓縮不足則會留下相對於訊號而言過多的雜訊,稀釋注意力。因此,最佳閾值會隨條件而異。我們將在「分析」章節中進一步討論。

Token 效率:顯著節省,特別是在較長的文件上

在所有三個資料集和四個閾值下,壓縮始終降低了 token 的使用量。在最佳閾值下,Worker 的中位數 token 使用量下降了 42–57%,總中位數 token 使用量下降了 21–31%,同時準確度提升了 3 pp。

相對於基準線的中位數 token 減少量:

最佳閾值下的中位數 token 節省量:


分析:為什麼不同的閾值在不同的情境下勝出

最佳準確度閾值在不同條件下有系統性差異:

較長的文件 → 輕度壓縮勝出 (t=−1.0)。在 32k–100k 的文件中,最佳準確度來自 t=−1.0(18% 壓縮)。較長的文件包含更多分散的資訊。輕度壓縮在保持廣泛覆蓋範圍的同時,仍能達成 57% 的 Worker token 節省。

較難的問題 → 積極壓縮勝出 (t=2.0)。在困難問題上,t=2.0(79% 壓縮)比基準線高出 3 個百分點。困難問題會引導 Orchestrator 探索許多假設,在其軌跡中產生推測性推理。這些推測性內容稀釋了 Worker 的訊號。積極壓縮充當了相關性過濾器,去除了雜訊,為 Worker 提供了更清晰、更聚焦的訊號。

簡短、容易的文件 → 中度壓縮 (t=1.0)。在 Easy <32k 的情況下,t=1.0(68% 壓縮)提供了最佳準確度。Orchestrator 的軌跡較短且更聚焦,因此中度壓縮能有效去除冗餘,而不會冒著資訊遺失的風險。

從概念上講,這有點像做筆記。有時你試圖隨著時間建立知識體系,細節很重要,因為它們會累積成更大的東西。在這種情況下,你希望保留上下文,而不是過早地壓縮它。面對較難的問題時,你通常是在勾勒想法、探索方向、跟隨可能通往任何地方的線索。在這個過程中寫下的內容,大部分並非為了長久保存。


限制

Orchestrator 變異性。Claude Sonnet 4 Orchestrator 是非確定性的,導致針對同一個問題在不同運行中會產生不同的分解策略。在每個條件下 n=42 的情況下,個別結果存在雜訊,儘管總體趨勢是一致的。

單一基準測試。我們僅在 LongBench v2 上進行了測試。其他任務類型可能具有不同的注意力模式和壓縮特性,例如程式碼生成、多文件合成、數學推理。


結論

Latent Briefing 透過直接操作 Worker 模型的內部表徵,減少了多 Agent 系統中的 token 使用量。在我們的實驗中,它在不降低準確度的情況下實現了顯著的 token 節省,同時在 Agent 流水線中部署起來依然切實可行。

該方法具有以下特點:

  • 快速:每次壓縮約 1.7 秒,比序列化 AM 快約 20 倍,比 LLM 摘要快 10–30 倍
  • 任務自適應:不同的查詢會以不同方式壓縮相同的上下文
  • 有效:在減少 token 使用量的同時保持或提升準確度
  • 可預測:MAD 標準化的閾值處理產生了一致的壓縮率

隨著 Agent 架構變得越來越深、越來越廣,跨 Agent 的上下文管理成為了瓶頸。Token 使用量在 Agent 呼叫中不斷累積,使得效率成為系統設計中的首要考量。除了提升個別 Agent 內部的每個 token 的智慧程度外,系統整體中 token 使用的效率也日益重要,這能節省時間與金錢。


¹ 我們透過將 β 初始化為 1(對於分佈良好的注意力而言是一個自然的先驗)來改進收斂,而不是使用預設的「最小二乘法後截斷」方法,後者通常會產生許多負值並被截斷為接近零,這是一個糟糕的起點。

² 資料集中包含難度分數。


研究作者 — Ben Geist @b_geist


想要跟上我們接下來的 AI 實驗嗎?請在此訂閱並在 @RampLabs 追蹤我們。我們也在 Ramp 招募各類職位。