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

夢(mèng)晨 2025-01-14 15:07:43 來(lái)源：量子位

夢(mèng)晨 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

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

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

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

一作Ali Behrouz表示：

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

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

新的長(zhǎng)期記憶模塊

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

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

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

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

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

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

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

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

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

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

MAC，記憶作為上下文

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

MAG，記憶作為門

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

MAL，記憶作為層

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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