新架構，再次向Transformer發(fā)起挑戰(zhàn)！

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

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

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

圖片

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

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

圖片

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

圖片

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

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

圖片

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

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

圖片

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

傳統(tǒng)RNN，隱藏狀態(tài)固定大小表達能力受限，也不好并行訓練。

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

最近一系列基于RNN的架構創(chuàng)新中：

RWKV，用線性注意力結合RNN和Transformer的優(yōu)點，在訓練時可以并行計算。

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

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

圖片

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

就像Transformer模型作為一個整體在壓縮互聯(lián)網數(shù)據(jù)到參數(shù)中一樣，可學習的隱藏狀態(tài)模型也在少量參數(shù)上不斷縮上下文信息。

這種“隱藏狀態(tài)模型”隨著時間的推移仍然具有固定的大?。ü潭ǖ哪Ｐ蛥?shù)），但表達能力更強了。

論文的聯(lián)合指導UCSD助理教授王小龍認為：

Transformer顯式地儲存所有輸入token，如果你認為個神經網絡是壓縮信息的好方法，那么壓縮這些token也將是有意義的。

圖片

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

圖片

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

外循環(huán)的參數(shù)變成了內循環(huán)的超參數(shù)，也就是元學習的一個變種了。

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

圖片

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

圖片

在測試時學習

在TTT層里，使用自監(jiān)督學習方法將上下文壓縮到隱藏狀態(tài)。

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

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

圖片

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

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

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

圖片

TTT能否成為“Transformer殺手”？

理論上都走的通了，那么TTT在實驗中表現(xiàn)到底如何？

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

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

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

最終在Pile數(shù)據(jù)集短上下文測試中：

• 2k上下文時，TTT-Linear、Mamba和Transform具有相當?shù)男阅?，TTT-MLP的表現(xiàn)略差。
• 8k上下文時，TTT-Linear和TTT-MLP都優(yōu)于Mamba和Transformer，應用在Transformer骨干的TTT-MLP（T）在1.3B參數(shù)左右也略好與Mamba。

總的來說，隨著上下文長度的增長，TTT層相對于Mamba的優(yōu)勢也會擴大。

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

圖片

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

• 32k上下文，TTT-Linear和TTT-MLP的表現(xiàn)都優(yōu)于曼巴，類似于Pile 8k的觀察。即使是帶有Transformer骨干的TTT-MLP（T）表現(xiàn)也略好于曼巴。
• 1.3B參數(shù)尺度上，TTT-MLP（T）僅比TTT-MLP（M）稍差，Transformer骨干可能更適合論文評估范圍之外的更大模型和更長的上下文。

圖片

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

圖片

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

隱藏狀態(tài)可以是任意模型，但目前的研究只涉及了線性模型和小型MLP，更復雜的還有待研究。

隱藏狀態(tài)模型的學習可以用Adam代替普通的梯度下降等等。

還可用于視頻建模

三位共同一作中：

Yu Sun博士畢業(yè)于UC Berkeley，目前是斯坦福大學博士后。

圖片

Xinhao Li是電子科技大學校友，碩士畢業(yè)于UCSD。

圖片

Karan Dalel本科畢業(yè)于UC Berkley，正在機器人初創(chuàng)公司1X實習。

圖片

最后，聯(lián)合指導UCSD助理教授王小龍還透露，TTT方法除了語言模型，還適用于視頻。

TTT就是“Transformer殺手”，我仍然認為我們需要繼續(xù)努力。

將來在對長視頻進行建模時，我們可以密集地采樣幀而不是采樣1 FPS，這些密集幀對Transformer來說是一種負擔，但對TTT層來說是一種福音。

圖片

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

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