谷歌新架構(gòu)一戰(zhàn)成名，打破Transformer記憶瓶頸，姚班校友鐘沛林新作

夢晨 2025-01-14 15:07:43 來源：量子位

夢晨 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

想挑戰(zhàn) Transformer 的新架構(gòu)有很多，來自谷歌的“正統(tǒng)”繼承者?Titan?架構(gòu)更受關(guān)注。

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

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

一作Ali Behrouz表示：

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

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

新的長期記憶模塊

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

• 要么將數(shù)據(jù)壓縮到固定大小的隱狀態(tài)，容量有限
• 要么可以捕捉長程依賴，但計算開銷隨序列長度平方級增長

并且，僅僅記住訓(xùn)練數(shù)據(jù)在實(shí)際使用時可能沒有幫助，因?yàn)闇y試數(shù)據(jù)可能在分布外。

為此，Titans 團(tuán)隊打算將過去信息編碼到神經(jīng)網(wǎng)絡(luò)的參數(shù)中，訓(xùn)練了一個在線元模型（Online meta-model），該模型學(xué)習(xí)如何在測試時記住/忘記特定數(shù)據(jù)。

他們從神經(jīng)心理學(xué)中汲取靈感，設(shè)計了一個神經(jīng)長期記憶模塊，它借鑒了人腦原理：

• 意料之外的事件（即“驚喜”）更容易被記住。
• 驚喜程度由記憶模塊對輸入的梯度來衡量，梯度越大說明輸入越出人意料。
• 引入動量機(jī)制和遺忘機(jī)制，前者將短期內(nèi)的驚喜累積起來形成長期記憶，后者可以擦除不再需要的舊記憶，防止記憶溢出。
• 記憶模塊由多層 MLP 組成，可以存儲深層次的數(shù)據(jù)抽象，比傳統(tǒng)的矩陣記憶更強(qiáng)大。

這種在線元學(xué)習(xí)范式，避免了模型記住無用的訓(xùn)練數(shù)據(jù)細(xì)節(jié)，而是學(xué)到了如何根據(jù)新數(shù)據(jù)調(diào)整自己，具有更好的泛化能力。

另外，團(tuán)隊還驗(yàn)證了這個模塊可以并行計算。

如何將這個強(qiáng)大的記憶模塊融入深度學(xué)習(xí)架構(gòu)中呢？

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

MAC，記憶作為上下文

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

MAG，記憶作為門

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

MAL，記憶作為層

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

在實(shí)驗(yàn)中，發(fā)現(xiàn)每種方法都有自己的優(yōu)缺點(diǎn)。

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

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

證明了即使沒有短期記憶（也就是 Attention），該機(jī)制也具備獨(dú)立學(xué)習(xí)的能力。

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

但團(tuán)隊認(rèn)為，這些通用的測試已經(jīng)體現(xiàn)不出 Titans 在長文本上的優(yōu)勢。

在另一項(xiàng)需要對分布在極長文檔中的事實(shí)做推理的任務(wù)中，Titans 表現(xiàn)超過了 GPT4 、Mamba 等，以及 Llama3.1 + RAG 的系統(tǒng)。

另外在時間序列預(yù)測、DNA 序列建模等特定任務(wù)中，Titans 也取得不錯的表現(xiàn)。

三位作者來自 Google Research NYC 算法和優(yōu)化團(tuán)隊，目前還沒被合并到 Google DeepMind。

一作是?Ali Behrouz?來自康奈爾大學(xué)的實(shí)習(xí)生。

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

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

領(lǐng)隊的?Vahab Mirrokni?是 Google Fellow 以及 VP。

團(tuán)隊表示 Titians 是用 Pytorch 和 Jax 中實(shí)現(xiàn)的，打算很快提供用于訓(xùn)練和評估模型的代碼。

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

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