研究人員來自Sakana AI推出NAMMs：針對高效能與高效能變壓器模型的最佳化記憶體管理

變壓器（Transformers）已成為深度學習模型的核心，特別是在需要處理序列數據的任務中，例如自然語言理解、計算機視覺和強化學習。這些模型非常依賴自注意力機制，讓它們能夠捕捉輸入序列中的複雜關係。然而，隨著任務和模型的擴展，對於更長的上下文窗口的需求顯著增加。有效管理這個擴展的上下文窗口非常重要，因為它會影響性能和計算成本。儘管變壓器很強大，但在處理長上下文輸入時仍面臨效率的挑戰，這使得這成為一個活躍的研究領域。

其中一個主要挑戰是平衡性能與資源效率。變壓器會將之前計算的表示存儲在稱為鍵值（Key-Value, KV）快取的記憶體中，這樣可以有效地參考過去的輸入。然而，對於長上下文任務，這個KV快取的大小會指數性增長，消耗大量的記憶體和計算資源。現有的方法試圖通過刪除不太重要的標記來減少KV快取的大小，但這些方法依賴於手動設計的啟發式規則。這些方法的限制很明顯：它們經常導致性能下降，因為標記刪除策略並未優化以保留下游任務所需的重要信息。

目前的工具，例如H2O和L2方法，試圖通過引入像L2範數和熵這樣的指標來量化標記的重要性，來緩解這個問題。這些方法旨在有選擇性地修剪KV快取中的標記，以減少記憶體使用量，同時保持模型性能。儘管這些方法取得了一些成功，但它們引入了一個固有的權衡——減少記憶體佔用會導致性能損失。使用這些技術的模型在不同任務之間的泛化能力較弱，且其基於啟發式的設計使得在性能和效率上無法同時取得顯著改進。

來自日本Sakana AI的研究團隊引入了神經注意力記憶模型（Neural Attention Memory Models, NAMMs）。NAMMs是一種新的記憶管理模型，能動態優化變壓器中的KV快取。NAMMs不依賴於手動設計的規則，而是通過進化優化來學習標記的重要性。通過對變壓器的注意力矩陣進行條件處理，NAMMs使每一層僅保留最相關的標記，從而提高效率和性能，而不改變基礎的變壓器架構。這種通用性使NAMMs適用於任何基於變壓器的模型，因為它們的設計僅依賴於從注意力矩陣中提取的特徵。

NAMMs的工作原理涉及使用基於頻譜圖的技術從注意力矩陣中提取有意義的特徵。研究人員應用短時傅立葉變換（Short-Time Fourier Transform, STFT）將注意力值壓縮為頻譜圖表示。這種緊湊的表示捕捉了標記重要性如何隨著注意力範圍的變化而演變。然後，使用指數移動平均（Exponential Moving Average, EMA）操作來減少頻譜圖特徵的複雜性。NAMMs使用輕量級神經網絡來評估這些壓縮的特徵，並為每個標記分配選擇分數。選擇分數較低的標記會被從KV快取中移除，釋放記憶體，同時確保性能不受影響。

NAMMs的一個關鍵創新是引入了反向注意力機制。這種設計使得網絡能夠有效比較標記，僅保留最相關的出現，同時丟棄冗餘的標記。通過利用跨標記的通信，NAMMs動態優化各層的記憶體使用，確保變壓器為每個任務保留關鍵的長距離信息。

NAMMs的性能在多個基準測試中進行了嚴格評估，顯示出其優於現有方法的能力。在LongBench基準測試中，NAMMs的標準化性能提高了11%，同時將KV快取的大小減少到原始模型的25%。同樣，在挑戰性的InfiniteBench基準測試中，平均輸入長度超過200,000個標記，NAMMs的性能從1.05%提高到11%。這一結果突顯了NAMMs在長上下文任務中有效擴展的能力，而不犧牲準確性。此外，NAMMs在InfiniteBench上的記憶體佔用量減少到原始大小的約40%，顯示了它們在管理長序列方面的效率。

研究人員進一步通過零樣本轉移實驗驗證了NAMMs的多樣性。僅在自然語言任務上訓練的NAMMs被應用於新的變壓器和輸入模式，包括計算機視覺和強化學習模型。例如，在長視頻理解任務中，使用Llava Next Video 7B模型進行測試時，NAMMs提高了基礎模型的性能，同時保持了減少的記憶體佔用。在使用Decision Transformers進行連續控制任務的強化學習實驗中，NAMMs在多個任務中平均獲得了9%的性能提升，顯示了它們能夠丟棄無用信息並改善決策能力。

總結來說，NAMMs為變壓器中的長上下文處理挑戰提供了一個強大的解決方案。通過進化優化學習高效的記憶管理策略，NAMMs克服了手動設計啟發式的限制。結果顯示，配備NAMMs的變壓器在顯著降低計算成本的同時，實現了更優越的性能。它們的通用適用性和在多種任務中的成功，突顯了它們在推進基於變壓器的模型在多個領域中的潛力，標誌著朝向高效長上下文建模的重要一步。

查看論文和詳細資訊。所有的研究成果都歸功於這個項目的研究人員。此外，別忘了在Twitter上關注我們，加入我們的Telegram頻道和LinkedIn小組。也別忘了加入我們的60k+ ML SubReddit。

🚨 熱門消息：LG AI研究發布EXAONE 3.5：三個開源雙語前沿AI級模型，提供無與倫比的指令跟隨和長上下文理解，為生成AI卓越的全球領導地位奠定基礎……。