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

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

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

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

近日，谷歌研究團(tuán)隊迎難而上，提出支持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)的情況下有效獲取有用記憶的過程。

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

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

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

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

（1）通過將鍵和值附加到記憶中來更新記憶（無壓縮）；

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

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

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

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

2. 是否需要一個深度記憶模塊，以有效存儲和記住長期歷史信息？

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

該模塊不是無差別地記住所有信息，而是會通過「驚訝度」來選擇性地記住那些重要或令人驚訝的信息。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

實驗結(jié)果

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

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

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

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

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

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

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

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

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

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

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

作者介紹

Ali Behrouz

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

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

Peilin Zhong

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

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

Vahab Mirrokni

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