首頁 > 科技要聞 > 科技> 正文

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

量子位 整合編輯:太平洋科技 發(fā)布于:2025-01-21 11:09

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

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

OpenAI 把它用在推理(Reasoning),谷歌這次把它用在了記憶(Memory)。

一作Ali Behrouz表示:

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

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

新的長期記憶模塊

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

要么將數(shù)據(jù)壓縮到固定大小的隱狀態(tài),容量有限

要么可以捕捉長程依賴,但計算開銷隨序列長度平方級增長

并且,僅僅記住訓(xùn)練數(shù)據(jù)在實際使用時可能沒有幫助,因為測試數(shù)據(jù)可能在分布外。

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

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

意料之外的事件(即“驚喜”)更容易被記住。

驚喜程度由記憶模塊對輸入的梯度來衡量,梯度越大說明輸入越出人意料。

引入動量機制遺忘機制,前者將短期內(nèi)的驚喜累積起來形成長期記憶,后者可以擦除不再需要的舊記憶,防止記憶溢出。

記憶模塊由多層 MLP 組成,可以存儲深層次的數(shù)據(jù)抽象,比傳統(tǒng)的矩陣記憶更強大。

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

另外,團隊還驗證了這個模塊可以并行計算。

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

為此,Titans 提出了三種變體:

MAC,記憶作為上下文

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

MAG,記憶作為門

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

MAL,記憶作為層

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

本文來源:量子位

網(wǎng)友評論

聚超值•精選

推薦 手機 筆記本 影像 硬件 家居 商用 企業(yè) 出行 未來
二維碼 回到頂部