超越Transformer和Mamba的新架構(gòu)，剛剛誕生了。斯坦福UCSD等機(jī)構(gòu)研究者提出的TTT方法，直接替代了注意力機(jī)制，語(yǔ)言模型方法從此或?qū)氐赘淖儭?/p>

一覺醒來(lái)，超越Transformer和Mamba的新架構(gòu)誕生了？

斯坦福、UCSD、UC伯克利和Meta的研究人員提出了一種全新架構(gòu)，用機(jī)器學(xué)習(xí)模型取代RNN的隱藏狀態(tài)。

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

這個(gè)模型通過對(duì)輸入token進(jìn)行梯度下降來(lái)壓縮上下文，這種方法被稱為「測(cè)試時(shí)間訓(xùn)練層（Test-Time-Training layers，TTT）」。

TTT層直接替代了注意力機(jī)制，解鎖了具有表現(xiàn)力記憶的線性復(fù)雜度架構(gòu)，使我們能夠在上下文中訓(xùn)練包含數(shù)百萬(wàn)（未來(lái)可能是數(shù)十億）個(gè)token的LLM。

作者相信，這個(gè)研究了一年多的項(xiàng)目，將從根本上改變我們的語(yǔ)言模型方法。

而結(jié)果證明，TTT-Linear和TTT-MLP直接趕超或擊敗了最強(qiáng)的Transformer和Mamba！

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

更令人興奮的是，雖然目前TTT只應(yīng)用于語(yǔ)言建模，但在未來(lái)，它也可以用在長(zhǎng)視頻上，可謂前景遠(yuǎn)大。

在將來(lái)，當(dāng)我們對(duì)長(zhǎng)視頻進(jìn)行建模時(shí)，就可以對(duì)幀進(jìn)行密集采樣，而不是采樣1FPS了。這些密集幀對(duì)Transformer是一種負(fù)擔(dān)，但對(duì)于TTT層來(lái)說(shuō)，這卻是一種福音！

01 一個(gè)5年多的想法，終于實(shí)現(xiàn)了

作者表示，在過去的1.5年里，團(tuán)隊(duì)一直在開發(fā)一種新的LLM架構(gòu)，可以具有線性復(fù)雜度和更強(qiáng)的隱藏狀態(tài)，用于長(zhǎng)上下文建模。

而這個(gè)測(cè)試時(shí)訓(xùn)練（TTT）的想法，已經(jīng)研究了超過5年。

Xiaolong清晰記得，在剛開始做博士后時(shí)，Alyosha曾讓自己去找Yu Sun討論TTT。

這次會(huì)面，就是這項(xiàng)研究的起點(diǎn)。

序列模型會(huì)把歷史上下文存儲(chǔ)在一個(gè)隱藏狀態(tài)中。

像Mamba這樣的RNN層，會(huì)隨著時(shí)間的推移壓縮成一個(gè)固定大小的狀態(tài)，它們雖然效率很高，但性能受限于其表達(dá)能力。

注意力機(jī)制有一個(gè)KV緩存，它會(huì)隨著時(shí)間的推移不斷增長(zhǎng)。這個(gè)狀態(tài)不會(huì)壓縮任何歷史上下文，但隨著上下文長(zhǎng)度的增加，成本也會(huì)越來(lái)越高。

團(tuán)隊(duì)成員想：既然這樣，為什么不把上下文壓縮到模型的權(quán)重中——就像LLM處理互聯(lián)網(wǎng)數(shù)據(jù)那樣呢？

這種「隱藏狀態(tài)模型」既能在時(shí)間上保持固定大小，又能大大增強(qiáng)表達(dá)能力。

研究人員使用了自監(jiān)督學(xué)習(xí)來(lái)更新隱藏狀態(tài)的權(quán)重，對(duì)每個(gè)token進(jìn)行一次梯度下降。在處理一個(gè)序列時(shí)，該狀態(tài)已經(jīng)在其上下文窗口中的token上「訓(xùn)練」過了。

值得注意的是，隱藏狀態(tài)只存在于端到端架構(gòu)中的一層。其他組件，比如QKV投影矩陣，是在預(yù)訓(xùn)練期間通過標(biāo)準(zhǔn)的交叉熵目標(biāo)函數(shù)學(xué)習(xí)的。

因此，端到端架構(gòu)實(shí)際上是在進(jìn)行元學(xué)習(xí)，尋找壓縮上下文的最佳方式，以便更好地預(yù)測(cè)下一個(gè)token，也就是在「學(xué)習(xí)如何在測(cè)試時(shí)學(xué)習(xí)」。

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

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

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

在8k上下文時(shí)，TTT-Linear比Transformer更快，與Mamba相當(dāng)。

02 RNN的尷尬現(xiàn)實(shí)

2020年，OpenAI縮放定律論文表明LSTM（RNN的一種）無(wú)法像Transformer那樣進(jìn)行縮放，或有效地使用長(zhǎng)上下文。

真的是這樣嗎？

在這個(gè)項(xiàng)目中，研究人員重新評(píng)估了圖2中的這些發(fā)現(xiàn)。

在左側(cè)，可以觀察到Mamba（當(dāng)今最流行的RNN之一）的擴(kuò)展性與強(qiáng)大的Transformer類似，這是自2020年的LSTM以來(lái)顯示出的巨大進(jìn)步。

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

平均而言，序列中靠后的token應(yīng)該更容易預(yù)測(cè)，因?yàn)樗鼈円愿嘈畔闂l件。

對(duì)Transformer來(lái)說(shuō)確實(shí)如此，每個(gè)token索引的平均復(fù)雜度在其32k上下文中不斷減少。相比之下，Mamba在16k后就出現(xiàn)了同樣的情況。

對(duì)于現(xiàn)有的RNN來(lái)說(shuō)，這個(gè)結(jié)果代表了一個(gè)尷尬的現(xiàn)實(shí)——

一方面，RNN（相對(duì)于Transformer）的主要優(yōu)勢(shì)就是它們的線性（相對(duì)于二次）復(fù)雜性。這種漸進(jìn)優(yōu)勢(shì)實(shí)際上只會(huì)在長(zhǎng)上下文中實(shí)現(xiàn)。

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

長(zhǎng)上下文的困難是RNN層本質(zhì)上的問題：與自注意力機(jī)制不同，RNN層必須將上下文壓縮為固定大小的隱藏狀態(tài)。

作為一種壓縮啟發(fā)式，更新規(guī)則需要發(fā)現(xiàn)成千上萬(wàn)甚至數(shù)百萬(wàn)個(gè)token之間的底層結(jié)構(gòu)和關(guān)系。

研究人員首先觀察到，自監(jiān)督學(xué)習(xí)可以將大量訓(xùn)練集壓縮為L(zhǎng)LM等模型的權(quán)重，該模型通常表現(xiàn)出對(duì)其訓(xùn)練數(shù)據(jù)之間語(yǔ)義聯(lián)系的深刻理解，而這，恰恰是他們所需要的。

1. TTT層

受此啟發(fā)，研究人員設(shè)計(jì)了一類新的序列建模層，其中隱藏狀態(tài)是模型，更新規(guī)則是自監(jiān)督學(xué)習(xí)的一個(gè)步驟。

由于更新測(cè)試序列上隱藏狀態(tài)的過程，相當(dāng)于在測(cè)試時(shí)訓(xùn)練模型，因此此類新層稱為測(cè)試時(shí)訓(xùn)練（TTT）層。

研究人員引入兩個(gè)簡(jiǎn)單的實(shí)例：TTT-Linear和TTT-MLP，其中隱藏狀態(tài)分別是線性模型和兩層MLP。TTT層可以集成到任何網(wǎng)絡(luò)架構(gòu)中并進(jìn)行端到端優(yōu)化，類似于RNN層和自注意力。

2. 實(shí)際運(yùn)行時(shí)間

TTT層在FLOP方面已經(jīng)非常高效，研究人員則更進(jìn)一步地提出了兩項(xiàng)創(chuàng)新，使其在實(shí)際運(yùn)行時(shí)間內(nèi)也能保持高效。

首先，與在常規(guī)訓(xùn)練中對(duì)mini-batch序列采取梯度步進(jìn)以實(shí)現(xiàn)更好的并行性類似，他們也在TTT中使用了mini-batch的token。

其次，研究人員為每個(gè)TTT mini-batch內(nèi)的操作開發(fā)了一種對(duì)偶形式，以更好地利用現(xiàn)代GPU和TPU。這種對(duì)偶形式的輸出與原始實(shí)現(xiàn)相當(dāng)，但訓(xùn)練速度卻快了5倍以上。

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

03 Transformer殺手——TTT

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

比如，RNN層——如LSTM、RWKV和Mamba層——將上下文壓縮成一個(gè)固定大小的狀態(tài)，這個(gè)狀態(tài)隨時(shí)間變化。

這種壓縮帶來(lái)了兩種結(jié)果：優(yōu)勢(shì)是處理效率高，因?yàn)槊總€(gè)token的處理時(shí)間是恒定的。劣勢(shì)是在處理長(zhǎng)上下文時(shí)，RNN性能受限于隱藏狀態(tài)的「表達(dá)能力」。

自注意力機(jī)制（Self-attention）也可以從如上角度來(lái)理解。

不同之處在于，它的隱藏狀態(tài)，通常稱為鍵值（KV）緩存是一個(gè)隨t增長(zhǎng)的線性list。

它可以存儲(chǔ)所有的上下文，并且不會(huì)進(jìn)行壓縮，具有很好的表達(dá)能力，不過其處理時(shí)間隨上下文長(zhǎng)度線性增長(zhǎng)。

因此，為了在長(zhǎng)上下文中既保持效率，又具有表達(dá)能力，需要一個(gè)更好的“壓縮啟發(fā)式”（compression heuristic）方法。

具體來(lái)說(shuō)，就需要將數(shù)百萬(wàn)個(gè)token壓縮成一個(gè)能有效捕捉其底層結(jié)構(gòu)和關(guān)系的隱藏狀態(tài)。

1. TTT隱藏狀態(tài)

研究人員的關(guān)鍵思想是，使用自監(jiān)督學(xué)習(xí)來(lái)將歷史上下文x1,…,xt壓縮成一個(gè)隱藏狀態(tài)St。

方法是將上下文視為一個(gè)無(wú)標(biāo)簽數(shù)據(jù)集，而將狀態(tài)視為一個(gè)模型。

具體來(lái)說(shuō)，隱藏狀態(tài)St現(xiàn)在等同于一個(gè)模型f的權(quán)重Wt，這個(gè)模型f可以是線性模型、小型神經(jīng)網(wǎng)絡(luò)或其他任何形式。輸出規(guī)則簡(jiǎn)單地表示為：

直觀來(lái)講，輸出token就是由更新后權(quán)重Wt的模型f對(duì)xt所做的預(yù)測(cè)。更新規(guī)則是在某個(gè)自監(jiān)督損失?上進(jìn)行的一步梯度下降：

其中學(xué)習(xí)率為η。從壓縮的角度來(lái)看，每種啟發(fā)式方法都需要決定記住/忘記哪些輸入。W會(huì)記住那些產(chǎn)生大梯度的輸入——直觀地說(shuō)，就是那些使W學(xué)習(xí)很多的輸入。

?的一種選擇是重構(gòu)xt本身。為了使學(xué)習(xí)問題變得非平凡，作者首先將xt處理成一個(gè)被破壞的輸入x?t，然后優(yōu)化：

類似于去噪自編碼器，f需要發(fā)現(xiàn)xt各維度之間的相關(guān)性，以便從部分信息x?t中重構(gòu)出xt。

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

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

即使在測(cè)試時(shí)，新層仍然為每個(gè)輸入序列訓(xùn)練一個(gè)不同的權(quán)重序列W1,…,WT。

因此，研究人員將其稱之為測(cè)試-時(shí)間訓(xùn)練層（TTT）。

4. 使用TTT層訓(xùn)練神經(jīng)網(wǎng)絡(luò)

TTT層的前向傳播，也有相應(yīng)的后向傳播。

TTT層與RNN層、自注意力機(jī)制有著相同的接口，因此可以在任何更大的神經(jīng)網(wǎng)絡(luò)架構(gòu)中替換它們。

值得一提的是，訓(xùn)練帶有TTT層神經(jīng)網(wǎng)絡(luò)的方式，與訓(xùn)練任何其他Transformer模型相同。

可以使用相同的數(shù)據(jù)、方法和目標(biāo)（如下一個(gè)token預(yù)測(cè)）來(lái)優(yōu)化網(wǎng)絡(luò)其余部分的參數(shù)。

在此，研究人員將訓(xùn)練更大的神經(jīng)網(wǎng)絡(luò)稱為外循環(huán)（outer loop），而在每個(gè)TTT層內(nèi)訓(xùn)練W稱為內(nèi)循環(huán)（inner loop）。

它們之間梯度計(jì)算的區(qū)別是，內(nèi)循環(huán)針對(duì)的是W（即模型f的參數(shù)），外循環(huán)針對(duì)的是網(wǎng)絡(luò)其余部分的參數(shù)θrest。

5. TTT學(xué)習(xí)自監(jiān)督任務(wù)

可以說(shuō)，TTT最重要的部分是自監(jiān)督任務(wù)，因?yàn)樗鼪Q定了W從測(cè)試序列中學(xué)習(xí)的特征類型。

在這個(gè)任務(wù)的設(shè)計(jì)上，研究人員采取了更加端到端的方法——直接優(yōu)化自監(jiān)督任務(wù)以實(shí)現(xiàn)下一個(gè)token預(yù)測(cè)的最終目標(biāo)。

具體來(lái)說(shuō)，研究者將自監(jiān)督任務(wù)的學(xué)習(xí)，作為外循環(huán)的一部分。

從如上公式3中的簡(jiǎn)單重構(gòu)任務(wù)開始，添加了一些外循環(huán)參數(shù)來(lái)讓這個(gè)任務(wù)可學(xué)習(xí)。最新的自監(jiān)督損失是：

在內(nèi)循環(huán)中，只有W被優(yōu)化，因此作為?的參數(shù)寫出；θ們是這個(gè)損失函數(shù)的“超參數(shù)”。在外循環(huán)中，θK,θV,θQ與θrest一起被優(yōu)化，而W僅僅是一個(gè)隱藏狀態(tài)，不是參數(shù)。

圖6用代碼說(shuō)明了這種區(qū)別，其中θK和θQ被實(shí)現(xiàn)為TTT層的參數(shù)，類似于自注意力中的KV參數(shù)。

總的來(lái)說(shuō)，θK,θV,θQ所有可能的選擇構(gòu)成了一系列多視圖重構(gòu)任務(wù)，外循環(huán)可以被理解為從這個(gè)任務(wù)組中選擇一個(gè)具體任務(wù)。為了簡(jiǎn)單起見，研究人員在這里將所有視圖設(shè)計(jì)為線性投影。

6. mini-batch TTT并行化

目前，開發(fā)的原生TTT層在浮點(diǎn)運(yùn)算（FLOP）次數(shù)方面已經(jīng)非常高效。

然而，其更新規(guī)則：

無(wú)法實(shí)現(xiàn)并行化，因?yàn)閃t在兩個(gè)位置上依賴于Wt-1：負(fù)號(hào)和▽l。

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

研究中使用Gt = ▽l(Wt’;xt)，其中t’ = t – mod(t,b)，其中代表著前一個(gè)mini-batch的最后一個(gè)時(shí)間步（或者第一個(gè)mini-batch 0），因此，可以一次并行b個(gè)梯度計(jì)算。

7. 對(duì)偶形式

上面介紹的并行化是必要的，但對(duì)于“實(shí)際運(yùn)行時(shí)間”（wall-clock time）的效率來(lái)說(shuō)還不夠。

然而，現(xiàn)實(shí)中，是無(wú)法對(duì)單個(gè)matmul來(lái)計(jì)算GtS所有的b。相反，需要b個(gè)外積來(lái)對(duì)其進(jìn)行一一計(jì)算。更糟糕的是，對(duì)于每個(gè)

Gt是d×d，這會(huì)比大dXt產(chǎn)生更大的內(nèi)存占用和I/O成本。

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

現(xiàn)在，就可以用上面簡(jiǎn)化的TTT-Linear情況來(lái)演示這些計(jì)算，表示X = [x1, . . . , xb]：

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

表示

和矩陣

可以得出：

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

8. 理論等價(jià)

前面已經(jīng)提到f可以是線性模型，也可以是神經(jīng)網(wǎng)絡(luò)。還有更新規(guī)則的三種變體：online GD、batch GD和mini-batch GD。

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

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

圖10總結(jié)了所有序列建模層的更廣泛范圍內(nèi)TTT層的一般定義。

9. 兩種變體

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

對(duì)于TTT-Linear，

，其中W是平方。對(duì)于TTT-MLP，有兩層，類似于Transfomer的MLP。

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

即，

，其中，可以是或。

04 實(shí)驗(yàn)

通過與兩個(gè)基線Transformer和Mamba（現(xiàn)代RNN）比較，研究人員評(píng)估了TTT-Linear和TTT-MLP。

數(shù)據(jù)集

繼續(xù)Mamba論文之后，研究人員在Pile上執(zhí)行了2k和8k上下文長(zhǎng)度的標(biāo)準(zhǔn)實(shí)驗(yàn)，Pile是一個(gè)用于訓(xùn)練開源LLM的流行文檔數(shù)據(jù)集。

主架構(gòu)

Transformer和Mamba使用不同的，除非另有說(shuō)明，TTT-Linear和TTT-MLP始終使用Mamba架構(gòu)。

1. 短上下文：the Pile

在2k上下文中，TTT-Linear（M）、Mamba和Transformer具有相當(dāng)?shù)男阅?，線條大部分重疊。

TTT-MLP（M）在較大的FLOP預(yù)算下表現(xiàn)稍差。盡管TTT-MLP在每個(gè)模型大小上，都比TTT-Linear具有更好的復(fù)雜度，但FLOP的額外成本抵消了這種優(yōu)勢(shì)。

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

另外，研究人員還觀察到了一個(gè)非常明顯的現(xiàn)象：隨著上下文長(zhǎng)度變長(zhǎng)，TTT層相對(duì)于Mamba的優(yōu)勢(shì)就更大了。

2. 長(zhǎng)上下文：Books

為了評(píng)估長(zhǎng)上下文中的功能，研究人員使用了Pile的一個(gè)流行子集——Books，對(duì)從1k到32k以2個(gè)增量的上下文長(zhǎng)度進(jìn)行了實(shí)驗(yàn)。

根據(jù)上圖，可以觀察到——

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

在32k上下文中，TTT-Linear（M）和TTT-MLP（M）的性能均優(yōu)于Mamba，與Pile 8k的觀察結(jié)果類似。即使具有Transformer架構(gòu)的TTT-MLP（T），在32k上下文中的表現(xiàn)也比Mamba稍好。

在1.3B尺度上，TTT-MLP（T）僅比TTT-MLP（M）稍差。由于缺之清晰的線性擬合，很難推導(dǎo)出經(jīng)驗(yàn)縮放定律。然而，TTT-MLP（T）的強(qiáng)勁趨勢(shì)表明，Transformer架構(gòu)可能更適合超出評(píng)估的更大模型和更長(zhǎng)上下文。

上下文長(zhǎng)度作為超參數(shù)

雖然輸入序列的長(zhǎng)度由用戶確定，但語(yǔ)言模型處理輸入的上下文長(zhǎng)度可以由工程師確定。因此，上下文長(zhǎng)度也是一個(gè)可以選擇的超參數(shù)。

對(duì)于具有線性復(fù)雜度的LLM，研究人員選擇了困惑度中的argmin，因?yàn)槊總€(gè)上下文長(zhǎng)度都有相同的FLOP。

從圖13中，可以觀察到以下結(jié)果——

• 性能最好的方法TTT-Linear和TTT-MLP的線幾乎完全重疊。Mamba和TF Finetune的線在10^20 FLOP后也大部分重疊。
• TF Finetune的性能明顯優(yōu)于TF Pretrain，因?yàn)樗芤嬗陂L(zhǎng)上下文，而不會(huì)在訓(xùn)練FLOP中產(chǎn)生極大的成本。
• 對(duì)于所有從頭開始訓(xùn)練的方法（包括TF預(yù)訓(xùn)練），一旦上下文長(zhǎng)度變得太大，困惑度就會(huì)變得更糟。

從上圖可見，與TTT-Linear相比，TTT-MLP在短上下文中表現(xiàn)稍差，但在長(zhǎng)上下文中表現(xiàn)更好。

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

3. 實(shí)際運(yùn)行時(shí)間

LLM訓(xùn)練和推理可以分解為前向、后向和生成。

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

由于資源限制，這項(xiàng)實(shí)驗(yàn)是用JAX編寫并在TPU上運(yùn)行的。

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

具體來(lái)說(shuō)，研究人員在ThunderKittens中編寫了一個(gè)用于前向的GPU內(nèi)核。從歷史上看，由于并行性和矩陣相乘的使用不當(dāng)，RNN在前向和后向過程中效率低下。

這個(gè)前向內(nèi)核的目標(biāo)，是證明mini-batch TTT和這些問題對(duì)偶形式的有效性。

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

對(duì)于Transformer，每個(gè)token的時(shí)間隨著上下文長(zhǎng)度的增加而線性增長(zhǎng)，但對(duì)于其他方法則大致保持不變。

此外，研究人員在Triton中編寫了另一個(gè)用于生成的GPU內(nèi)核，并在圖15的右圖中對(duì)批大小為512的速度進(jìn)行了基準(zhǔn)測(cè)試。

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

Mamba之后，又看到TTT這么能打的新架構(gòu)誕生，少不了AI社區(qū)的熱議。

有網(wǎng)友稱，這會(huì)不會(huì)是最接近實(shí)時(shí)上下文的方法？很想聽聽大家的想法。這意味著TTT甚至在使用過程中，也能夠?qū)W習(xí)和適應(yīng)，為長(zhǎng)上下文提供更好的性能，而不會(huì)產(chǎn)生通常與Transformer相關(guān)的高昂計(jì)算成本。

OpenAI視頻生成研究人員對(duì)此表示，這項(xiàng)研究看起來(lái)很有趣。

如果scaling law依然存在，TTT將帶來(lái)難以置信的影響。對(duì)于長(zhǎng)序列，Transformer的計(jì)算成本往往很高，當(dāng)長(zhǎng)序列變得更長(zhǎng)時(shí)，RNN會(huì)遺忘。TTT訓(xùn)練巧妙地利用神經(jīng)網(wǎng)絡(luò)解決RNN的不足。

作者介紹

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

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

Yu Sun

Yu Sun是斯坦福大學(xué)計(jì)算機(jī)專業(yè)的博士后，導(dǎo)師是Carlos Guestrin、Tatsu Hashimoto和Sanmi Koyejo。

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

個(gè)人主頁(yè)中，他介紹自己的研究重點(diǎn)是一種名為測(cè)試時(shí)間訓(xùn)練（test-time training）的算法框架。其核心思想是，每個(gè)測(cè)試實(shí)例都定義了自己的學(xué)習(xí)問題，都有自己的泛化目標(biāo)。這通常使用自監(jiān)督學(xué)習(xí)，為每個(gè)實(shí)例即時(shí)訓(xùn)練一個(gè)不同的模型來(lái)實(shí)現(xiàn)的。

在最新研究中，Yu Sun與Xinhao Li在2022年11月共同啟動(dòng)了這一項(xiàng)目。自2023年6月起，Yu Sun專職負(fù)責(zé)該項(xiàng)目。

他提出了項(xiàng)目的概念框架，設(shè)計(jì)了mini-batch TTT和對(duì)偶形式（dual form）。

Xinhao Li

Xinhao Li是UC San Diego研二的學(xué)生，導(dǎo)師是Xiaolong Wang教授。他本人的研究興趣主要是深度學(xué)習(xí)和計(jì)算機(jī)視覺。

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

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

Karan Dalal

Karan Dalal是UC Berkeley電子工程科學(xué)系的本科生。他于2023年6月全職加入該項(xiàng)目，與Xinhao Li合作共同領(lǐng)導(dǎo)了當(dāng)前代碼庫(kù)的開發(fā)工作。

參考資料：

https://x.com/karansdalal/status/1810338845659131940

https://x.com/xiaolonw/status/1810387662060269668

https://arxiv.org/abs/2407.04620