近日，Google Research 發佈了一篇突破性論文《Titans: Learning to Memorize at Test Time》，提出了新一代神經長時記憶模組，為 AI 模型賦予類似人腦的記憶能力。這一研究不僅突破了 Transformer 的上下文長度瓶頸，更重新定義了長序列任務中的 AI 能力。
Titans 通過引入三種互相協作的記憶模組，建構了“智能大腦”：
- 短期記憶（Short-term memory）：基於注意力機制，負責當前任務中的即時資料處理。這部分與 Transformer 的傳統注意力機制類似，但通過最佳化後的窗口限制解決了資源消耗問題。
- 長期記憶（Long-term memory）：通過參數化記憶歷史資訊，將過去的重要上下文抽象化儲存。這種記憶獨立於短期記憶，專注於長期資料的積累和提取。
- 持久記憶（Persistent memory）：作為任務相關知識的儲存單元，它是固定的、與具體輸入資料無關的模組，確保模型在任務切換時具備更高的泛化能力。
這些記憶模組彼此獨立又相互協作，使 Titans 能像人腦一樣高效利用多層次記憶，適應複雜任務。
Titans 的“驚訝”機制是其記憶系統的核心創新之一。當模型在推理過程中遇到超出預期的輸入時，會觸發“驚訝度”計算：
- 驚訝定義：輸入資料的“驚訝度”由其與歷史資料的差異決定，通過模型對輸入的梯度變化量計算得出。
- 驚訝優先順序：驚訝度越高的資訊會被優先儲存到長期記憶中。
- 驚訝衰減：隨著時間推移，過去的驚訝事件會逐漸衰減其影響力，從而避免記憶的過載和資源浪費。
這一機制模仿了人類記憶的特點——突發事件（如車禍）往往更容易被記住，但隨著時間的推移，其細節會逐漸模糊。
在處理大規模資料或超長文字時，模型可能因記憶累積過多而性能下降。Titans 的自適應遺忘機制，通過以下步驟最佳化記憶管理：
- 遺忘機制：根據當前儲存容量和輸入資訊的驚訝度，動態選擇需要保留或捨棄的資訊。
- 自適應控制：遺忘機制的強度由輸入資料的上下文相關性和重要性決定，確保模型保留關鍵記憶的同時，不浪費儲存資源。
實驗顯示，這種機制使 Titans 能高效處理超過 200 萬 Token 的超長文字，解決了 Transformer 在長序列任務中的計算瓶頸。
Titans 提出了三種記憶整合方案，靈活適配不同任務需求：
1. 記憶作為上下文（MAC）
- 工作原理：通過將歷史和當前上下文結合，提供更全面的資訊給注意力模組。模型在處理新資料時，既能利用歷史記憶，又能專注於當前任務。
- 優勢：適用於需要完整歷史記錄的任務，例如複雜的時間序列預測或跨長時間段的資料分析。
2. 記憶作為門控（MAG）
- 工作原理：採用門控機制平衡短期記憶和長期記憶的影響。短期記憶處理當前輸入，長期記憶提供歷史背景，門控模組動態調整兩者的權重。
- 優勢：更靈活的記憶管理，適合任務需求頻繁變化的場景。
3. 記憶作為層（MAL）
- 工作原理：將記憶模組作為深度網路的一層，負責在當前輸入前壓縮並提煉過去的上下文。
- 優勢：適合對上下文壓縮和抽象要求較高的任務，如基因組學資料分析。
Titans 在多項任務中表現出了強大的競爭力：
- 語言建模與常識推理：在多個基準任務中，Titans 的表現全面超越現有 Transformer 和其他線性模型。
- “大海撈針”檢索任務：在長達 16K 的序列中，Titans 能精準定位關鍵資訊，其檢索精準率遠高於對比模型。
- 基因組學與時間序列分析：通過深度記憶模組，Titans 能有效處理高度複雜和長時間跨度的資料集。
這些結果證明，Titans 不僅可以擴展上下文窗口長度，還能在處理複雜任務時保持高效性。
Google Research 的 Titans 代表了 AI 系統的一次重要進化。通過重新設計記憶模組，它解決了 Transformer 的核心瓶頸，為長序列任務提供了前所未有的性能提升。未來，Titans 有望在自然語言處理、科學計算和即時推理系統中發揮更廣泛的作用。 (本原智數)