感謝網友 刺客 的線索投遞！

一覺醒來，超越 Transformer 和 Mamba 的新架構誕生了？

斯坦福、UCSD、UC 伯克利和 meta 的研究人員提出了一種全新架構，用機器學習模型取代 RNN 的隱藏狀態。

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

這個模型通過對輸入 token 進行梯度下降來壓縮上下文，這種方法被稱為“測試時間訓練層（Test-Time-Training layers，TTT）”。

TTT 層直接替代了注意力機制，解鎖了具有表現力記憶的線性復雜度架構，使我們能夠在上下文中訓練包含數百萬（未來可能是數十億）個 token 的 LLM。

作者相信，這個研究了一年多的項目，將從根本上改變我們的語言模型方法。

而結果證明，TTT-Linear 和 TTT-MLP 直接趕超或擊敗了最強的 Transformer 和 Mamba！

作者之一的 Xiaolong Wang 驚喜地表示：不敢相信，我們真的做到了。

更令人興奮的是，雖然目前 TTT 只應用于語言建模，但在未來，它也可以用在長視頻上，可謂前景遠大。

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

一個 5 年多的想法，終于實現了

作者表示，在過去的 1.5 年里，團隊一直在開發一種新的 LLM 架構，可以具有線性復雜度和更強的隱藏狀態，用于長上下文建模。

而這個測試時訓練（TTT）的想法，已經研究了超過 5 年。

Xiaolong 清晰記得，在剛開始做博士后時，Alyosha 曾讓自己去找 Yu Sun 討論 TTT。這次會面，就是這項研究的起點。

序列模型會把歷史上下文存儲在一個隱藏狀態中。

像 Mamba 這樣的 RNN 層，會隨著時間的推移壓縮成一個固定大小的狀態，它們雖然效率很高，但性能受限于其表達能力。

注意力機制有一個 KV 緩存，它會隨著時間的推移不斷增長。這個狀態不會壓縮任何歷史上下文，但隨著上下文長度的增加，成本也會越來越高。

團隊成員想：既然這樣，為什么不把上下文壓縮到模型的權重中 —— 就像 LLM 處理互聯網數據那樣呢？

這種「隱藏狀態模型」既能在時間上保持固定大小，又能大大增強表達能力。

研究人員使用了自監督學習來更新隱藏狀態的權重，對每個 token 進行一次梯度下降。在處理一個序列時，該狀態已經在其上下文窗口中的 token 上「訓練」過了。

值得注意的是，隱藏狀態只存在于端到端架構中的一層。其他組件，比如 QKV 投影矩陣，是在預訓練期間通過標準的交叉熵目標函數學習的。

因此，端到端架構實際上是在進行元學習，尋找壓縮上下文的最佳方式，以便更好地預測下一個 token，也就是在「學習如何在測試時學習」。

結果顯示，與 Mamba 相比，TTT-Linear 具有更好的困惑度和更少的 FLOP（左），并且更好地利用了長上下文（右）。

下圖顯示了批大小為 16 的情況下，隨著上下文長度的變化，每個 token 的前向時間（延遲）。所有模型的參數都是 1.3B（Mamba 為 1.4B）。

可以看到，隨著上下文長度的增加，Transformer 每個 token 的前向時間呈線性增長，但其他兩種方法的前向時間基本保持不變。

在 8k 上下文時，TTT-Linear 比 Transformer 更快，與 Mamba 相當。

RNN 的尷尬現實

2020 年，OpenAI 縮放定律論文表明 LSTM（RNN 的一種）無法像 Transformer 那樣進行縮放，或有效地使用長上下文。

真的是這樣嗎？

在這個項目中，研究人員重新評估了圖 2 中的這些發現。

在左側，可以觀察到 Mamba（當今最流行的 RNN 之一）的擴展性與強大的 Transformer 類似，這是自 2020 年的 LSTM 以來顯示出的巨大進步。

然而，在右側，可以觀察到與 OpenAI 相同的 Mamba 問題。

平均而言，序列中靠后的 token 應該更容易預測，因為它們以更多信息為條件。

對 Transformer 來說確實如此，每個 token 索引的平均復雜度在其 32k 上下文中不斷減少。相比之下，Mamba 在 16k 后就出現了同樣的情況。

對于現有的 RNN 來說，這個結果代表了一個尷尬的現實 ——

一方面，RNN（相對于 Transformer）的主要優勢就是它們的線性（相對于二次）復雜性。這種漸進優勢實際上只會在長上下文中實現。

另一方面，一旦上下文足夠長，現有的 RNN（如 Mamba）就很難真正利用額外的條件信息。

長上下文的困難是 RNN 層本質上的問題：與自注意力機制不同，RNN 層必須將上下文壓縮為固定大小的隱藏狀態。

作為一種壓縮啟發式，更新規則需要發現成千上萬甚至數百萬個 token 之間的底層結構和關系。

研究人員首先觀察到，自監督學習可以將大量訓練集壓縮為 LLM 等模型的權重，該模型通常表現出對其訓練數據之間語義聯系的深刻理解，而這，恰恰是他們所需要的。

TTT 層

受此啟發，研究人員設計了一類新的序列建模層，其中隱藏狀態是模型，更新規則是自監督學習的一個步驟。

由于更新測試序列上隱藏狀態的過程，相當于在測試時訓練模型，因此此類新層稱為測試時訓練（TTT）層。

研究人員引入兩個簡單的實例：TTT-Linear 和 TTT-MLP，其中隱藏狀態分別是線性模型和兩層 MLP。TTT 層可以集成到任何網絡架構中并進行端到端優化，類似于 RNN 層和自注意力。

實際運行時間

TTT 層在 FLOP 方面已經非常高效，研究人員則更進一步地提出了兩項創新，使其在實際運行時間內也能保持高效。

首先，與在常規訓練中對 mini-batch 序列采取梯度步進以實現更好的并行性類似，他們也在 TTT 中使用了 mini-batch 的 token。

其次，研究人員為每個 TTT mini-batch 內的操作開發了一種對偶形式，以更好地利用現代 GPU 和 TPU。這種對偶形式的輸出與原始實現相當，但訓練速度卻快了 5 倍以上。

正如圖 3 所示，TTT-Linear 在 8k 上下文中比 Transformer 更快，并且與 Mamba 相當。

Transformer 殺手 ——TTT

如圖 4 所示，所有的序列建模層，都可以從將歷史上下文存儲到隱藏狀態的角度來看待。

比如，RNN 層 —— 如 LSTM、RWKV 和 Mamba 層 —— 將上下文壓縮成一個固定大小的狀態，這個狀態隨時間變化。

這種壓縮帶來了兩種結果：優勢是處理效率高，因為每個 token 的處理時間是恒定的。劣勢是在處理長上下文時，RNN 性能受限于隱藏狀態的「表達能力」。

自注意力機制（Self-attention）也可以從如上角度來理解。

不同之處在于，它的隱藏狀態，通常稱為鍵值（KV）緩存是一個隨 t 增長的線性 list。

它可以存儲所有的上下文，并且不會進行壓縮，具有很好的表達能力，不過其處理時間隨上下文長度線性增長。

因此，為了在長上下文中既保持效率，又具有表達能力，需要一個更好的「壓縮啟發式」（compression heuristic）方法。

具體來說，就需要將數百萬個 token 壓縮成一個能有效捕捉其底層結構和關系的隱藏狀態。

TTT 隱藏狀態

研究人員的關鍵思想是，使用自監督學習來將歷史上下文 x1,...,xt 壓縮成一個隱藏狀態 St。方法是將上下文視為一個無標簽數據集，而將狀態視為一個模型。

具體來說，隱藏狀態 St 現在等同于一個模型 f 的權重 Wt，這個模型 f 可以是線性模型、小型神經網絡或其他任何形式。輸出規則簡單地表示為：zt=f(xt;wt)。

直觀講，輸出 token 就是由更新后權重 Wt 的模型 f 對 xt 所做的預測。更新規則是在某個自監督損失?上進行的一步梯度下降：Wt=Wt-1-ηΔ?（Wt-1;xt）。其中學習率為 η。

從壓縮的角度來看，每種啟發式方法都需要決定記住 / 忘記哪些輸入。W 會記住那些產生大梯度的輸入 —— 直觀地說，就是那些使 W 學習很多的輸入。

?的一種選擇是重構 xt 本身。為了使學習問題變得非平凡，作則首先將 xt 處理成一個被破壞的輸入

，然后優化：

類似于去噪自編碼器，f 需要發現 xt 各維度之間的相關性，以便從部分信息

中重構出 xt。

如圖 5 所示，梯度下降能夠減少?，但無法將其降至零。

與其他 RNN 層和自注意力機制一樣，研究人員將輸入序列 x1,...,xT 映射到輸出序列 z1,...,zt 的算法可以被編程到序列建模層的前向傳播中，使用上述的隱藏狀態、更新規則和輸出規則。

即使在測試時，新層仍然為每個輸入序列訓練一個不同的權重序列 W1,...,Wt。因此，研究人員將其稱之為測試-時間訓練層（TTT）。

使用 TTT 層訓練神經網絡

TTT 層的前向傳播，也有相應的后向傳播。

TTT 層與 RNN 層、自注意力機制有著相同的接口，因此可以在任何更大的神經網絡架構中替換它們。

值得一提的是，訓練帶有 TTT 層神經網絡的方式，與訓練任何其他 Transformer 模型相同。

可以使用相同的數據、方法和目標（如下一個 token 預測）來優化網絡其余部分的參數。

在此，研究人員將訓練更大的神經網絡稱為外循環（outer loop），而在每個 TTT 層內訓練 W 稱為內循環（inner loop）。

它們之間梯度計算的區別是，內循環針對的是 W（即模型 f 的參數），外循環針對的是網絡其余部分的參數 θrest。

TTT 學習自監督任務

可以說，TTT 最重要的部分是自監督任務，因為它決定了 W 從測試序列中學習的特征類型。

在這個任務的設計上，研究人員采取了更加端到端的方法 —— 直接優化自監督任務以實現下一個 token 預測的最終目標。

具體來說，研究者將自監督任務的學習，作為外循環的一部分。

從如上公式 3 中的簡單重構任務開始，添加了一些外循環參數來讓這個任務可學習。最新的自監督損失是：

在內循環中，只有 W 被優化，因此作為?的參數寫出；θ 們是這個損失函數的「超參數」。在外循環中，θK,θV,θQ 與 θrest 一起被優化，而 W 僅僅是一個隱藏狀態，不是參數。

圖 6 用代碼說明了這種區別，其中 θK 和 θV 被實現為 TTT 層的參數，類似于自注意力中的 KV 參數。

總的來說，θK,θV,θQ 所有可能的選擇構成了一系列多視圖重構任務，外循環可以被理解為從這個任務組中選擇一個具體任務。為了簡單起見，研究人員在這里將所有視圖設計為線性投影。

mini-batch TTT 并行化

目前，開發的原生 TTT 層在浮點運算（FLOP）次數方面已經非常高效。

然而，其更新規則 Wt=Wt-1-ηΔ?（Wt-1;xt）無法實現并行化，因為 Wt 在兩個位置上依賴于 Wt-1：負號和 Δ?。

對此，研究人員提出了 mini-batch 梯度下降，用 b 表示 TTT 批大小。

研究中使用 Gt=Δ?（Wt';xt），其中 t'=t-mod（t,d）代表著前一個 mini-batch 的最后一個時間步（或者第一個 mini-batch 0），因此，可以一次并行 b 個梯度計算。

對偶形式

上面介紹的并行化是必要的，但對于「實際運行時間」（wall-clock time）的效率來說還不夠。

正如之前所述，可以對于 t = 1,...,b 進行并行計算：

然而，現實中，是無法對單個 matmul 來計算 GtS 所有的 b。

相反，需要 b 個外積來對其進行一一計算。更糟糕的是，對于每個

，Gt 是 d×d，這會比大 d xt 產生更大的內存占用和 I / O 成本。

為了解決這兩個問題，研究人員觀察到：我們實際上并不需要具體化 G1,...,Gb，只要要我們可以在 mini-batch 結束時計算 Wb，并且輸出 token z1,...,zb（如上圖 7 所示）。

現在，就可以用上面簡化的 TTT-Linear 情況來演示這些計算，表示 X = [x1,...,xb]：

所以 Wb 可以用 matmul 方便地計算出來。為了計算 Z = [z1,...,zb]，我們知道：

表示

和矩陣

，可以得出：

如上過程，研究人員將其稱為「對偶形式」。

理論等價

前面已經提到 f 可以是線性模型，也可以是神經網絡。還有更新規則的三種變體：online GD、batch GD 和 mini-batch GD。

如下圖所示，在這些 2×3 組合中，每一種都會引起 TTT 層的不同實例化。

研究中，作者分別從 2 個定理證明了在這些誘導實例中，具有線性模型和 batch GD 的 TTT 層等同于線性注意力 —— 一個廣為人知的 RNN 層。

圖 10 總結了所有序列建模層的更廣泛范圍內 TTT 層的一般定義。

兩種變體

研究中，作者提出了 TTT 層的兩種變體 TTT-Linear 和 TTT-MLP，僅在 f 的實例化方面有所不同。

對于 TTT-Linear，flin（x）=Wx，其中 W 是平方。對于 TTT-MLP，fMLP 有兩層，類似于 Transfomer 的 MLP。

具體來說，隱藏維度是 4× 輸入維度，然后是 GELU 激活。為了在 TTT 期間獲得更好的穩定性，f 始終包含層歸一化 (LN) 和殘差連接。

即，f（x）=x + LN（fres(x)），其中，fres 可以是 flin 或 fMLP。

實驗

通過與兩個基線 Transformer 和 Mamba（現代 RNN）比較，研究人員評估了 TTT-Linear 和 TTT-MLP。

數據集

繼續 Mamba 論文之后，研究人員在 Pile 上執行了 2k 和 8k 上下文長度的標準實驗，Pile 是一個用于訓練開源 LLM 的流行文檔數據集。

主架構

Transformer 和 Mamba 使用不同的，除非另有說明，TTT-Linear 和 TTT-MLP 始終使用 Mamba 架構。

短上下文：the Pile

在 2k 上下文中，TTT-Linear（M）、Mamba 和 Transformer 具有相當的性能，線條大部分重疊。

TTT-MLP（M）在較大的 FLOP 預算下表現稍差。盡管 TTT-MLP 在每個模型大小上，都比 TTT-Linear 具有更好的復雜度，但 FLOP 的額外成本抵消了這種優勢。

在 8k 上下文中，TTT-Linear（M）和 TTT-MLP（M）的表現均明顯優于 Mamba。即使是具有 Transformer 架構的 TTT-MLP（T），性能也比 Mamba 略好。

另外，研究人員還觀察到了一個非常明顯的現象：隨著上下文長度變長，TTT 層相對于 Mamba 的優勢就更大了。

長上下文：Books

為了評估長上下文中的功能，研究人員使用了 Pile 的一個流行子集 ——Books，對從 1k 到 32k 以 2 個增量的上下文長度進行了實驗。

根據上圖，可以觀察到 ——

在 Books 的 2k 上下文中，Pile 2k 的所有觀察結果仍然成立，唯一的例外是 Mamba 的表現略好于 TTT-Linear。

在 32k 上下文中，TTT-Linear（M）和 TTT-MLP（M）的性能均優于 Mamba，與 Pile 8k 的觀察結果類似。即使具有 Transformer 架構的 TTT-MLP（T），在 32k 上下文中的表現也比 Mamba 稍好。

在 1.3B 尺度上，TTT-MLP（T）僅比 TTT-MLP（M）稍差。由于缺之清晰的線性擬合，很難推導出經驗縮放定律。然而，TTT-MLP（T）的強勁趨勢表明，Transformer 架構可能更適合超出評估的更大模型和更長上下文。

上下文長度作為超參數

雖然輸入序列的長度由用戶確定，但語言模型處理輸入的上下文長度可以由工程師確定。因此，上下文長度也是一個可以選擇的超參數。

對于具有線性復雜度的 LLM，研究人員選擇了困惑度中的 argmin，因為每個上下文長度都有相同的 FLOP。

從圖 13 中，可以觀察到以下結果 ——

- 性能最好的方法 TTT-Linear 和 TTT-MLP 的線幾乎完全重疊。Mamba 和 TF Finetune 的線在 10^20 FLOP 后也大部分重疊。

- TF Finetune 的性能明顯優于 TF Pretrain，因為它受益于長上下文，而不會在訓練 FLOP 中產生極大的成本。

- 對于所有從頭開始訓練的方法（包括 TF 預訓練），一旦上下文長度變得太大，困惑度就會變得更糟。

從上圖可見，與 TTT-Linear 相比，TTT-MLP 在短上下文中表現稍差，但在長上下文中表現更好。

這一觀察結果正符合研究人員的預期，即作為隱藏狀態的 MLP 比線性模型更具表現力。同樣，所有方法都具有與 Mamba 1.4B 相同的訓練 FLOP。

實際運行時間

LLM 訓練和推理可以分解為前向、后向和生成。

由于前向（在訓練和推理期間）和后向都可以并行化，因此研究人員使用對偶形式。生成新 token（也稱為解碼）本質上是順序的，因此研究人員使用原始形式。

由于資源限制，這項實驗是用 JAX 編寫并在 TPU 上運行的。

然而，由于 Mamba（在 PyTorch、Triton 和 CUDA 中實現）只能在 GPU 上運行，因此為了公平比較，研究人員還重寫了方法，以在 GPU 上運行。

具體來說，研究人員在 ThunderKittens 中編寫了一個用于前向的 GPU 內核。從歷史上看，由于并行性和矩陣相乘的使用不當，RNN 在前向和后向過程中效率低下。

這個前向內核的目標，是證明 mini-batch TTT 和這些問題對偶形式的有效性。

圖 15 的左圖顯示了前向內核批大小為 16 的延遲。所有模型參數均為 1.3B（Mamba 為 1.4B）。

對于 Transformer，每個 token 的時間隨著上下文長度的增加而線性增長，但對于其他方法則大致保持不變。

此外，研究人員在 Triton 中編寫了另一個用于生成的 GPU 內核，并在圖 15 的右圖中對批大小為 512 的速度進行了基準測試。

可以看出，TTT-Linear 和 Mamba 的延遲幾乎相同，明顯小于 Transformer 和 TTT-MLP。

Mamba 之后，又看到 TTT 這么能打的新架構誕生，少不了 AI 社區的熱議。

有網友稱，這會不會是最接近實時上下文的方法？很想聽聽大家的想法。這意味著 TTT 甚至在使用過程中，也能夠學習和適應，為長上下文提供更好的性能，而不會產生通常與 Transformer 相關的高昂計算成本。

OpenAI 視頻生成研究人員對此表示，這項研究看起來很有趣。

如果 scaling law 依然存在，TTT 將帶來難以置信的影響。對于長序列，Transformer 的計算成本往往很高，當長序列變得更長時，RNN 會遺忘。TTT 訓練巧妙地利用神經網絡解決 RNN 的不足。

作者介紹

論文最后，分別列出了這篇研究的作者貢獻。

其中的核心作者是，Yu Sun、Xinhao Li 和 Karan Dalal。

Yu Sun

Yu Sun 是斯坦福大學計算機專業的博士后，導師是 Carlos Guestrin、Tatsu Hashimoto 和 Sanmi Koyejo。

此前，他曾在加州大學伯克利分校完成了電子工程科學博士學位，導師是 Alyosha Efros 和 Moritz Hardt。他還在康奈爾大學拿到了學士學位。

個人主頁中，他介紹自己的研究重點是一種名為測試時間訓練（test-time training）的算法框架。其核心思想是，每個測試實例都定義了自己的學習問題，都有自己的泛化目標。這通常使用自監督學習，為每個實例即時訓練一個不同的模型來實現的。

在最新研究中，Yu Sun 與 Xinhao Li 在 2022 年 11 月共同啟動了這一項目。自 2023 年 6 月起，Yu Sun 專職負責該項目。

他提出了項目的概念框架，設計了 mini-batch TTT 和對偶形式（dual form）。

Xinhao Li

Xinhao Li 是 UC San Diego 研二的學生，導師是 Xiaolong Wang 教授。他本人的研究興趣主要是深度學習和計算機視覺。

他在斯坦福大學 Tatsunori Hashimoto 教授的團隊中作為訪問學生，與 Yu Sun 博士和其他導師朋友一起工作。在此之前，他曾在電子科技大學獲得了學士學位。

在 2024 年 3 月之前，Xinhao Li 是 TTT 早期代碼庫的主要貢獻者，這些代碼庫塑造了最新項目。

Karan Dalal

Karan Dalal 是 UC Berkeley 電子工程科學系的本科生。他于 2023 年 6 月全職加入該項目，與 Xinhao Li 合作共同領導了當前代碼庫的開發工作。

參考資料：

https://x.com/karansdalal/status/1810338845659131940 https://x.com/xiaolonw/status/1810387662060269668 https://arxiv.org/abs/2407.04620

廣告聲明：文內含有的對外跳轉鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節省甄選時間，結果僅供參考，所有文章均包含本聲明。