Google Titan 架構突破 Transformer 記憶瓶頸

Google 新架構 Titan 的問世

科技界正熱烈討論 Google 新推出的 Titan 架構。此架構旨在挑戰 Transformer 模型在記憶處理上的限制。由於是由 Google 團隊開發，Titan 被視為 Transformer 的潛在繼承者，引起廣泛關注。

現有模型的記憶挑戰

傳統模型如 LSTM 和 Transformer 雖然創新，但在模擬人類記憶方面仍面臨挑戰：

容量有限： 資料通常被壓縮成固定大小的隱藏狀態，限制了可保留的資訊量。
計算負擔： 雖然能夠捕捉長程依賴關係，但計算成本隨著序列長度呈二次方增長，導致處理超長序列時效率低下。
過度依賴訓練資料： 單純記憶訓練資料並不總能有效應用於現實世界，因為測試資料可能超出訓練分佈範圍。

Titan 的方法：神經啟發式記憶模組

Titan 團隊採取不同的方法，嘗試將資訊編碼到神經網路的參數中。他們開發了一個線上元模型，旨在學習如何在測試期間記住和遺忘特定資料。這個模型受到神經心理學原理的啟發，包含以下關鍵元素：

驚奇作為觸發器： 意外事件更容易被記住。「驚奇」通過輸入到記憶模組的梯度來衡量。梯度越大，輸入越出乎意料。
動量和遺忘機制： 動量機制將短期驚奇累積到長期記憶中，而遺忘機制則清除舊記憶，防止記憶溢出。
基於多層感知器 (MLP) 的記憶： 記憶模組由多個 MLP 層組成，使其能夠儲存資料的深度抽象，比傳統的基於矩陣的記憶更強大。

這種線上元學習方法有助於模型專注於學習如何適應新資料，而不僅僅是記憶訓練資料。該模組還專為平行計算而設計，提高了其效率。

將記憶模組整合到深度學習架構中

Titan 研究團隊提出了三種將記憶模組整合到深度學習架構中的變體：

MAC (Memory as Context)： 此方法將長期和持久記憶（編碼任務知識）結合為輸入到注意力機制的上下文。
MAG (Memory as Gate)： 此方法使用記憶模組與跨兩個分支的滑動視窗注意力機制進行門控融合。
MAL (Memory as Layer)： 在此方法中，記憶模組被實現為一個獨立層，在將歷史資訊饋送到注意力機制之前壓縮歷史資訊。

該團隊發現每個變體都有其優缺點。

Titan 的性能與優勢

Titan 在包括語言建模、常識推理和時間序列預測在內的各種任務中都展現出卓越的性能。它超越了 Transformer 和 Mamba 等最先進的模型。值得注意的是，長期記憶模組 (LMM) 本身在多項任務中都優於基準模型，展示了其獨立的學習能力，無需短期記憶（注意力）。

在旨在尋找長文本中細微線索的「大海撈針」測試中，即使序列長度從 2k 增加到 16k，Titan 仍保持約 90% 的準確度。該團隊表示，標準測試並未完全展現 Titan 在處理長文本方面的優勢。在需要從分佈在極長文件中的事實進行推斷的任務中，Titan 的表現也優於 GPT4、Mamba 甚至使用 RAG 的 Llama3.1。

Titan 在時間序列預測和 DNA 序列建模等特定領域也表現出令人印象深刻的性能。

Titan 背後的團隊

這項研究由 Google Research NYC 算法和優化團隊進行，目前不屬於 Google DeepMind。

Ali Behrouz，康乃爾大學的實習生，是該論文的第一作者。
鐘沛林，清華大學校友，哥倫比亞大學博士畢業生，自 2021 年起擔任 Google 的研究科學家。他以大學生身份在 STOC 2016 上發表了第一作者論文而聞名。
Vahab Mirrokni，Google 研究員和副總裁，領導該團隊。

該團隊使用 Pytorch 和 Jax 開發了 Titan，並計劃很快發佈用於訓練和評估的程式碼。