Sakana AI 推出通用 Transformer 記憶體，新技術最佳化 LLM 記憶體成本降低高達 75%！

總部位於東京的新創公司 Sakana AI 的研究人員開發了一種新技術，使語言模型能夠更有效地利用記憶體，幫助企業降低在大型語言模型 (LLM) 和其他基於 Transformer 的模型之上構建應用程式的成本。

這項技術名為「通用 Transformer 記憶體（Universal Transformer Memory）」，利用特殊的神經網路來最佳化 LLM，保留重要的資訊，並從上下文中移除冗餘的細節。

最佳化 Transformer 記憶體

Transformer 模型是 LLM 的核心，其回應依賴於「上下文視窗」的內容，也就是從使用者那裡接收到的輸入內容。

上下文視窗可視為模型的工作記憶體。調整上下文視窗的內容會對模型的效能產生重大影響，這也催生了一個完整的「提示工程」領域。

目前的模型支援非常長的上下文視窗，可容納數十萬甚至數百萬個 token（LLM 將使用者輸入的文字、詞彙部分、短語、概念及數字轉換成的數值表示）。

這使得使用者能夠將更多資訊塞進他們的提示中。然而，更長的提示會導致更高的運算成本和更慢的效能。最佳化提示以刪除不必要的 token，同時保留重要的資訊，可以降低成本並提高速度。

目前的提示最佳化技術既耗費資源，又需要使用者手動測試不同的配置，以減少提示的大小。

類神經注意力記憶體模組

通用 Transformer 記憶體透過神經注意力記憶模組（Neural Attention Memory Models, NAMMs）來最佳化提示，這是一種簡單的神經網路，可以決定是否「記住」或「遺忘」儲存在 LLM 記憶體中的每個 token。

「這項新功能讓 Transformer 能夠移除無用或冗餘的細節，並專注於最關鍵的資訊，對於需要長上下文推理的任務來說，這尤為重要，」研究人員表示。

▲ Universal transformer memory 。圖片來源：Sakana AI)

NAMM 是在與 LLM 分開的環境中訓練的，並於推理時與預訓練模型結合，這使其具有靈活性且易於部署。然而，NAMM 需要存取模型的內部啟動，這意味著它只能應用於開源模型。

與 Sakana AI 研發的其他技術一樣，NAMM 是透過演化演算法訓練的，而非基於梯度的最佳化方法。透過反覆變異與選擇最具表現力的模型，演化演算法以試錯方式最佳化 NAMM 的效率與性能。這尤其重要，因為 NAMM 正試圖實現一個不可微分的目標：保留或丟棄 token。

NAMM 在 Transformer 的注意力層運作，這是 Transformer 架構的關鍵組成部分，負責確定模型上下文視窗中每個 token 的關係及重要性。根據注意力值，NAMM 決定哪些 token 應保留，哪些應從模型的上下文視窗中移除。這種基於注意力的機制使得訓練過的 NAMM 可以在不同模型間使用，而無需進一步修改。例如，針對純文本資料訓練的 NAMM，可以應用於視覺或多模態模型，而無需額外訓練。

▲ 類神經注意力記憶體模組 (NAMM) 會檢視注意力層，以決定應從上下文視窗中保留或丟棄哪些 token。圖片來源：Sakana AI 

通用記憶體的應用

為測試通用 Transformer 記憶體的概念，研究人員在開放原始碼的 Meta Llama 3-8B 模型之上訓練了一個 NAMM。他們的實驗表明，使用 NAMM，基於 Transformer 的模型在非常長的序列上的自然語言和程式碼問題上表現更好。同時，透過丟棄不必要的 token，NAMM 使 LLM 模型在執行任務時節省了高達 75% 的快取記憶體。

研究人員寫道：「在我們的基準測試中，NAMM 為 Llama 3-8B 變壓器提供了明顯的效能改進。」「此外，我們的記憶體系統產生了顯著的額外好處，減少了每一層的上下文大小，同時從未針對記憶體效率進行明確的優化。」

▲ NAMM 模型在提升模型效能的同時，也能與領先的提示最佳化技術相抗衡。圖片來源：Sakana AI

他們還在 70B 版本的 Llama 以及為其他模態和任務設計的 Transformer 模型上測試了該模型，例如 Llava（電腦視覺）和 Decision Transformer（強化學習）。

研究人員寫道：「即使在這些分佈外的設定中，NAMM 仍然透過丟棄諸如冗餘的影片幀和次佳的動作等 token 來保持其優勢，使其新的基礎模型能夠專注於最相關的資訊以提高效能。」

任務相關的行為

另一個有趣的發現是，NAMM 會根據任務自動調整其行為。

例如，對於程式碼編寫任務，模型會丟棄對應於不影響程式碼執行的註解和空白的連續 token 塊。

另一方面，在自然語言任務中，模型會移除表達語法冗餘的 token，這些 token 並不影響序列的意義。

研究人員已釋出用於創建自有 NAMM 的程式碼。像通用 Transformer 記憶體這樣的技術，對於處理數百萬個 token 且可以從速度提升和成本降低中受益的企業應用程式非常有用。經過訓練的 NAMM 的可重複使用性也使其成為在企業中的不同應用程式中使用的多功能工具。

對於未來，研究人員建議使用更先進的技術，例如在 LLM 的訓練期間使用 NAMM 以進一步擴展其記憶體功能。

「我們的新型記憶模型才剛剛開始展現其潛力，我們期待這將為未來的 Transformer 世代帶來更多突破，」研究人員表示。

• 延伸閱讀：Google 推出 Infini-Transformer模型，可以讓大型語言模型無限長輸入，而不增加記憶體和運算需求
• 延伸閱讀：將330億參數大模型「塞進」單個消費級GPU，加速15%、性能不減
• 延伸閱讀：蘋果論文透露可將大型語言模型裝進快閃記憶體，日後iPhone可直接流暢跑AI模型
• 延伸閱讀：你手上的GPU能跑Llama 2等大模型嗎？用這個開源的專案測測看
• 資料來源：venturebeat