量子位
整合編輯：太平洋科技
發(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)練和評估模型的代碼。 論文地址： 參考鏈接： 本文來源：量子位 |
原創(chuàng)欄目
IT百科
網(wǎng)友評論
聚超值•精選
量子位
硬件編年史
