新架構RNN反超Transformer：每個隱藏狀態都是一個模型

新架構，再次向Transformer發起挑戰！

核心思想：將RNN中的隱藏狀態換成可學習的模型。

甚至在測試時都可以學習，所以該方法稱爲TTT（Test-Time Training）。

共同一作UC伯克利的Karen Dalal表示：我相信這將從根本上改變語言模型。

一個TTT層擁有比RNN表達能力更強的隱藏狀態，可以直接取代Transformer中昂貴的自注意力層。

在實驗中，隱藏狀態是線性模型的TTT-Linear表現超過了Transformer和Mamba，用更少的算力達到更低的困惑度（左），也能更好利用長上下文（右）。

此外，隱藏狀態是MLP模型的TTT-MLP在32k長上下文時表現還要更好。

Karen Dalel還指出，理論上可學習的隱藏狀態可以是任意模型，對於更長上下文來說，可以是CNN、甚至可以是完整的Transformer來套娃。

目前剛剛出爐的TTT論文已經在學術界引起關注和討論，斯坦福博士生Andrew Gao認爲，這篇論文或許能成爲下一篇Attention is all you need。

另外有人表示，衆多新架構能否真正擊敗Transformer，還要看能不能擴展到更大規模。

Karen Dalel透露，馬上就會推出7B模型。

用機器學習模型來壓縮上下文

傳統RNN，隱藏狀態固定大小表達能力受限，也不好並行訓練。

Transformer強大，但自注意力機制隨上下文長度呈平方複雜度，非常昂貴。

最近一系列基於RNN的架構創新中：

RWKV，用線性注意力結合RNN和Transformer的優點，在訓練時可以並行計算。

Mamba，賦予模型選擇性記住或遺忘信息的能力來壓縮上下文，同時設計了面向硬件的高效並行算法。

它們的表現在短上下文時追上甚至超越了Transformer，但在32k超長上下文以上，Trasformer依舊稱霸。

TTT團隊的想法來自於：與其讓隱藏狀態被動地儲存信息，不如讓它主動學習。

就像Transformer模型作爲一個整體在壓縮互聯網數據到參數中一樣，可學習的隱藏狀態模型也在少量參數上不斷縮上下文信息。

這種“隱藏狀態模型”隨着時間的推移仍然具有固定的大小（固定的模型參數），但表達能力更強了。

論文的聯合指導UCSD助理教授王小龍認爲：

如此一來，整個框架的時間複雜度還是線性的，

至此，序列建模被拆解爲兩個嵌套的學習循環，外循環負責整體的語言建模，內循環通過自監督學習壓縮上下文信息。

外循環的參數變成了內循環的超參數，也就是元學習的一個變種了。

標準的元學習是訓練一個適應不同任務的模型，而TTT是讓模型去適應每一個測試樣本。單個樣本雖然信息量小，但用來訓練隱藏狀態模型也綽綽有餘。

特別的，在內循環是一個線性模型時，相當於線性注意力。當內循環是一個Nadaraya-Watson estimator時，TTT等價於自注意力。

在測試時學習

在TTT層裡，使用自監督學習方法將上下文壓縮到隱藏狀態。

上下文就是未標記的數據集，隱藏狀態不再是一個固定的向量，可以是線性模型、小型神經網絡或任何機器學習模型，更新規則採用了在自監督損失上的一步梯度下降。

這樣一來，隱藏狀態模型可以記住產生大梯度的輸入，並且可以獲得比選擇性遺忘機制更強的擬合和泛化能力，並且在測試時仍然爲每個輸入序列訓練不同的參數。

到目前爲止，樸素的TTT層已經有效了，但還無法並行化。

團隊提出的解決方案爲mini-batch梯度下降，把一個batch內的梯度計算並行化。

再通過Dual form方法，只在mini-batch結束時計算權重以及輸出token，避免冗餘計算。在JAX版實現中快了5倍以上。

TTT能否成爲“Transformer殺手”？

理論上都走的通了，那麼TTT在實驗中表現到底如何？

最簡單幹淨的測試方法，應該是直接替換掉Transformer中的自注意力層。

但是在研究過程中，團隊發現Mamba等現代RNN的骨幹中在RNN層之前還包含時間卷積，對TTT也有幫助。

所以實驗中TTT-Linear和TTT-MLP主要應用到Mamba骨幹上，其他訓練細節也嚴格遵照Mamba論文中的設置。

最終在Pile數據集短上下文測試中：

總的來說，隨着上下文長度的增長，TTT層相對於Mamba的優勢也會擴大。

另外團隊猜測，線性模型比MLP表達能力差，因此從Mamba骨幹的卷積中受益更多。

長上下文實驗使用Pile的子集Books3：

在A100上測試速度，TTT-Linear在預填充階段比Mamba稍快，解碼階段幾乎與Mamba速度相同。TTT-MLP相比Transformer整體上也有線性複雜度的優勢。

共同一作Karan Dala表示：我一直被問到的一個問題是，我們是否相信TTT就是“Transformer殺手”，我仍然認爲我們需要繼續努力。

還可用於視頻建模

三位共同一作中：

Yu Sun博士畢業於UC Berkeley，目前是斯坦福大學博士後。

Xinhao Li是電子科技大學校友，碩士畢業於UCSD。

Karan Dalel本科畢業於UC Berkley，正在機器人初創公司1X實習。

最後，聯合指導UCSD助理教授王小龍還透露，TTT方法除了語言模型，還適用於視頻。

TTT就是“Transformer殺手”，我仍然認爲我們需要繼續努力。

論文地址：https://arxiv.org/abs/2407.04620

參考鏈接：[1]https://x.com/karansdalal/status/1810338845659131940[2]https://x.com/xiaolonw/status/1810387662060269668