清華+阿里發(fā)布ViT3：解鎖視覺(jué)TTT，突破Transformer復(fù)雜度瓶頸

本文第一作者為清華大學(xué)博士生韓東辰，主要研究方向包括基礎(chǔ)模型架構(gòu)、持續(xù)學(xué)習(xí)等。通訊作者為黃高副教授。

序列建模是大語(yǔ)言模型、計(jì)算機(jī)視覺(jué)等領(lǐng)域的基礎(chǔ)共性問(wèn)題。當(dāng)前通用的 Transformer 模型計(jì)算復(fù)雜度隨序列長(zhǎng)度平方增長(zhǎng)，在長(zhǎng)序列任務(wù)中面臨顯著的計(jì)算挑戰(zhàn)。因此，研究者們一直在探索具有線性計(jì)算復(fù)雜度的高效序列建模方法。

測(cè)試時(shí)訓(xùn)練（Test-Time Training，TTT）模型是一種新興的序列建模范式，它將注意力操作重新定義為一個(gè)在線學(xué)習(xí)過(guò)程，在每次推理時(shí)用 Key-Value 對(duì)構(gòu)建一個(gè)輕量化的內(nèi)部模型。不同于一個(gè)確定性的模型設(shè)計(jì)，這一范式開(kāi)啟了一個(gè)豐富的線性復(fù)雜度設(shè)計(jì)空間，有望實(shí)現(xiàn)高精度的高效序列建模。

然而，TTT 靈活的設(shè)計(jì)空間是一把雙刃劍：當(dāng)前測(cè)試時(shí)訓(xùn)練模型缺少系統(tǒng)性的理解和設(shè)計(jì)原則，如何從豐富的設(shè)計(jì)空間中構(gòu)建一個(gè)理想的視覺(jué) TTT 模型、如何實(shí)現(xiàn)進(jìn)一步提升，都是目前極具挑戰(zhàn)的問(wèn)題。

為此，本文在視覺(jué)領(lǐng)域系統(tǒng)性地研究了測(cè)試時(shí)訓(xùn)練模型的構(gòu)建和效果。通過(guò)一系列的實(shí)驗(yàn)和分析，我們總結(jié)了高精度、高效視覺(jué) TTT 模型的六條設(shè)計(jì)原則，并討論了一些未來(lái)改進(jìn)方向。

基于這些發(fā)現(xiàn)，本文構(gòu)建了純 TTT 架構(gòu) Vision Test-Time Training （ViT3）模型，它在各類(lèi)視覺(jué)任務(wù)中表超越了現(xiàn)有線性復(fù)雜度模型，同時(shí)保持并行計(jì)算和高推理速度，為未來(lái)視覺(jué) TTT 的研究提供了一個(gè)有力的基線。

• 論文鏈接：https://arxiv.org/abs/2512.01643
• 代碼鏈接：https://github.com/LeapLabTHU/ViTTT

序列建模的新視角

在傳統(tǒng)的理解視角下，我們通常將注意力計(jì)算理解為用注意力分?jǐn)?shù)對(duì) value 進(jìn)行加權(quán)求和。實(shí)際上，我們可以從一個(gè)用 context 構(gòu)建模型的視角來(lái)理解注意力。

圖 1：Softmax 注意力、線性注意力、TTT 模塊示意圖

因此，兩種經(jīng)典的注意力范式都可以理解為使用 key 和 value 構(gòu)造了一個(gè)小的模型，進(jìn)而將 query 輸入這個(gè)模型以實(shí)現(xiàn)建模。

于是，一個(gè)核心問(wèn)題是：我們能否在實(shí)現(xiàn)壓縮的同時(shí)保持模型的性能，同時(shí)實(shí)現(xiàn)高精度和高效性？

對(duì)于這個(gè)問(wèn)題，研究者們回答是肯定的。回顧整個(gè)機(jī)器學(xué)習(xí)領(lǐng)域的發(fā)展，最為成功的壓縮算法莫過(guò)于深度學(xué)習(xí)。深度學(xué)習(xí)通過(guò)梯度下降，把數(shù)據(jù)信息壓縮到更緊的模型權(quán)重空間中，通常能夠?qū)崿F(xiàn)十倍甚至百倍的壓縮率。因此，測(cè)試時(shí)訓(xùn)練模型正是嘗試在序列建模任務(wù)上再次套用深度學(xué)習(xí)的成功。

如圖 1（c）所示，測(cè)試時(shí)訓(xùn)練模型嘗試將 key 和 value 的信息壓縮到一個(gè)小的內(nèi)部模型中，這個(gè)內(nèi)部模型的結(jié)構(gòu)幾乎是任意的。為了實(shí)現(xiàn)這一點(diǎn)，TTT 將 key 和 value 視為一個(gè)小型的 “數(shù)據(jù)集”，要求模型在輸入 key 時(shí)重建對(duì)應(yīng)的 value，通過(guò)梯度下降對(duì)內(nèi)部模型權(quán)重進(jìn)行更新：

其中，?是一個(gè)自監(jiān)督的重建損失函數(shù)，比如 L2 損失。在幾次更新之后，我們認(rèn)為 key 和 value 的信息已經(jīng)較好地壓縮到了 TTT 內(nèi)部模型權(quán)重之中。這時(shí)，我們將更新后的內(nèi)部模型應(yīng)用于 query，通過(guò)一次簡(jiǎn)單的前向傳播，讓 query 從內(nèi)部模型的權(quán)重中讀取 key 和 value 的信息。整個(gè)測(cè)試時(shí)訓(xùn)練計(jì)算過(guò)程的計(jì)算復(fù)雜度正比于內(nèi)部模型的計(jì)算復(fù)雜度。

因此，當(dāng)選取一個(gè)簡(jiǎn)單的兩層 MLP 等線性復(fù)雜度模型作為內(nèi)部模型時(shí)，TTT 具有線性空間和計(jì)算復(fù)雜度；同時(shí)，基于學(xué)習(xí)的壓縮和非線性狀態(tài)空間使得 TTT 具有良好的表達(dá)能力。

視覺(jué) TTT 如何設(shè)計(jì)

如前所述，測(cè)試時(shí)訓(xùn)練模型在每個(gè) Block 內(nèi)部包含一次小的 “深度學(xué)習(xí)”，具有極高的設(shè)計(jì)自由度。當(dāng)前 TTT 缺乏充分的理解和設(shè)計(jì)原則，阻礙了其在視覺(jué)領(lǐng)域的發(fā)展。本文在視覺(jué)領(lǐng)域系統(tǒng)性地探索了 TTT 的設(shè)計(jì)空間，聚焦于兩個(gè)核心方面：（1）內(nèi)部訓(xùn)練設(shè)置，如損失函數(shù)、學(xué)習(xí)率、batch size、訓(xùn)練輪數(shù)等（2）內(nèi)部模型設(shè)計(jì)，如內(nèi)部模型架構(gòu)、寬度、深度等。本文基于 DeiT-S 進(jìn)行探索，將原始的 Softmax 注意力替換為 TTT 模塊，在 ImageNet-1K 上進(jìn)行實(shí)驗(yàn)。

1. 內(nèi)部訓(xùn)練設(shè)置

圖 2：TTT 內(nèi)部訓(xùn)練損失函數(shù)的選擇

TTT 模型的內(nèi)部訓(xùn)練所需要的 key、value 等，是在整個(gè)外部網(wǎng)絡(luò)的端到端訓(xùn)練過(guò)程中優(yōu)化。在進(jìn)行外部端到端訓(xùn)練時(shí)，我們將內(nèi)部梯度視為表達(dá)式，讓整個(gè)模型實(shí)現(xiàn)端到端的反向傳播。在這個(gè)過(guò)程中，我們需要對(duì)內(nèi)部訓(xùn)練梯度的表達(dá)式再次求導(dǎo)，計(jì)算 “梯度的梯度”，這是元學(xué)習(xí)中的經(jīng)典操作。在端到端優(yōu)化過(guò)程中，value 的投影矩陣的外部梯度為：

觀察 2：?jiǎn)未稳績(jī)?nèi)部訓(xùn)練（batch gradient）在視覺(jué)任務(wù)中效果較好。

圖 3：TTT 內(nèi)部訓(xùn)練 batch size 和 epoch

如圖 3 所示，B=N 表示使用全部 N 個(gè) key-value 對(duì)進(jìn)行全批量?jī)?nèi)部梯度下降，而 B=N/2、B=N/3、B=N/4 則對(duì)應(yīng)小批次梯度下降，即把數(shù)據(jù)集劃分為 2、3、4 個(gè)連續(xù)的小批次，每個(gè) epoch 的內(nèi)部訓(xùn)練執(zhí)行 2、3、4 次參數(shù)更新。

本文認(rèn)為，連續(xù)小批次梯度下降會(huì)引入顯式的因果關(guān)系：(1) 前序小批次的更新會(huì)改變內(nèi)模型權(quán)重，影響后續(xù)批次的梯度計(jì)算；(2) 后續(xù)批次的更新可能覆蓋前序批次產(chǎn)生的參數(shù)更新。這種因果關(guān)系非常適合自然語(yǔ)言這類(lèi)因果數(shù)據(jù)，但在視覺(jué)任務(wù)中通常是不必要的。

多個(gè) epoch 的內(nèi)部訓(xùn)練能夠提升性能，但是會(huì)顯著降低模型推理速度。

觀察 3：相對(duì)較大的內(nèi)部學(xué)習(xí)率（如 1.0）更適合 TTT。

圖 4：TTT 內(nèi)部訓(xùn)練學(xué)習(xí)率

在圖 4 中，我們驗(yàn)證了不同內(nèi)部學(xué)習(xí)率的影響。過(guò)小的學(xué)習(xí)率不足以快速更新原始的內(nèi)部模型權(quán)重，過(guò)大的學(xué)習(xí)率會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定。在本文驗(yàn)證的視覺(jué)任務(wù)中，內(nèi)部學(xué)習(xí)率 1.0 較為合適。

2. 內(nèi)部模型設(shè)計(jì)

圖 5：TTT 內(nèi)部模型結(jié)構(gòu)。r 和 l 分別表示模型寬度比例和模型深度

觀察 4：擴(kuò)大內(nèi)部模型容量能夠持續(xù)提升 TTT 的序列建模能力。

如圖 5 第一部分所示，我們將 TTT 內(nèi)部模型實(shí)現(xiàn)為一個(gè) SiLU 激活函數(shù)的兩層 MLP，并逐漸將寬度比例（隱藏維度 / 輸入維度）從 1.0 增加到 4.0。我們可以觀察到，模型的性能隨著 TTT 內(nèi)部模型的容量持續(xù)提升。這是測(cè)試時(shí)訓(xùn)練范式的一個(gè)重要優(yōu)勢(shì)，即在外部模型尺寸固定的條件下，我們可以通過(guò)簡(jiǎn)單地?cái)U(kuò)展內(nèi)部模型尺寸來(lái)實(shí)現(xiàn)更好的序列建模效果。

但是，值得指出的是，在推理時(shí)，一個(gè)內(nèi)部模型引入的計(jì)算量大于相同尺寸的外部模塊。在推理階段，一個(gè)外部模塊僅需要一次前向傳播；而一個(gè)內(nèi)部模型則至少需要進(jìn)行（1）key 前向傳播，1 倍前向傳播計(jì)算量（2）內(nèi)部損失反向傳播，約 2 倍前向傳播計(jì)算量（3）將更新后的內(nèi)部模型應(yīng)用于 query，1 倍前向傳播計(jì)算量。一個(gè)內(nèi)部模型的計(jì)算量是同尺寸外部模塊的近 4 倍。因此，盡管簡(jiǎn)單擴(kuò)展內(nèi)部模型能顯著提升模型能力，本文認(rèn)為設(shè)計(jì)輕量的高效內(nèi)部模型仍然是一個(gè)重要的研究問(wèn)題。

觀察 5：在當(dāng)前的 TTT 實(shí)現(xiàn)下，較深的內(nèi)部模型存在優(yōu)化困難。

在上一部分中，我們擴(kuò)展了內(nèi)部模型的寬度。另一個(gè)常見(jiàn)的方案是擴(kuò)展模型的深度。如圖 5 第二部分所示，我們分別將內(nèi)部模型實(shí)現(xiàn)為單層線性層、兩層 MLP、三層 MLP，以探索內(nèi)部模型深度對(duì)性能的影響。結(jié)果表明，擴(kuò)展內(nèi)部模型深度時(shí)，模型性能持續(xù)下降。在一定程度上，這是一個(gè)不符合預(yù)期的結(jié)果。因此理論上更深的內(nèi)部模型具有更大的容量，應(yīng)該能夠?qū)崿F(xiàn)更好的序列建模。

圖 6：較深的 TTT 內(nèi)部模型的訓(xùn)練、測(cè)試 loss 均更高

本文認(rèn)為，優(yōu)化問(wèn)題是導(dǎo)致較深內(nèi)部模型性能不佳的原因。為了驗(yàn)證這一點(diǎn)，我們?cè)趫D 6 中展示了上述三個(gè)模型的訓(xùn)練和測(cè)試 loss。可以看到，采用更深內(nèi)部模型時(shí)，模型的訓(xùn)練 loss 更高，這說(shuō)明性能衰退不是因?yàn)檫^(guò)擬合，而恰恰是欠擬合。更深的內(nèi)部模型具有更強(qiáng)的擬合能力，而在真實(shí)訓(xùn)練中相比于淺內(nèi)部模型欠擬合，說(shuō)明深內(nèi)部模型存在優(yōu)化問(wèn)題，遠(yuǎn)遠(yuǎn)沒(méi)有達(dá)到其理論表征能力。在圖 5 第二部分的結(jié)果進(jìn)一步驗(yàn)證了這一觀點(diǎn)：將兩層 MLP、SwiGLU 的輸出層強(qiáng)行限制為 identity 能夠讓模型性能顯著提升，說(shuō)明較深的內(nèi)部模型顯著欠優(yōu)化。

這種優(yōu)化困難與 TTT 的元學(xué)習(xí)屬性有關(guān)。一方面，內(nèi)部模型的初始權(quán)重是外部訓(xùn)練過(guò)程中的可學(xué)參數(shù)，較深的內(nèi)部模型可能導(dǎo)致這個(gè)初始參數(shù)難以優(yōu)化。另一方面，即便有較好的初始參數(shù)，深內(nèi)部模型也可能在內(nèi)部訓(xùn)練過(guò)程中遇到梯度消失和爆炸問(wèn)題，影響效果。

解決較深內(nèi)部模型的優(yōu)化問(wèn)題，是 TTT 的一個(gè)重要未來(lái)方向。理論工作表明，神經(jīng)網(wǎng)絡(luò)的擬合能力隨深度指數(shù)增長(zhǎng)，這正是當(dāng)前神經(jīng)網(wǎng)絡(luò)成功的核心原因。因此，較深的內(nèi)部模型在實(shí)現(xiàn)高精度測(cè)試時(shí)訓(xùn)練序列建模中具有突出潛力。

觀察 6：在視覺(jué)任務(wù)中，內(nèi)部模型非常適合采用卷積架構(gòu)。

在 Transformer 出現(xiàn)之前，卷積操作一直是計(jì)算機(jī)視覺(jué)領(lǐng)域的基石。在 TTT 中，我們完全可以將內(nèi)部模型實(shí)現(xiàn)為一個(gè)小的卷積網(wǎng)絡(luò)，而不需要將其限定為 MLP、SwiGLU 等單 token 組件。在圖 5 第三部分中，我們展示了兩種最簡(jiǎn)單的實(shí)現(xiàn)，將內(nèi)部模型實(shí)現(xiàn)為 3x3 卷積或 Depthwise 卷積。可以看到，這兩種實(shí)現(xiàn)均取得了顯著的性能增益。

本文認(rèn)為，采用卷積網(wǎng)絡(luò)作為 TTT 內(nèi)部模型，能夠?qū)崿F(xiàn)局部信息和全局信息的完美融合。具體而言，在 TTT 的內(nèi)部訓(xùn)練過(guò)程中，key 和 value 的全局信息被壓縮到卷積核的權(quán)重中。在 query 通過(guò)內(nèi)部模型前向傳播時(shí)，它很自然地同時(shí)得到局部信息（卷積核感受野）和全局信息（包含全局信息的卷積核權(quán)重），因此完美契合視覺(jué)任務(wù)。

ViT3：視覺(jué)測(cè)試時(shí)訓(xùn)練模型

基于以上探索，本文提出純 TTT 架構(gòu) Vision Test-Time Training （ViT3）模型。具體而言，對(duì)于內(nèi)部訓(xùn)練，采用一次全批量梯度下降、學(xué)習(xí)率 1.0、點(diǎn)積損失函數(shù)；對(duì)于內(nèi)部模型，采用簡(jiǎn)化的 SwiGLU 和 Depthwise 卷積。作為 TTT 架構(gòu)的一個(gè) baseline，ViT3 可以應(yīng)用于各類(lèi)視覺(jué) Transformer 框架，完成各類(lèi)任務(wù)。

本文在分類(lèi)、檢測(cè)、分割、生成等代表性視覺(jué)任務(wù)上驗(yàn)證了 ViT3 的性能。實(shí)驗(yàn)結(jié)果表明，ViT3 能夠超越各類(lèi)先進(jìn)的線性復(fù)雜度設(shè)計(jì)，如線性注意力、視覺(jué) Mamba 模型。這一結(jié)果驗(yàn)證了 ViT3 設(shè)計(jì)的有效性，同時(shí)凸顯了測(cè)試時(shí)訓(xùn)練模型在高效、可擴(kuò)展的線性復(fù)雜度視覺(jué)序列建模領(lǐng)域的突出潛力。

以下是部分實(shí)驗(yàn)結(jié)果，詳細(xì)結(jié)果請(qǐng)見(jiàn)論文原文：

圖 7：ImagNet-1K 圖片分類(lèi)實(shí)驗(yàn)

圖 8：高分辨率下游任務(wù) —— 物體檢測(cè)

圖 9：圖片生成實(shí)驗(yàn)

圖 10：推理 Throughput 與顯存隨輸入分辨率變化

總結(jié)與展望

設(shè)計(jì)高效、高表達(dá)能力的序列建模范式一直是計(jì)算機(jī)視覺(jué)領(lǐng)域的核心課題。

本文在視覺(jué)領(lǐng)域系統(tǒng)性地探索了測(cè)試時(shí)訓(xùn)練（Test-Time Training，TTT）這一新興序列建模范式的設(shè)計(jì)空間。

通過(guò)大量實(shí)驗(yàn)，本文將觀察與分析總結(jié)為六條設(shè)計(jì)見(jiàn)解，為構(gòu)建高效的視覺(jué) TTT 模型明確了設(shè)計(jì)準(zhǔn)則，并討論了一些未來(lái)的研究方向。

基于上述研究，本文提出了視覺(jué)測(cè)試時(shí)訓(xùn)練模型（ViT3），一種面向視覺(jué)序列建模的純 TTT 架構(gòu)。ViT3 在多項(xiàng)視覺(jué)任務(wù)上取得了具有競(jìng)爭(zhēng)力的結(jié)果，為線性復(fù)雜度模型提供一個(gè)有力的基線。

我們期望本文的研究結(jié)論與分析能夠幫助視覺(jué) TTT 模型的后續(xù)研究工作。