夢晨 發自 凹非寺
量子位 | 公眾號 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— 完 —
量子位年度AI主題策劃正在徵集中！
歡迎投稿專題 一千零一個AI應用，365行AI落地方案
或與我們分享你在尋找的AI產品，或發現的AI新動向
點這裡👇關注我，記得標星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~