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

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

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

一作Ali Behrouz表示：

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

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

新的長期記憶模塊

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

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

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

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

他們從神經(jīng)心理學(xué)中汲取靈感，設(shè)計(jì)了一個(gè)神經(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)隊(duì)還驗(yàn)證了這個(gè)模塊可以并行計(jì)算。

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

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

MAC，記憶作為上下文

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

MAG，記憶作為門

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

MAL，記憶作為層

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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