用 20% 記憶體超越 RAG：ASM 在 NBA Benchmark 的延遲優化

by
目錄

🌏 Read this article in English

長上下文推理（Long Context Inference）的挑戰，往往源於注意力計算的物理限制。

在標準 Transformer 架構中，模型推論時對前綴（Prefix）的處理面臨兩項結構性限制：一是前綴注意力開銷隨長度線性擴張，二是隨著生成推進，前綴對模型行為的影響力逐漸衰減。現有方案通常在「保留完整前綴的高延遲」與「透過梯度下降內化資訊的高維護成本」之間做出取捨。

Attention-State Memory (ASM) 提供了一種無需重新訓練的替代路徑：它不壓縮前綴，也不修改模型權重，而是將前綴轉化為輕量級、基於查找的外部記憶體。這種設計試圖在降低推論開銷的同時，維持甚至超越傳統 In-Context Learning (ICL) 與 RAG 的性能表現。

兩項結構限制：線性開銷與資訊衰減

要理解 ASM 的設計動機，必須先釐清長上下文推理中的兩個物理限制。

第一個限制是注意力運算的線性擴張。在標準 Transformer 架構中，KV Cache 的記憶體佔用以及每個生成 token（Decode step）對既有前綴的注意力計算成本，會隨著前綴長度增加而線性增長。這不僅影響初始預填充（Prefill）階段的延遲，也會拖慢後續的解碼速度。

第二個限制是前綴影響力的衰減。研究指出，隨著生成過程推進，模型的注意力會分佈到更多 token 上，導致初始前綴對模型行為的影響力逐漸降低。這意味著在長上下文場景下，單純增加前綴長度不一定能等比例提升指令遵循或邏輯能力。

現有的路徑存在各自的適用情境：

  1. 壓縮方法（如 KV Cache Compression）：這類方法適合容許近似且以記憶體為主要限制的情境。然而，實驗顯示像 KVZip 這類的壓縮方法在某些設置下表現不如 ASM 或 ICL，這表明過度的 KV cache 壓縮可能會導致 ICL 提示詞中的關鍵標籤資訊發生損耗。
  2. 參數化方法（Training-intensive approach）：透過梯度下降將前綴資訊內化到模型權重中。這種方式需要頻繁的訓練週期，且不利於前綴頻繁更新的應用情境；對於需要動態調整前綴的場景而言，其靈活性較低。

核心機制：注意力狀態記憶體 (ASM)

ASM 提供了一種無需重新訓練（Training-free）的方案，其核心思想是將前綴轉化為一個輕量級、基於查找（Lookup-based）的外部記憶體。

構建階段：從查詢到質心

ASM 的構建過程並非對所有可能的 token 進行預計算，而是透過以下步驟建立結構：

  1. 收集代表性查詢：執行一系列具有代表性的前向傳播（Forward passes），針對特定的查詢集合收集它們在前綴上的注意力輸出。
  2. 聚類與壓縮：將這些注意力輸出進行聚類，提取出「質心」（Centroids）作為記憶體的條目。

推論階段：層級化查找與無損合併

在推論時，ASM 不再對原始前綴執行完整的注意力運算，而是利用以下機制實現效能提升：

1. 層級化查找（Hierarchical Lookup） 透過索引質心，將查找成本降低至 O(logK)。這種解耦機制讓檢索成本與記憶體條目數量的關係由線性轉為對數增長，允許記憶體在不線性增加延遲的情況下增長。這意味著即使記憶體規模擴大，推論延遲也不會按比例上升。

2. 在線 Softmax 恆等式（Online-softmax Identity） 這是 ASM 能實現「無損」合併的核心數學基礎。利用數學上的恆等性質，ASM 可以將查詢與預計算質心的注意力狀態進行合併。這個過程能夠在不實際對前綴進行完整注意力運算的情況下，精確恢復出完整的注意力輸出結果。需注意：此處的「無損」僅限於 online-softmax 的合併步驟；質心聚類與查找仍是近似選擇，因此 ASM 並非對任意查詢都能精確還原原始完整前綴的注意力。

實證數據與性能分析

研究針對 LLaMA-3.1-8B 模型在單張 NVIDIA RTX Ada 4500 GPU 上進行了測試（Batch size 1, 問題長度 512 tokens）。以下是基於實驗的關鍵發現：

準確性與記憶體效率

  • ManyICLBench：在 1K–8K 的記憶體預算下，ASM 在 LLaMA-3.1-8B 上展現出優於 In-Context Learning (ICL) 的準確性。同時，在 8K 記憶體預算下，論文報告其 attention latency 改善 1.36×（沿用論文定義的倍率計算）。
  • NBA Benchmark：ASM 使用僅為全量注意力 RAG 20% 的記憶體佔用，即可超越後者的性能表現。
  • Banking77：ASM 在各種記憶體大小下，其表現均能匹配或超過相同預算的 ICL 基線。

關鍵技術觀察

  1. 非單調性（Non-monotonicity）：實驗發現準確性隨記憶體條目數量的增加並非線性增長，而是在中間的代碼本大小（Codebook size）達到峰值。這意味著最佳的記憶體條目數是一個與前綴長度無關、需針對特定任務進行調優的超參數。
  2. 白化變體（Whitening Variant）：在對分詞敏感的推理任務（如 gpqa_cot, nba）中，使用白化變體能獲得更好的表現。論文作者將此歸因於前導空格 token 在區別答案選項時發揮了關鍵作用。
  3. 長前綴校準的峰值記憶體：論文亦指出，ASM 的可組合結構能以大幅降低的峰值記憶體完成長前綴校準，同時保留完整前綴構建的性能。這補上部署判準的一部分：20% 記憶體結果來自 NBA Benchmark，而長前綴校準主張針對構建階段的峰值記憶體；兩者並非同一項測量，不宜合併解讀。

取捨與邊界 (Tradeoffs & Boundaries)

ASM 的應用並非無限制的優化，其效能受限於以下維度：

1. 離線成本 vs. 線上效率 ASM 以增加離線構建（需要執行前向傳播進行聚類）為代價，換取了線上推論時的對數級查找成本。在單張 NVIDIA RTX Ada 4500 GPU 的測試環境下，這種以預計算換取延遲降低的策略，在需要頻繁重複使用相同長前綴的場景中具備潛力。然而，來源未量化記憶體構建、聚類與更新成本；建議讀者以重用次數、建構成本、線上延遲及準確率共同驗證端到端收益。

2. 適用邊界：查詢分佈的代表性 ASM 的效能建立在「代表性查詢」能覆蓋目標查詢分佈的假設之上。如果線上查詢與離線構建時的查詢分佈存在巨大偏差，其預計算的質心可能無法精確表徵注意力狀態。這意味著 ASM 適合前綴相對靜態或可預測的場景（如固定規則庫），而在高度動態、不可預測的前綴場景下，離線構建的質心可能失效。

3. 適用邊界：前綴更新與重建成本 ASM 的效能亦受前綴動態更新的影響。若線上前綴頻繁變動，預計算的質心將無法直接適用，系統可能需要重新執行離線構建或進行記憶體重建。這增加了工程落地的複雜度，也限制了其在高動態場景下的即時性優勢。

4. 任務敏感性與超參數調優 ASM 的效果在不同任務間存在差異。對於需要精細區分 token 特徵的推理任務（如 gpqa_cot），選擇正確的變體至關重要。同時，由於準確性隨記憶體條目數量呈現非單調性，最佳的代碼本大小需針對特定任務進行調優。

結論

ASM 提供了一種將長前綴從「運算負擔」轉化為「查找任務」的技術路徑。透過無損的在線合併與層級化索引，它在 NBA Benchmark 等特定場景下顯著降低了記憶體佔用，並在多個基準測試中維持甚至提升準確性。然而，其效能受限於查詢分佈的代表性以及代碼本大小（Codebook size）的超參數調優。

這提供了一個值得考慮的方向：在部署前，驗證查詢分佈的穩定性是否足以支撐離線預計算。對於高動態場景，這種離線預計算的邊界究竟在哪裡，可能需要透過更多真實流量的驗證來釐清。

Sources