想挑戰 Transformer 的新架構有很多，來自谷歌的“正統”繼承者Titan架構更受關注。

英偉達把測試時間計算（Test-time Computing）稱為大模型的第三個 Scaling Law。

OpenAI 把它用在推理（Reasoning），谷歌這次把它用在了記憶（Memory）。

一作Ali Behrouz表示：

Titans 比 Transformers 和現代線性 RNN 更高效，并且可以有效地擴展到超過 200 萬上下文窗口，性能比 GPT4、Llama3 等大模型更好。

他還解釋了這篇研究的動機，團隊認為 Transformer 中的注意力機制表現為短期記憶，因此還需要一個能記住很久以前信息的神經記憶模塊。

新的長期記憶模塊

提到記憶，大家可能會想到 LSTM、Transformer 等經典模型，它們從不同角度模擬了人腦記憶，但仍有局限性：

• 要么將數據壓縮到固定大小的隱狀態，容量有限

• 要么可以捕捉長程依賴，但計算開銷隨序列長度平方級增長

并且，僅僅記住訓練數據在實際使用時可能沒有幫助，因為測試數據可能在分布外。

為此，Titans 團隊打算將過去信息編碼到神經網絡的參數中，訓練了一個在線元模型（Online meta-model），該模型學習如何在測試時記住/忘記特定數據。

他們從神經心理學中汲取靈感，設計了一個神經長期記憶模塊，它借鑒了人腦原理：

• 意料之外的事件（即“驚喜”）更容易被記住。?

• 驚喜程度由記憶模塊對輸入的梯度來衡量，梯度越大說明輸入越出人意料。

• 引入動量機制和遺忘機制，前者將短期內的驚喜累積起來形成長期記憶，后者可以擦除不再需要的舊記憶，防止記憶溢出。

• 記憶模塊由多層 MLP 組成，可以存儲深層次的數據抽象，比傳統的矩陣記憶更強大。

這種在線元學習范式，避免了模型記住無用的訓練數據細節，而是學到了如何根據新數據調整自己，具有更好的泛化能力。

另外，團隊還驗證了這個模塊可以并行計算。

如何將這個強大的記憶模塊融入深度學習架構中呢？

為此，Titans 提出了三種變體：

MAC，記憶作為上下文

將長期記憶和持久記憶（編碼任務知識的不變參數）作為當前輸入的上下文，一起輸入給 attention。

MAG，記憶作為門

在記憶模塊和滑動窗口 attention 兩個分支上進行門控融合。

MAL，記憶作為層

將記憶模塊作為獨立的一層，壓縮歷史信息后再輸入給 attention。

在實驗中，發現每種方法都有自己的優缺點。

Titans 在語言建模、常識推理、時間序列預測等任務上全面超越 Transformer 和 Mamba 等各路架構的 SOTA 模型。

并且僅靠長期記憶模塊（LMM，Long-term Memory Module）本身，就在多個任務上擊敗基線。

證明了即使沒有短期記憶（也就是 Attention），該機制也具備獨立學習的能力。

在長文本中尋找細粒度線索的“大海撈針”測試中，序列長度從 2k 增加到 16k，準確率保持在 90% 左右。

但團隊認為，這些通用的測試已經體現不出 Titans 在長文本上的優勢。

在另一項需要對分布在極長文檔中的事實做推理的任務中，Titans 表現超過了 GPT4 、Mamba 等，以及 Llama3.1 + RAG 的系統。

另外在時間序列預測、DNA 序列建模等特定任務中，Titans 也取得不錯的表現。

三位作者來自 Google Research NYC 算法和優化團隊，目前還沒被合并到 Google DeepMind。

一作是Ali Behrouz來自康奈爾大學的實習生。

鐘沛林是清華姚班校友，博士畢業于哥倫比亞大學，2021 年起加入谷歌任研究科學家。

2016 年，鐘沛林本科期間的一作論文被頂會 STOC 2016 接收，是首次有中國本科生在 STOC 上發表一作論文。

領隊的Vahab Mirrokni是 Google Fellow 以及 VP。

團隊表示 Titians 是用 Pytorch 和 Jax 中實現的，打算很快提供用于訓練和評估模型的代碼。

論文地址：
https://arxiv.org/abs/2501.00663v1

參考鏈接：
[1]https://x.com/behrouz_ali/status/1878859086227255347