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

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

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

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

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

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

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

Titans是什么

研究者認為大多數(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)更新機制類似于人類的學(xué)習(xí)過程，使得模型能夠不斷適應(yīng)新的數(shù)據(jù)和任務(wù)需求。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

因此，研究者認為，訓(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)一直讓我們感到驚訝。

因此為了改進上述驚訝度度量，作者將驚訝度度量分為兩部分：（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é)會如何在測試時記憶鍵與值之間的映射。

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

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

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

實驗結(jié)果

在語言建模及常識推理任務(wù)中，對340M、400M、760M等不同參數(shù)規(guī)模下的Titans變體與多種基線模型進行對比。非混合模型里，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)了卓越的性能，能夠有效擴展到超過200萬的上下文窗口，超越了GPT-4、Llama3+RAG和Llama3-70B等大模型。

Titans (MAC) 的參數(shù)量遠少于基線模型，展現(xiàn)出在長序列推理方面的高效性和強大能力。在微調(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)也展示了其強大的長序列處理能力。實驗結(jié)果表明，Titans架構(gòu)在這些任務(wù)中能夠有效地利用歷史信息，從而提高模型的性能。

消融研究表明，神經(jīng)記憶模塊的所有組件對模型性能均有積極貢獻，特別是權(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é)計算機科學(xué)系的二年級博士生，同時也是Google Research的研究實習(xí)生。在加入康奈爾之前，他在哥倫比亞大學(xué)師從Margo Seltzer教授，取得計算機科學(xué)碩士學(xué)位。

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

Peilin Zhong

Peilin Zhong現(xiàn)為谷歌紐約市算法與優(yōu)化團隊的研究科學(xué)家，該團隊由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ā)表一作論文的記錄。

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

Vahab Mirrokni

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