張量緩存：用於Transformer的基於驅逐條件的關聯記憶
張量緩存是一種兩層級緩存架構，結合滑動窗口注意力作為一級緩存（L1）和固定大小的外積快速權重記憶作為二級緩存（L2），L2由窗口驅逐的KV對填充。該方法通過線性注意力恆等式實現高效讀取，並引入可學習的門控融合L1和L2輸出。實驗表明，張量緩存在記憶-質量邊界上優於有狀態基線。
文章情報
要點
- 提出張量緩存，一種兩層級緩存機制，結合精確局部注意力和壓縮記憶。
- 二級緩存使用外積快速權重記憶，僅由滑動窗口驅逐的KV對填充。
- 識別並修復了常見分塊均值訓練捷徑中的虛假交叉標記外積問題。
- 在多個基準測試中，張量緩存改進了記憶-質量邊界。
為甚麼重要
這條新聞值得關注，因為提出張量緩存，一種兩層級緩存機制，結合精確局部注意力和壓縮記憶。
技術影響
可能影響模型選型、推理成本、產品能力和評測基準。
近日，一篇題為“Tensor Cache: Eviction-conditioned Associative Memory for Transformers”的論文在arXiv上發佈，提出了一種新穎的兩層級緩存架構，旨在解決自迴歸Transformer模型中鍵值（KV）緩存隨上下文長度線性增長的問題。傳統的滑動窗口緩存雖然能夠限制內存使用，但會完全丟棄被驅逐的令牌，導致窗口外的重要信息無法被訪問，從而限制了模型處理長上下文的能力。
張量緩存（Tensor Cache）通過兩層級設計巧妙地將精確的局部注意力與壓縮的關聯記憶結合起來。第一級緩存（L1）採用標準的滑動窗口softmax注意力，保留最近令牌的精確注意力計算。第二級緩存（L2）則是一個固定大小的外積快速權重記憶，其內容完全由從滑動窗口中驅逐的鍵值對填充。每當一個鍵值對被驅逐出窗口時，它會被壓縮到每層的一個矩陣A中。未來查詢可以通過一次矩陣乘法高效地讀取這些壓縮記憶，這得益於線性注意力恆等式q_t(k_i ⊗ v_i) = ⟨q_t, k_i⟩v_i。這種設計使得模型能夠以最小的內存開銷訪問窗口外的歷史信息。
論文的作者為Kabir Swain等五位研究人員，他們貢獻了多個關鍵創新。首先，他們引入了可學習的標量門來融合L1和L2的輸出，並且每個注意力頭都有獨立的衰減率和寫入率參數，這些參數通過端到端訓練得到。其次，他們識別並修正了常見分塊均值訓練捷徑中存在的問題。該捷徑使用A←λA+η(ķ⊗v)進行更新，但這樣會在每個分塊中引入C²-C個虛假的跨令牌外積，從而破壞了記憶的準確性。作者通過並行加權和掃描（等同於每令牌寫入，誤差在float32 epsilon內）來彌補這一差距，恢復了正確的記憶更新。
在系統擴展性、受控關聯回憶、長上下文語言建模和記憶容量診斷等一系列實驗中，張量緩存相比有狀態基線顯著改進了記憶-質量邊界。實驗結果表明，該方法能夠有效地平衡內存效率與長上下文表現，為Transformer模型的推理優化提供了新思路。這項工作可能對需要長上下文理解的應用程序（如文檔分析、對話系統和代碼生成）產生重要影響。