近8年后，谷歌Transformer繼任者「Titans」來了，上下文記憶瓶頸被打破

作者：機器之心 2025-01-16 08:20:00

正如論文一作所說，「新架構 Titans 既比 Transformer 和現代線性 RNN 更有效，也比 GPT-4 等超大型模型性能更強。」

終于，在 2017 年推出影響 AI 行業長達 8 年的 Transformer 架構之后，谷歌帶來了全新的架構 Titans。這次，谷歌的重點是將推理領域非常重要的測試時（test-time）計算用在了記憶（memory）層面。

在談到推出 Titans 的初衷時，論文一作 Ali Behrouz 表示，「注意力機制一直是大多數 LLM 進展的重要組成部分，不過它無法擴展到長上下文。因此，Titans 應運而出，它成為了一種同時具備注意力機制和元上下文記憶的結構，可以在測試時學習記憶。該架構可以將上下文窗口擴展到 200 萬 tokens。」

圖源：https://x.com/behrouz_ali/status/1878859086227255347

這意味著，谷歌 Transformer 迎來了它的「繼任者」。

圖源：https://x.com/mark_k/status/1878896628654022993

多年來，研究人員一直在廣泛探究如何有效地利用循環模型和注意力機制，其中循環模型旨在將數據壓縮到固定大小的記憶（稱為隱狀態）中，而注意力機制允許處理整個上下文窗口，捕捉所有 token 的直接依賴。不過，更準確的依賴建模往往伴隨著二次成本，導致模型只能處理固定長度的上下文。

因此，谷歌提出了一種新的長期神經記憶模塊（neural memory module），它能夠學習記憶歷史上下文，并幫助注意力機制在利用過去已久信息的同時處理當前上下文。結果表明，這種神經記憶具有快速并行化訓練的優勢，同時還能保持快速推理。

從記憶的角度來看，谷歌認為注意力機制雖然受限于上下文但可以更準確地建模依賴關系，因此可以起到短期記憶的作用；而神經記憶能夠對數據進行記憶，起到了長期、更持久的記憶作用。基于這兩個模塊，谷歌引入了一個全新的系列架構 —— Titans，通過三種變體有效地將記憶融合到該系統架構中，它們分別是記憶作為上下文（Memory as a Context，MAC）、記憶作為門（Memory as a Gate，MAG）和記憶作為層（Memory as a Layer，MAL）。

在語言建模、常識推理、基因組學和時序預測任務上的實驗結果表明，Titans 架構比 Transformer 和近年來的現代線性循環模型更有效。另外，在大海撈針（needle-in-haystack）中，Titans 架構能夠有效地擴展到超過 200 萬 tokens 的上下文窗口，并且比基準模型實現了更高的準確性。

論文標題：Titans: Learning to Memorize at Test Time
論文地址：https://arxiv.org/pdf/2501.00663v1

另外，論文作者之一 Peilin Zhong 為谷歌 NYC 算法與優化團隊的研究科學家，2021 年加入谷歌。他本科畢業于清華姚班，博士畢業于哥倫比亞大學。

目前，已經有人搞出了有關 Titans 架構的非官方實現，感興趣的讀者可以去看一下。

GitHub 地址：https://github.com/lucidrains/titans-pytorch

學習測試時記憶

谷歌詳細介紹了長期神經記憶模塊，它成為了一種可以在測試時學習記憶的元模型。

長期記憶

為了設計一個長期神經記憶模塊，我們需要模型能夠將過去歷史的抽象編碼到其參數中。因此，一個簡單的思路是訓練神經網絡并期望它能夠記住自己的訓練數據，然而記憶幾乎一直是神經網絡中令人頭疼的現象，它限制了模型的泛化能力，還引發隱私問題，因此導致測試時性能不佳。

基于此，谷歌認為需要一個在線元模型來學習如何在測試時記憶或忘記數據。在這種設置下，模型學習一個能夠記憶的函數，但不會過擬合訓練數據，從而在測試時實現更好的泛化性能。

學習過程和意外指標（Learning Process and Surprise Metric）。訓練長期記憶的關鍵思路是將訓練視為在線學習問題，其中將過去信息 x_1, …, x_t-1 壓縮到長期神經記憶模塊中。人類往往能夠記住背離預期（令人驚訝）的事件，受此啟發，模型意外可以簡單定義為它相對于輸入的梯度。梯度越大，輸入數據與過去數據的偏差就越大。因此，使用這個意外分數，可以將記憶更新如下：

這一意外指標可以導致在重大意外時刻之后出現重要信息缺失。從人類記憶的角度來看，即使一個事件令人難忘，但它可能不會在長時間內持續讓我們感到驚訝。為了改進這一現象，谷歌將意外指標分解為了（1）過去意外，它衡量最近過去的意外程度；（2）瞬時意外，它衡量傳入數據的意外。

這些意外指標基于一個損失函數，它就是我們的記憶在測試時學習充當的目標。也就是說，記憶模塊是一個元模型，它基于損失函數來學習一個函數。

在本文中，谷歌則專注于聯想記憶，目的是將過去的數據存儲為鍵（keys）和值（values）對。類似于 Transformer，在給定 x_t 的情況下，谷歌使用兩個線性層將 x_t 投影到鍵和值中：

接下來，谷歌希望記憶模塊可以學習鍵和值之間的關聯，為此將損失定義如下：

遺忘機制（Forgetting Mechanism）。在處理非常大的序列（比如百萬 tokens）時，管理哪些過去信息應該被遺忘非常重要，即使使用深度或者非常大的矩陣值記憶時也是如此。因此，谷歌使用了一種自適應遺忘機制，允許記憶忘記不再需要的信息，從而更好地管理有限的記憶容量。也就是說，給定下一個 token x_t，谷歌將更新規則做如下修改：

記憶架構（Memory Architecture）。谷歌重點將具有 L_M≥1 層的簡單 MLP 作為長期記憶架構，選擇它們的原因在于希望能夠更好地激勵長期記憶設計以及將其融入架構的方法。谷歌表示，本文的架構開辟了一個新的研究方向，有助于設計更有效且高效記憶數據的神經架構。

檢索記憶（Retrieving a Memory）。在探討如何設計和訓練一個可以在測試時學習記憶的長期記憶模塊之后，剩下的關鍵問題便是如何從記憶中檢索信息？谷歌僅僅使用了沒有更新權重的前向傳遞（即推理）來檢索與查詢相對應的記憶。在形式上，給定一個輸入 x_t，谷歌使用線性層 W_Q 來投影輸入，即 q_t = x_tW_Q，并通過以下公式從記憶 y_t 中檢索相應（或有用）的信息。

并行化長期記憶訓練

理論上，長期記憶模塊的訓練需要 FLOPS，其中 N 為序列長度。不過在實踐中，我們需要并行化訓練過程并充分利用 TPU、GPU 等硬件加速器，同時需要張量化該過程并使用更多矩陣乘法（matmuls）。

接下來，谷歌表示，使用小批量梯度下降、數據學習率和權重衰減來計算內循環權重的方式可以重新來表示，以便它只使用矩陣乘法和求和（sums）。這里將序列拆分為大小為 b ≥ 1 的塊，并將小批量梯度下降寫做：

此外，谷歌解釋了 M_t = W_t 為線性時的情況。對于 N_p ≥ 2 的 MLP，過程類似。使用本文的損失函數可以得到如下：

最后，谷歌擴展這一表示，以便可以合并動量項。在具有動量的塊式梯度下降中，如果觀察動量項則可以得到如下：

作為函數塊的參數（Parameters as the Function of Chunks）。谷歌沒有讓參數 a_t、θ_t 和 η_t 依賴于輸入，而是讓它們成為函數塊。盡管失去了表達能力，但可以幫助更快地訓練。在這種情況下，谷歌在每個塊中對每一個 a、θ 和 η 都使用了相同的值。在實驗中，谷歌將這些參數作為了 token 的函數，并表示，這種簡化（即作為塊函數）可能是未來工作感興趣的地方，以便以更高效的方式訓練更大的模型。

下圖 1 展示了如何并行并在使用矩陣乘法時完成神經記憶訓練。

如何融合記憶？

接下來需要解決的一個重要問題是：如何有效且高效地將神經記憶融合到深度學習架構中？

從記憶的角度來看，Transformer 中的 K 和 V 矩陣對可以解釋為聯想記憶塊。由于它們對依賴關系的精確建模以及有限的上下文窗口，它們可以被用作短期記憶模塊，以處理當前上下文窗口大小。另一方面，神經記憶能夠不斷從數據中學習并存儲在其權重中，因而可以發揮長期記憶的作用。谷歌通過三個不同的 Titans 變體來回答以上問題。

記憶作為上下文（Memory as a Context，MAC）

Titans 的第一個變體 MAC 的架構設計如下圖 2 所示，將記憶作為當前信息的上下文。

也就是說，給定一個長序列，首先將序列分成固定大小的片段 S^(??)，其中 ?? = 1,...,??/??。給定傳入片段 S^(??)，谷歌將它視為當前上下文，將其過去的片段視為歷史信息。因此，谷歌讓 M_??-1 成為片段 S^(??) 之前的長期記憶狀態，使用輸入上下文作為對記憶 M^??-1 的查詢，以從長期記憶中檢索相應的信息。谷歌如下所示檢索與 S^(??) 相對應的過去信息：

接下來將這些歷史信息與持久記憶參數一起作為注意力模塊的輸入序列：

該架構具有兩個關鍵優勢：一是注意力模塊同時具有歷史和當前上下文，能夠根據當前數據決定是否需要長期記憶信息，二是注意力模塊幫助長期記憶只存儲來自當前上下文的有用信息。這意味著，并非每個片段中的所有 token 都是有用的，記憶所有 token 可能會導致內存溢出。因此，注意力模塊幫助記憶了解哪些信息是有用的，從而更好地管理內存容量。

另外，在測試時，（i）持久記憶參數是固定的，它們編碼了有關任務的知識，不應改變；（ii）注意力模塊權重是上下文學習器；（iii）長期記憶模塊在測試時仍然學習（記憶）信息。也就是說，即使在測試時，神經記憶的權重也會更新，這是因為權重對過去已久的抽象進行了編碼。

記憶作為門（Memory as a Gate，MAG）

Titans 第二個變體 MAG 的架構設計如下圖 4 所示：

在其中一個分支中，谷歌直接使用輸入數據來更新長期記憶；在第二個分支中，谷歌使用了滑動窗口注意力（SWA）：

該架構的整體注意力掩碼如下圖 3b 所示，其中滑動窗口注意力（SWA）充當精確的短期記憶，而神經記憶模塊充當模型的衰減記憶。該設計也可以看作是多頭架構，其中各頭的結構不同。

記憶作為層（Memory as a Layer，MAL）

Titans 的第三個變體 MAL 使用了深度神經網絡，這種架構設計在文獻中更為常見，其中混合模型堆疊具有完整或滑動窗口注意力的循環模型。

給定輸入 x，可以得到以下:

其中 SW-Attn 是滑動窗口注意力。

無注意力記憶（Memory Without Attention）。從記憶的角度來看，谷歌期望記憶系統的每個組件都能獨立工作，即使其他組件受到了干擾。因此，即使沒有短期記憶（即注意力），長期記憶模塊仍然應該是一個強大的模型。谷歌在實驗中將這種變體稱為 Titans (LMM)。

架構細節

在所有塊中，谷歌使用了殘差連接；在實現中，谷歌使用 SiLU (.) 激活函數作為計算查詢、鍵和值的非線性激活，并使用對查詢和鍵進行歸一化。

卷積（Convolution）。遵循最近的現代線性循環模型，谷歌在每個查詢、鍵和值投影后都融合了一個 1D 深度可分離卷積層。這些 1D 卷積可以提升性能，并且計算高效。

門控（Gating）。谷歌還在最終輸出投影之前利用線性層進行歸一化和門控。

實驗結果

谷歌在實驗部分關注上述三種 Titans 變體，分別是 MAC、MAG 和 MAL，以及單獨的神經記憶模塊。對于每個模型，谷歌使用了四種尺寸的模型，參數分別是 (i) 170M、(ii) 340M、(iii) 400M 和 (iv) 760M。

語言建模

谷歌首先關注模型在語言建模和常識推理任務中的困惑度。下表 1 報告了 Titans 變體和三種不同大小（340M、400M 和 760M）基線的結果。在包括 Transformer++ 在內的非混合模型中，神經記憶模塊在困惑度和準確度測量方面均取得了最佳性能。

谷歌還發現，Titans 的三種變體（MAC， MAG 和 MAL）都優于 Samba （Mamba + 注意力）和 Gated DeltaNet-H2（Gated DeltaNet + 注意力）。

大海撈針

下表 2 結果顯示，與基線相比，神經記憶模塊均取得了最佳結果。

谷歌將這種卓越的表現歸因于 Titans 與現有序列模型的三個關鍵差異：（1）與 TTT 相比，神經記憶能夠通過使用動量和遺忘機制（即權重衰減）更好地處理記憶容量。因此，隨著序列長度的增加，神經記憶的性能不會下降，呈現出一致的趨勢；（2）與具有門控（遺忘）機制的 Mamba2 相比，Titans 具有深度非線性記憶，從而實現了更好的記憶管理。此外，與神經記憶和 DeltaNet 不同，Mamba2 無法移除記憶，因此在增加序列長度時，其性能會出現顯著下降；（3）與 DeltaNet 相比，盡管它能夠使用增量規則移除記憶，但無法擦除記憶，缺乏遺忘機制。

最終，正如預期的那樣，使用 Titans 變體時能看到相當或更好的結果，其中最佳結果來自 MAC。