谷歌新架構終結Transformer，長序列處理王者誕生？清華姚班校友新作Titans

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達

點選進入—>【Mamba/多模態/擴散】交流群

新增微信：CVer2233，助手會拉你進群！

掃描下方二維碼，加入CVer學術星球！可獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪必備！

轉載自：新智元｜編輯：澤正英智

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

Transformer後繼者終於現身！

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

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

正所謂，「風浪越大，魚越貴！」

近日，谷歌研究團隊迎難而上，提出支援200K上下文處理視窗的新架構——Titans。

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

論文連結：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模組，並在ETT、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現為谷歌紐約市演算法與最佳化團隊的研究科學家，該團隊由Vahab Mirrokni領導。他在哥倫比亞大學獲得博士學位（導師為Alex Andoni、Cliff Stein和Mihalis Yannakakis）。此前，他是清華大學交叉資訊研究院（姚班）的本科生。2016年，他以第一作者發表的論文被頂會STOC 2016接收，創下首次有中國本科生在STOC上發表一作論文的記錄。

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

Vahab Mirrokni

Vahab Mirrokni在紐約的谷歌研究院領導演算法與最佳化團隊。團隊包括市場演算法、圖挖掘和大規模最佳化小組。此外，他還在紐約大學庫朗研究所擔任兼職副教授，講授網際網路演算法與經濟學。

參考資料：

https://arxiv.org/abs/2501.00663

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

何愷明在MIT授課的課件PPT下載

在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集

CVPR 2024 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2024，即可下載CVPR 2024論文和程式碼開源的論文合集

Mamba、多模態和擴散模型交流群成立

掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-Mamba、多模態學習或者擴散模型微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。

一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者擴散模型+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群