Transformer后繼者終于現(xiàn)身！

2017年，Attention Is All You Need首次引入注意力機(jī)制，成為現(xiàn)代LLM誕生標(biāo)志。

雖然Transformer依然大行其道，但其架構(gòu)的缺陷卻飽受詬病，尤其是無(wú)法擴(kuò)展更長(zhǎng)上下文。

正所謂，「風(fēng)浪越大，魚(yú)越貴！」

近日，谷歌研究團(tuán)隊(duì)迎難而上，提出支持200K上下文處理窗口的新架構(gòu)——Titans。

最重要的是，輕松擴(kuò)展到2M上下文的Titans架構(gòu)，要比Transformer和線性RNN更加有效。

論文鏈接：https://arxiv.org/abs/2501.00663

Titans是什么

研究者認(rèn)為大多數(shù)現(xiàn)有架構(gòu)將記憶視為由輸入引起的神經(jīng)更新，并將學(xué)習(xí)定義為在給定目標(biāo)的情況下有效獲取有用記憶的過(guò)程。

從這個(gè)角度來(lái)看，循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）可以被定義為具有向量值記憶模塊?（也稱(chēng)為隱藏狀態(tài)）的模型，其主要步驟包括：在時(shí)間t給定新輸入時(shí)，

（1）使用函數(shù)更新記憶（帶有壓縮）；

（2）使用函數(shù)檢索輸入的相應(yīng)記憶。

類(lèi)似地，Transformer可以被視為具有不斷增長(zhǎng)的記憶和兩個(gè)相似步驟的架構(gòu)。即：

（1）通過(guò)將鍵和值附加到記憶中來(lái)更新記憶（無(wú)壓縮）；

（2）通過(guò)查找查詢(xún)向量與鍵向量的相似性來(lái)檢索查詢(xún)向量的相應(yīng)記憶，然后將其用于加權(quán)值向量以生成輸出。

由于記憶分為短期記憶、工作記憶和長(zhǎng)期記憶，而其中每個(gè)部分都相互獨(dú)立地服務(wù)于不同的場(chǎng)景，也具有不同的神經(jīng)結(jié)構(gòu)。

受此啟發(fā)，研究者提出了兩個(gè)問(wèn)題：

1. 如何設(shè)計(jì)一個(gè)高效架構(gòu)，將不同且相互關(guān)聯(lián)的記憶模塊整合起來(lái)？

2. 是否需要一個(gè)深度記憶模塊，以有效存儲(chǔ)和記住長(zhǎng)期歷史信息？

本研究旨在通過(guò)設(shè)計(jì)一個(gè)長(zhǎng)期神經(jīng)記憶模塊來(lái)解決上述問(wèn)題，神經(jīng)長(zhǎng)期記憶模塊的設(shè)計(jì)受到人類(lèi)長(zhǎng)期記憶系統(tǒng)的啟發(fā)，能存儲(chǔ)和檢索過(guò)去的信息。

該模塊不是無(wú)差別地記住所有信息，而是會(huì)通過(guò)「驚訝度」來(lái)選擇性地記住那些重要或令人驚訝的信息。

并且其記憶不是靜態(tài)的，可以根據(jù)新的信息動(dòng)態(tài)更新。這種動(dòng)態(tài)更新機(jī)制類(lèi)似于人類(lèi)的學(xué)習(xí)過(guò)程，使得模型能夠不斷適應(yīng)新的數(shù)據(jù)和任務(wù)需求。

為了更好地管理有限的內(nèi)存，模塊引入了衰減機(jī)制。該機(jī)制根據(jù)記憶的大小和數(shù)據(jù)的驚訝程度來(lái)調(diào)整記憶的權(quán)重，從而優(yōu)化內(nèi)存管理。

長(zhǎng)期神經(jīng)記憶模塊設(shè)計(jì)完成后，面臨的一個(gè)關(guān)鍵問(wèn)題是如何把記憶高效地整合進(jìn)深度學(xué)習(xí)架構(gòu)。研究者提出了Titans架構(gòu)，由三個(gè)模塊構(gòu)成：

1. 核心模塊（Core）：包含短期記憶，負(fù)責(zé)主要的數(shù)據(jù)處理流程，采用具有有限窗口大小的注意力機(jī)制。

2. 長(zhǎng)期記憶模塊（Long-term Memory）：此模塊是研究者設(shè)計(jì)的神經(jīng)長(zhǎng)期記憶模塊，負(fù)責(zé)存儲(chǔ)和記住遠(yuǎn)距離的歷史信息。

3. 持久記憶模塊（Persistent Memory）：這是一組可學(xué)習(xí)但與數(shù)據(jù)無(wú)關(guān)的參數(shù)，主要用于對(duì)任務(wù)知識(shí)進(jìn)行編碼，為模型提供先驗(yàn)知識(shí)儲(chǔ)備。

在此基礎(chǔ)上，研究者提出了Titans架構(gòu)的三種變體：

記憶作為上下文（MAC）架構(gòu)

核心分支把對(duì)應(yīng)的長(zhǎng)期記憶、持久記憶和當(dāng)前輸入信息拼接在一起，然后用注意力機(jī)制來(lái)處理上下文，并決定哪些信息應(yīng)存儲(chǔ)在長(zhǎng)期記憶中。

在測(cè)試時(shí)，與上下文記憶對(duì)應(yīng)的參數(shù)仍在學(xué)習(xí)，與核心分支對(duì)應(yīng)的參數(shù)負(fù)責(zé)上下文學(xué)習(xí)，而持久記憶的參數(shù)則負(fù)責(zé)存儲(chǔ)任務(wù)相關(guān)知識(shí)，因此是固定的。

記憶作為門(mén)控（MAG）架構(gòu)

在此架構(gòu)中，一個(gè)分支用輸入數(shù)據(jù)更新長(zhǎng)期記憶，另一個(gè)分支使用滑動(dòng)窗口注意力（SWA），最后將兩者結(jié)果通過(guò)門(mén)控機(jī)制組合。

在此設(shè)計(jì)中，滑動(dòng)窗口注意力充當(dāng)精確的短期記憶，而神經(jīng)記憶模塊則作為模型的衰減記憶。這種架構(gòu)設(shè)計(jì)也可視為一種多頭架構(gòu)，其中頭的結(jié)構(gòu)各不相同。

與MAC架構(gòu)不同的是，MAG架構(gòu)僅將持久記憶融入上下文，并通過(guò)門(mén)控機(jī)制將記憶與核心分支結(jié)合。門(mén)控機(jī)制決定了來(lái)自持久記憶的信息在多大程度上影響核心分支的處理結(jié)果。

記憶作為層（MAL）架構(gòu)

將神經(jīng)記憶模塊作為深度神經(jīng)網(wǎng)絡(luò)的一層，結(jié)合滑動(dòng)窗口注意力機(jī)制。記憶層的核心功能是對(duì)過(guò)去和當(dāng)前的上下文信息進(jìn)行壓縮處理，之后將處理結(jié)果傳遞給注意力模塊。

在測(cè)試時(shí)去學(xué)習(xí)記憶

神經(jīng)長(zhǎng)期記憶模塊

對(duì)于神經(jīng)網(wǎng)絡(luò)來(lái)講，記憶能力通常反而會(huì)限制模型的泛化能力，并可能引發(fā)隱私問(wèn)題，導(dǎo)致在測(cè)試時(shí)性能下降。

此外，由于測(cè)試數(shù)據(jù)可能屬于分布外數(shù)據(jù)，訓(xùn)練數(shù)據(jù)的記憶在測(cè)試時(shí)可能就并沒(méi)有什么效用。

因此，研究者認(rèn)為，訓(xùn)練長(zhǎng)期記憶的關(guān)鍵思想是將其訓(xùn)練視為一個(gè)在線學(xué)習(xí)問(wèn)題，學(xué)會(huì)在測(cè)試時(shí)如何記住或忘記數(shù)據(jù)。在這種設(shè)置中，模型學(xué)習(xí)的是一個(gè)能夠記憶的函數(shù)，但不會(huì)過(guò)擬合訓(xùn)練數(shù)據(jù)，從而在測(cè)試時(shí)實(shí)現(xiàn)更好的泛化。

學(xué)習(xí)過(guò)程與目標(biāo)函數(shù)：由于令人驚訝的事件對(duì)人類(lèi)來(lái)說(shuō)更易記住。受此啟發(fā)，作者將「驚訝度」定義為相對(duì)于輸入的梯度。梯度越大，輸入數(shù)據(jù)與過(guò)去數(shù)據(jù)的差異就越大。因此，利用這個(gè)驚訝度，我們可以更新記憶如下：

如此就能將過(guò)去的信息壓縮到長(zhǎng)期神經(jīng)記憶模塊的參數(shù)中。

然而，這種驚訝度的度量方法也可能會(huì)導(dǎo)致錯(cuò)過(guò)一些重要信息。也就是說(shuō)，在若干個(gè)驚訝步驟之后，梯度可能變得非常小，從而導(dǎo)致陷入平坦區(qū)域（即局部最小值），錯(cuò)失序列中的某些信息。

而從人類(lèi)記憶的角度來(lái)看，某個(gè)事件雖然值得記住，但可能不會(huì)在很長(zhǎng)時(shí)間內(nèi)一直讓我們感到驚訝。

因此為了改進(jìn)上述驚訝度度量，作者將驚訝度度量分為兩部分：（1）過(guò)去的驚訝，衡量最近過(guò)去的驚訝度；（2）瞬時(shí)驚訝，衡量即將到來(lái)的數(shù)據(jù)的驚訝度。

在此公式中，項(xiàng)是數(shù)據(jù)依賴(lài)的驚喜衰減，控制著驚喜隨時(shí)間如何衰減；而項(xiàng) θt 則控制著應(yīng)以數(shù)據(jù)依賴(lài)的方式將多少瞬時(shí)驚喜納入最終的驚喜度量中。

這種數(shù)據(jù)依賴(lài)性在此設(shè)計(jì)中尤為重要：雖然前一個(gè)標(biāo)記的驚喜可能影響下一個(gè)標(biāo)記的驚喜，但這主要在所有標(biāo)記相關(guān)且處于同一上下文時(shí)才有效。

因此，數(shù)據(jù)依賴(lài)的η可以控制記憶是否需要：

（1）通過(guò)設(shè)置→0忽略上一次的驚喜（可能由于上下文的變化）

（2）通過(guò)設(shè)置→1完全納入上一次的驚喜（可能因?yàn)樵摌?biāo)記與其最近的過(guò)去標(biāo)記高度相關(guān)）。

在本工作中，作者專(zhuān)注于關(guān)聯(lián)記憶，旨在將過(guò)去的數(shù)據(jù)存儲(chǔ)為鍵值對(duì)。即給定，類(lèi)似于Transformers，使用兩個(gè)線性層將其投影為鍵和值：

接著，作者希望記憶模塊能夠?qū)W習(xí)鍵和值之間的關(guān)聯(lián)。為此，定義損失函數(shù)如下：

通過(guò)在元模型的內(nèi)循環(huán)中優(yōu)化上述損失函數(shù)，模型就可以學(xué)會(huì)如何在測(cè)試時(shí)記憶鍵與值之間的映射。

遺忘機(jī)制：在處理非常大的序列（例如，數(shù)百萬(wàn)個(gè)標(biāo)記）時(shí)，明確哪些過(guò)去信息應(yīng)該被遺忘至關(guān)重要。為此，作者使用了一種自適應(yīng)遺忘機(jī)制，允許內(nèi)存遺忘不再需要的信息，從而更好地管理內(nèi)存的有限容量。也就是說(shuō)，給定下一個(gè)標(biāo)記，然后將更新規(guī)則修改為：

其中是靈活控制記憶的門(mén)控機(jī)制；即決定應(yīng)遺忘多少信息。例如，它可以通過(guò)讓來(lái)更新記憶而不影響過(guò)去的抽象，并可以通過(guò)讓來(lái)清除整個(gè)記憶。

檢索記憶：作者簡(jiǎn)單地使用不更新權(quán)重的前向傳遞（即推理）來(lái)檢索與查詢(xún)對(duì)應(yīng)的記憶。形式上，給定輸入，使用線性層 WQ 來(lái)投影輸入，即，并通過(guò)以下方式從記憶中檢索相應(yīng)（或有用的）信息：

實(shí)驗(yàn)結(jié)果

在語(yǔ)言建模及常識(shí)推理任務(wù)中，對(duì)340M、400M、760M等不同參數(shù)規(guī)模下的Titans變體與多種基線模型進(jìn)行對(duì)比。非混合模型里，Titans (LMM) 在困惑度和準(zhǔn)確率上表現(xiàn)優(yōu)異。

混合模型對(duì)比中，Titans的三個(gè)變體均比基線模型更好。MAC和MAG整體性能高于MAL，能更好地整合注意力和記憶模塊。

在S-NIAH任務(wù)里，基于RULER基準(zhǔn)測(cè)試，對(duì)2K、4K、8K 和 16K長(zhǎng)度序列予以評(píng)估。神經(jīng)記憶模塊相較基線模型優(yōu)勢(shì)顯著。在Titans變體中，MAC性能最佳。

在BABILong基準(zhǔn)測(cè)試中，Titans (MAC) 展現(xiàn)了卓越的性能，能夠有效擴(kuò)展到超過(guò)200萬(wàn)的上下文窗口，超越了GPT-4、Llama3+RAG和Llama3-70B等大模型。

Titans (MAC) 的參數(shù)量遠(yuǎn)少于基線模型，展現(xiàn)出在長(zhǎng)序列推理方面的高效性和強(qiáng)大能力。在微調(diào)設(shè)置環(huán)節(jié)，Titans（MAC）的表現(xiàn)更為出色。

研究發(fā)現(xiàn)，增加記憶深度可提升模型在較長(zhǎng)序列上的性能，并改善困惑度，但訓(xùn)練速度會(huì)因此降低，呈現(xiàn)出性能與效率之間的權(quán)衡。

通過(guò)在Simba框架中替換Mamba模塊，并在ETT、ECL、Traffic和Weather等基準(zhǔn)數(shù)據(jù)集上測(cè)試，神經(jīng)記憶模塊超越了所有的基線模型。這表明其在處理時(shí)間序列任務(wù)中的潛在優(yōu)勢(shì)。

在DNA建模任務(wù)中，Titans架構(gòu)也展示了其強(qiáng)大的長(zhǎng)序列處理能力。實(shí)驗(yàn)結(jié)果表明，Titans架構(gòu)在這些任務(wù)中能夠有效地利用歷史信息，從而提高模型的性能。

消融研究表明，神經(jīng)記憶模塊的所有組件對(duì)模型性能均有積極貢獻(xiàn)，特別是權(quán)重衰減和動(dòng)量。MAC和MAG在語(yǔ)言建模和常識(shí)推理上表現(xiàn)相近，但MAC在長(zhǎng)上下文任務(wù)中表現(xiàn)最佳。

Titans架構(gòu)通過(guò)引入神經(jīng)長(zhǎng)期記憶模塊，顯著提高了模型在處理長(zhǎng)序列數(shù)據(jù)時(shí)的性能和效率。

實(shí)驗(yàn)結(jié)果表明，Titans架構(gòu)在語(yǔ)言建模、常識(shí)推理、時(shí)間序列預(yù)測(cè)和DNA建模等任務(wù)中均表現(xiàn)出色，特別是在處理超200萬(wàn)上下文窗口任務(wù)中，能夠有效地利用歷史信息，提高模型的準(zhǔn)確性。

作者介紹

Ali Behrouz

Ali Behrouz目前是康奈爾大學(xué)計(jì)算機(jī)科學(xué)系的二年級(jí)博士生，同時(shí)也是Google Research的研究實(shí)習(xí)生。在加入康奈爾之前，他在哥倫比亞大學(xué)師從Margo Seltzer教授，取得計(jì)算機(jī)科學(xué)碩士學(xué)位。

他對(duì)多樣化的研究主題感興趣，目前正致力于深度學(xué)習(xí)架構(gòu)、圖表示學(xué)習(xí)、醫(yī)療保健中的機(jī)器學(xué)習(xí)以及計(jì)算神經(jīng)科學(xué)等領(lǐng)域的研究。

Peilin Zhong

Peilin Zhong現(xiàn)為谷歌紐約市算法與優(yōu)化團(tuán)隊(duì)的研究科學(xué)家，該團(tuán)隊(duì)由Vahab Mirrokni領(lǐng)導(dǎo)。他在哥倫比亞大學(xué)獲得博士學(xué)位（導(dǎo)師為Alex Andoni、Cliff Stein和Mihalis Yannakakis）。此前，他是清華大學(xué)交叉信息研究院（姚班）的本科生。2016年，他以第一作者發(fā)表的論文被頂會(huì)STOC 2016接收，創(chuàng)下首次有中國(guó)本科生在STOC上發(fā)表一作論文的記錄。

他對(duì)理論計(jì)算機(jī)科學(xué)有廣泛的興趣，主要集中在算法的設(shè)計(jì)與分析上。一些特別感興趣的領(lǐng)域包括并行和大規(guī)模并行算法、隱私算法、草圖算法、流算法、圖算法、機(jī)器學(xué)習(xí)、高維幾何、度量嵌入、數(shù)值線性代數(shù)、聚類(lèi)以及其他與大規(guī)模數(shù)據(jù)計(jì)算相關(guān)的算法。

Vahab Mirrokni

Vahab Mirrokni在紐約的谷歌研究院領(lǐng)導(dǎo)算法與優(yōu)化團(tuán)隊(duì)。團(tuán)隊(duì)包括市場(chǎng)算法、圖挖掘和大規(guī)模優(yōu)化小組。此外，他還在紐約大學(xué)庫(kù)朗研究所擔(dān)任兼職副教授，講授互聯(lián)網(wǎng)算法與經(jīng)濟(jì)學(xué)。