Google新架構終結Transformer,長序列處理王者誕生?清華姚班校友新作

新智元報導  

編輯:澤正 英智

【新智元導讀】Google團隊提出的Titans架構通過引入神經長期記憶模塊,突破了傳統Transformer架構在長序列處理中的局限。該架構通過創新的記憶整合和遺忘機制,在語言建模、常識推理、時間序列預測等任務中展現了顯著的性能提升,在長上下文任務中的優勢突出。

Transformer後繼者終於現身!

2017年,Attention Is All You Need首次引入注意力機制,成為現代LLM誕生標誌。

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

正所謂,「風浪越大,魚越貴!」

近日,Google研究團隊迎難而上,提出支持200K上下文處理窗口的新架構——Titans。

最重要的是,輕鬆擴展到2M上下文的Titans架構,要比Transformer和線性RNN更加有效。

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

Titans是什麼

研究者認為大多數現有架構將記憶視為由輸入引起的神經更新,並將學習定義為在給定目標的情況下有效獲取有用記憶的過程。

從這個角度來看,循環神經網絡(RNN)可以被定義為具有向量值記憶模塊ℳ(也稱為隱藏狀態)的模型,其主要步驟包括:在時間t給定新輸入時,

(1)使用函數

更新記憶(帶有壓縮);

(2)使用函數

檢索輸入的相應記憶。

類似地,Transformer可以被視為具有不斷增長的記憶和兩個相似步驟的架構。即:

(1)通過將鍵和值附加到記憶中來更新記憶(無壓縮);

(2)通過查找查詢向量與鍵向量的相似性來檢索查詢向量的相應記憶,然後將其用於加權值向量以生成輸出。

由於記憶分為短期記憶、工作記憶和長期記憶,而其中每個部分都相互獨立地服務於不同的場景,也具有不同的神經結構。

受此啟發,研究者提出了兩個問題:

1. 如何設計一個高效架構,將不同且相互關聯的記憶模塊整合起來?

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

本研究旨在通過設計一個長期神經記憶模塊來解決上述問題,神經長期記憶模塊的設計受到人類長期記憶系統的啟發,能存儲和檢索過去的信息。

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

並且其記憶不是靜態的,可以根據新的信息動態更新。這種動態更新機制類似於人類的學習過程,使得模型能夠不斷適應新的數據和任務需求。

為了更好地管理有限的內存,模塊引入了衰減機制。該機制根據記憶的大小和數據的驚訝程度來調整記憶的權重,從而優化內存管理。

長期神經記憶模塊設計完成後,面臨的一個關鍵問題是如何把記憶高效地整合進深度學習架構。研究者提出了Titans架構,由三個模塊構成:

1. 核心模塊(Core):包含短期記憶,負責主要的數據處理流程,採用具有有限窗口大小的注意力機制。

2. 長期記憶模塊(Long-term Memory):此模塊是研究者設計的神經長期記憶模塊,負責存儲和記住遠距離的歷史信息。

3. 持久記憶模塊(Persistent Memory):這是一組可學習但與數據無關的參數,主要用於對任務知識進行編碼,為模型提供先驗知識儲備。

在此基礎上,研究者提出了Titans架構的三種變體:

記憶作為上下文(MAC)架構

核心分支把對應的長期記憶、持久記憶和當前輸入信息拚接在一起,然後用注意力機制來處理上下文,並決定哪些信息應存儲在長期記憶中。

在測試時,與上下文記憶對應的參數仍在學習,與核心分支對應的參數負責上下文學習,而持久記憶的參數則負責存儲任務相關知識,因此是固定的。

記憶作為門控(MAG)架構

在此架構中,一個分支用輸入數據更新長期記憶,另一個分支使用滑動窗口注意力(SWA),最後將兩者結果通過門控機制組合。

在此設計中,滑動窗口注意力充當精確的短期記憶,而神經記憶模塊則作為模型的衰減記憶。這種架構設計也可視為一種多頭架構,其中頭的結構各不相同。

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

記憶作為層(MAL)架構

將神經記憶模塊作為深度神經網絡的一層,結合滑動窗口注意力機制。記憶層的核心功能是對過去和當前的上下文信息進行壓縮處理,之後將處理結果傳遞給注意力模塊。

在測試時去學習記憶

神經長期記憶模塊

對於神經網絡來講,記憶能力通常反而會限制模型的泛化能力,並可能引發隱私問題,導致在測試時性能下降。

此外,由於測試數據可能屬於分佈外數據,訓練數據的記憶在測試時可能就並沒有什麼效用。

因此,研究者認為,訓練長期記憶的關鍵思想是將其訓練視為一個在線學習問題,學會在測試時如何記住或忘記數據。在這種設置中,模型學習的是一個能夠記憶的函數,但不會過擬合訓練數據,從而在測試時實現更好的泛化。

學習過程與目標函數:由於令人驚訝的事件對人類來說更易記住。受此啟發,作者將「驚訝度」定義為相對於輸入的梯度。梯度越大,輸入數據與過去數據的差異就越大。因此,利用這個驚訝度,我們可以更新記憶如下:

如此就能將過去的信息

的參數中。

壓縮到長期神經記憶模塊

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

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

因此為了改進上述驚訝度度量,作者將驚訝度度量分為兩部分:(1)過去的驚訝,衡量最近過去的驚訝度;(2)瞬時驚訝,衡量即將到來的數據的驚訝度。

在此公式中,項

是數據依賴的驚喜衰減,控制著驚喜隨時間如何衰減;而項 θt 則控制著應以數據依賴的方式將多少瞬時驚喜納入最終的驚喜度量中。

這種數據依賴性在此設計中尤為重要:雖然前一個標記的驚喜可能影響下一個標記的驚喜,但這主要在所有標記相關且處於同一上下文時才有效。

因此,數據依賴的η可以控制記憶是否需要:

(1)通過設置

→0忽略上一次的驚喜(可能由於上下文的變化)

(2)通過設置

→1完全納入上一次的驚喜(可能因為該標記與其最近的過去標記高度相關)。

在本工作中,作者專注於關聯記憶,旨在將過去的數據存儲為鍵值對。即給定

,類似於Transformers,使用兩個線性層將其投影為鍵和值:

接著,作者希望記憶模塊能夠學習鍵和值之間的關聯。為此,定義損失函數如下:

通過在元模型的內循環中優化上述損失函數,模型就可以學會如何在測試時記憶鍵與值之間的映射。

遺忘機制:在處理非常大的序列(例如,數百萬個標記)時,明確哪些過去信息應該被遺忘至關重要。為此,作者使用了一種自適應遺忘機制,允許內存遺忘不再需要的信息,從而更好地管理內存的有限容量。也就是說,給定下一個標記

,然後將更新規則修改為:

其中

來清除整個記憶。

來更新記憶而不影響過去的抽像,並可以通過讓

是靈活控制記憶的門控機制;即決定應遺忘多少信息。例如,它可以通過讓

檢索記憶:作者簡單地使用不更新權重的前向傳遞(即推理)來檢索與查詢對應的記憶。形式上,給定輸入,使用線性層 WQ 來投影輸入,即

中檢索相應(或有用的)信息:

,並通過以下方式從記憶

實驗結果

在語言建模及常識推理任務中,對340M、400M、760M等不同參數規模下的Titans變體與多種基線模型進行對比。非混合模型里,Titans (LMM) 在困惑度和準確率上表現優異。

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

在S-NIAH任務里,基於RULER基準測試,對2K、4K、8K 和 16K長度序列予以評估。神經記憶模塊相較基線模型優勢顯著。在Titans變體中,MAC性能最佳。

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

Titans (MAC) 的參數量遠少於基線模型,展現出在長序列推理方面的高效性和強大能力。在微調設置環節,Titans(MAC)的表現更為出色。

研究發現,增加記憶深度可提升模型在較長序列上的性能,並改善困惑度,但訓練速度會因此降低,呈現出性能與效率之間的權衡。

通過在Simba框架中替換Mamba模塊,並在E湯臣、ECL、Traffic和Weather等基準數據集上測試,神經記憶模塊超越了所有的基線模型。這表明其在處理時間序列任務中的潛在優勢。

在DNA建模任務中,Titans架構也展示了其強大的長序列處理能力。實驗結果表明,Titans架構在這些任務中能夠有效地利用歷史信息,從而提高模型的性能。

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

Titans架構通過引入神經長期記憶模塊,顯著提高了模型在處理長序列數據時的性能和效率。

實驗結果表明,Titans架構在語言建模、常識推理、時間序列預測和DNA建模等任務中均表現出色,特別是在處理超200萬上下文窗口任務中,能夠有效地利用歷史信息,提高模型的準確性。

作者介紹

Ali Behrouz

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

他對多樣化的研究主題感興趣,目前正致力於深度學習架構、圖表示學習、醫療保健中的機器學習以及計算神經科學等領域的研究。

Peilin Zhong

Peilin Zhong現為Google紐約市算法與優化團隊的研究科學家,該團隊由Vahab Mirrokni領導。他在哥倫比亞大學獲得博士學位(導師為Alex Andoni、Cliff Stein和Mihalis Yannakakis)。此前,他是清華大學交叉信息研究院(姚班)的本科生。2016年,他以第一作者發表的論文被頂會STOC 2016接收,創下首次有中國本科生在STOC上發表一作論文的記錄。

他對理論計算機科學有廣泛的興趣,主要集中在算法的設計與分析上。一些特別感興趣的領域包括並行和大規模並行算法、隱私算法、草圖算法、流算法、圖算法、機器學習、高維幾何、度量嵌入、數值線性代數、聚類以及其他與大規模數據計算相關的算法。

Vahab Mirrokni

Vahab Mirrokni在紐約的Google研究院領導算法與優化團隊。團隊包括市場算法、圖挖掘和大規模優化小組。此外,他還在紐約大學庫朗研究所擔任兼職副教授,講授互聯網算法與經濟學。

參考資料:

https://arxiv.org/abs/2501.00663

https://x.com/behrouz_ali/status/1878859086227255347