Sakana AI推出LLM記憶管理技術NAMMs，可將內存成本降低75%

記憶是認知的關鍵組成部分，讓人類能夠從充斥我們生活的無盡噪音中選擇性存儲和提取重要信息。相比之下，大語言模型（Large Language Model，LLM）則缺乏這種能力，只能無差別地存儲和處理所有過去的輸入，這在長時間任務中對它們的性能和成本產生了嚴重影響。

就像人類大腦會選擇性地保留重要信息並逐漸淡忘不重要的細節一樣，人工智能系統也需要這樣的智能記憶管理機制。否則，隨著模型規模的不斷擴大，其訓練和部署對計算資源和內存的需求也會無限上升。

長期以來，研究人員一直在探索如何讓 AI 系統具備更接近人類的記憶能力。傳統的解決方案主要依賴預設的規則來管理模型的記憶，比如基於時間順序或注意力分數（Attention Score）來選擇性地保留或丟棄信息。

然而，這些方法往往過於機械，無法像人類記憶那樣智能地區分信息的重要性，導致在提高效率的同時往往會損害模型的性能。

在此背景下，日本初創公司 Sakana AI 的研究團隊提出了一種新的解決方案——神經注意力記憶模型（Neural Attention Memory Models，NAMMs）。

這一方案借鑒了自然進化在塑造人類記憶系統中發揮的關鍵作用，通過進化算法訓練一個專門的神經網絡來，這種方法能夠像人類大腦一樣主動選擇和保留重要信息，從而在提高效率的同時提升模型的性能。

就像人類大腦會根據信息的長期使用價值來決定是否保留一樣，NAMMs 通過對注意力模式的分析來評估信息的重要性。其核心機制包括三個關鍵組成部分：特徵提取系統、記憶管理網絡和進化優化策略。

圖丨 NAMMs 執行過程中的三個主要步驟（來源：arXiv）

首先是特徵提取機制。NAMMs 採用短時傅里葉變換（STFT）來處理注意力矩陣的列向量。具體來說，它使用大小為 32 的 Hann 窗口進行處理，生成 17 個複值頻率的頻譜圖表示。

這種表示方式非常巧妙，因為它既保留了注意力值隨時間變化的頻率特徵，又大大壓縮了數據量。研究團隊通過實驗發現，這種頻譜表示比直接使用原始注意力值或手工設計的特徵更有效。

其次是向後注意力記憶（BAM）架構的設計。這是 NAMMs 的核心創新之一，它引入了一種特殊的注意力機制，允許 token 只能關注其在 KV 緩存中的「未來」相關內容。

這種設計的妙處在於，它能夠建立 token 之間的競爭關係，使得模型能夠學會保留最有信息量的 token。例如，當出現重覆的句子或詞語時，模型會傾向於保留最新的出現，因為它包含了更完整的上下文信息。

圖丨 NAMMs 的反向掩碼使每個 token 僅關注 KV 緩存中的未來相關詞（來源：arXiv）

在優化策略上，研究團隊採用了 CMA-ES（協方差矩陣自適應進化策略）算法。傳統的梯度下降法難以處理記憶管理這種具有離散決策的問題，而 CMA-ES 通過模擬自然進化過程，能夠直接優化非可微的目標函數。

具體來說，團隊採用了增量進化的方式，從單個任務開始，逐步增加訓練任務的數量，這種方式能夠提供更好的正則化效果，提高模型的泛化能力。

研究團隊選擇 Llama 3-8b 作為基礎模型訓練了 NAMMs，並在 LongBench、InfiniteBench 以及 ChouBun 上進行了全面評估。結果顯示，NAMM 為 Llama 3-8b Transformer 帶來了明顯的性能提升，在總體表現上超越了已有的 H2O 和 L2 這兩種手工設計的內存管理方法。

例如，在 LongBench 基準測試中，NAMMs 不僅將 KV 緩存大小減少到原來的 25%，還實現了 11% 的性能提升。在 InfiniteBench 測試中模型性能從基線的 1.05% 提升到了 11%，同時將緩存大小減少到原來的 40%。