谷歌新架構一戰成名，打破Transformer記憶瓶頸，姚班校友鍾沛林新作_風聞

夢晨 發自 凹非寺

量子位 | 公眾號 QbitAI

想挑戰 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