0
| 本文作者: 陳淑瑜 | 2026-06-02 17:58 | 專題:CVPR 計算機視覺與模式識別會議 |
來源:公眾號“機器之心”
原文鏈接:https://mp.weixin.qq.com/s/yJDuMj9gY6JLQtjng2p_FQ?scene=1&click_id=35
本文第一作者為清華大學博士生韓東辰,主要研究方向包括基礎模型架構、持續學習等。通訊作者為黃高副教授。該論文已入選 CVPR 2026 Best Paper Candidates。
序列建模是大語言模型、計算機視覺等領域的基礎共性問題。當前通用的 Transformer 模型計算復雜度隨序列長度平方增長,在長序列任務中面臨顯著的計算挑戰。因此,研究者們一直在探索具有線性計算復雜度的高效序列建模方法。
測試時訓練(Test-Time Training,TTT)模型是一種新興的序列建模范式,它將注意力操作重新定義為一個在線學習過程,在每次推理時用 Key-Value 對構建一個輕量化的內部模型。不同于一個確定性的模型設計,這一范式開啟了一個豐富的線性復雜度設計空間,有望實現高精度的高效序列建模。
然而,TTT 靈活的設計空間是一把雙刃劍:當前測試時訓練模型缺少系統性的理解和設計原則,如何從豐富的設計空間中構建一個理想的視覺 TTT 模型、如何實現進一步提升,都是目前極具挑戰的問題。
為此,本文在視覺領域系統性地研究了測試時訓練模型的構建和效果。通過一系列的實驗和分析,我們總結了高精度、高效視覺 TTT 模型的六條設計原則,并討論了一些未來改進方向。
基于這些發現,本文構建了純 TTT 架構 Vision Test-Time Training (ViT3)模型,它在各類視覺任務中表超越了現有線性復雜度模型,同時保持并行計算和高推理速度,為未來視覺 TTT 的研究提供了一個有力的基線。

論文鏈接:https://arxiv.org/abs/2512.01643
代碼鏈接:https://github.com/LeapLabTHU/ViTTT
序列建模的新視角
在傳統的理解視角下,我們通常將注意力計算理解為用注意力分數對 value 進行加權求和。實際上,我們可以從一個用 context 構建模型的視角來理解注意力。

圖 1:Softmax 注意力、線性注意力、TTT 模塊示意圖
如圖 1 所示,對于經典的 Softmax 注意力,我們可以將其理解為用當前 context 的 key 和 value 構建了一個兩層 MLP 模型。其中,MLP 的第一個線性層的權重是 key 矩陣的轉置,激活函數是 Softmax,第二個線性層的權重是 value 矩陣。整個注意力計算的過程可以理解為用 query 輸入這樣一個由 key 和 value 構造的兩層 MLP 模型。類似地,線性注意力可以理解為用 key 和 value 構造了一個線性層權重
,進而把 query 輸入這個線性模型,實現線性注意力計算。
因此,兩種經典的注意力范式都可以理解為使用 key 和 value 構造了一個小的模型,進而將 query 輸入這個模型以實現建模。
在這種視角下,兩種注意力的本質區別是是否進行壓縮。Softmax 注意力用完整的 key 和 value 構造兩層 MLP,不進行壓縮,效果很好但是成本高昂。與此相反,線性注意力對 key 和 value 進行了極大地壓縮,直接通過矩陣乘法
將完整的 key 和 value 壓縮到 d×d 空間。因此,線性注意力的空間和計算復雜度都顯著低于不壓縮的 Softmax 注意力。然而,由于 d×d 的空間是完全線性的,并且壓縮方式過于粗糙,線性注意力會導致顯著的性能下降。
于是,一個核心問題是:我們能否在實現壓縮的同時保持模型的性能,同時實現高精度和高效性?
對于這個問題,研究者們回答是肯定的。回顧整個機器學習領域的發展,最為成功的壓縮算法莫過于深度學習。深度學習通過梯度下降,把數據信息壓縮到更緊的模型權重空間中,通常能夠實現十倍甚至百倍的壓縮率。因此,測試時訓練模型正是嘗試在序列建模任務上再次套用深度學習的成功。
如圖 1(c)所示,測試時訓練模型嘗試將 key 和 value 的信息壓縮到一個小的內部模型中,這個內部模型的結構幾乎是任意的。為了實現這一點,TTT 將 key 和 value 視為一個小型的 “數據集”,要求模型在輸入 key 時重建對應的 value,通過梯度下降對內部模型權重進行更新:

其中,?是一個自監督的重建損失函數,比如 L2 損失。在幾次更新之后,我們認為 key 和 value 的信息已經較好地壓縮到了 TTT 內部模型權重之中。這時,我們將更新后的內部模型應用于 query,通過一次簡單的前向傳播,讓 query 從內部模型的權重中讀取 key 和 value 的信息。整個測試時訓練計算過程的計算復雜度正比于內部模型的計算復雜度。
因此,當選取一個簡單的兩層 MLP 等線性復雜度模型作為內部模型時,TTT 具有線性空間和計算復雜度;同時,基于學習的壓縮和非線性狀態空間使得 TTT 具有良好的表達能力。
視覺 TTT 如何設計
如前所述,測試時訓練模型在每個 Block 內部包含一次小的 “深度學習”,具有極高的設計自由度。當前 TTT 缺乏充分的理解和設計原則,阻礙了其在視覺領域的發展。本文在視覺領域系統性地探索了 TTT 的設計空間,聚焦于兩個核心方面:(1)內部訓練設置,如損失函數、學習率、batch size、訓練輪數等(2)內部模型設計,如內部模型架構、寬度、深度等。本文基于 DeiT-S 進行探索,將原始的 Softmax 注意力替換為 TTT 模塊,在 ImageNet-1K 上進行實驗。
1. 內部訓練設置
觀察 1:二階混合偏導數
消失(近似為 0 或等于 0)的損失函數 L 不適合 TTT。

圖 2:TTT 內部訓練損失函數的選擇
TTT 模型的內部訓練所需要的 key、value 等,是在整個外部網絡的端到端訓練過程中優化。在進行外部端到端訓練時,我們將內部梯度視為表達式,讓整個模型實現端到端的反向傳播。在這個過程中,我們需要對內部訓練梯度的表達式再次求導,計算 “梯度的梯度”,這是元學習中的經典操作。在端到端優化過程中,value 的投影矩陣的外部梯度為:

若混合導數
(近似)為 0,則
的外部梯度消失,無法優化
。圖 2 驗證了這一點:MAE(L1)損失的一階導數是符號函數,混合二階導數幾乎處處為 0,效果最差。Smooth L1 的混合二階導數在特定區域為 0,同樣效果不佳。
觀察 2:單次全批量內部訓練(batch gradient)在視覺任務中效果較好。

圖 3:TTT 內部訓練 batch size 和 epoch
如圖 3 所示,B=N 表示使用全部 N 個 key-value 對進行全批量內部梯度下降,而 B=N/2、B=N/3、B=N/4 則對應小批次梯度下降,即把數據集劃分為 2、3、4 個連續的小批次,每個 epoch 的內部訓練執行 2、3、4 次參數更新。
本文認為,連續小批次梯度下降會引入顯式的因果關系:(1) 前序小批次的更新會改變內模型權重,影響后續批次的梯度計算;(2) 后續批次的更新可能覆蓋前序批次產生的參數更新。這種因果關系非常適合自然語言這類因果數據,但在視覺任務中通常是不必要的。
多個 epoch 的內部訓練能夠提升性能,但是會顯著降低模型推理速度。
觀察 3:相對較大的內部學習率(如 1.0)更適合 TTT。

圖 4:TTT 內部訓練學習率
在圖 4 中,我們驗證了不同內部學習率的影響。過小的學習率不足以快速更新原始的內部模型權重,過大的學習率會導致訓練不穩定。在本文驗證的視覺任務中,內部學習率 1.0 較為合適。
值得注意的是,在一些特殊情況下,內部學習率可以吸收為 key 和 value 的放縮。但是這并不內說明內部學習率的設置不重要。一個類似的例子是 Softmax 注意力中的
放縮,它也能夠被 query 和 key 吸收,但是依然重要。
2. 內部模型設計

圖 5:TTT 內部模型結構。r 和 l 分別表示模型寬度比例和模型深度
觀察 4:擴大內部模型容量能夠持續提升 TTT 的序列建模能力。
如圖 5 第一部分所示,我們將 TTT 內部模型實現為一個 SiLU 激活函數的兩層 MLP,并逐漸將寬度比例(隱藏維度 / 輸入維度)從 1.0 增加到 4.0。我們可以觀察到,模型的性能隨著 TTT 內部模型的容量持續提升。這是測試時訓練范式的一個重要優勢,即在外部模型尺寸固定的條件下,我們可以通過簡單地擴展內部模型尺寸來實現更好的序列建模效果。
但是,值得指出的是,在推理時,一個內部模型引入的計算量大于相同尺寸的外部模塊。在推理階段,一個外部模塊僅需要一次前向傳播;而一個內部模型則至少需要進行(1)key 前向傳播,1 倍前向傳播計算量(2)內部損失反向傳播,約 2 倍前向傳播計算量(3)將更新后的內部模型應用于 query,1 倍前向傳播計算量。一個內部模型的計算量是同尺寸外部模塊的近 4 倍。因此,盡管簡單擴展內部模型能顯著提升模型能力,本文認為設計輕量的高效內部模型仍然是一個重要的研究問題。
觀察 5:在當前的 TTT 實現下,較深的內部模型存在優化困難。
在上一部分中,我們擴展了內部模型的寬度。另一個常見的方案是擴展模型的深度。如圖 5 第二部分所示,我們分別將內部模型實現為單層線性層、兩層 MLP、三層 MLP,以探索內部模型深度對性能的影響。結果表明,擴展內部模型深度時,模型性能持續下降。在一定程度上,這是一個不符合預期的結果。因此理論上更深的內部模型具有更大的容量,應該能夠實現更好的序列建模。

圖 6:較深的 TTT 內部模型的訓練、測試 loss 均更高
本文認為,優化問題是導致較深內部模型性能不佳的原因。為了驗證這一點,我們在圖 6 中展示了上述三個模型的訓練和測試 loss。可以看到,采用更深內部模型時,模型的訓練 loss 更高,這說明性能衰退不是因為過擬合,而恰恰是欠擬合。更深的內部模型具有更強的擬合能力,而在真實訓練中相比于淺內部模型欠擬合,說明深內部模型存在優化問題,遠遠沒有達到其理論表征能力。在圖 5 第二部分的結果進一步驗證了這一觀點:將兩層 MLP、SwiGLU 的輸出層強行限制為 identity 能夠讓模型性能顯著提升,說明較深的內部模型顯著欠優化。
這種優化困難與 TTT 的元學習屬性有關。一方面,內部模型的初始權重是外部訓練過程中的可學參數,較深的內部模型可能導致這個初始參數難以優化。另一方面,即便有較好的初始參數,深內部模型也可能在內部訓練過程中遇到梯度消失和爆炸問題,影響效果。
解決較深內部模型的優化問題,是 TTT 的一個重要未來方向。理論工作表明,神經網絡的擬合能力隨深度指數增長,這正是當前神經網絡成功的核心原因。因此,較深的內部模型在實現高精度測試時訓練序列建模中具有突出潛力。
觀察 6:在視覺任務中,內部模型非常適合采用卷積架構。
在 Transformer 出現之前,卷積操作一直是計算機視覺領域的基石。在 TTT 中,我們完全可以將內部模型實現為一個小的卷積網絡,而不需要將其限定為 MLP、SwiGLU 等單 token 組件。在圖 5 第三部分中,我們展示了兩種最簡單的實現,將內部模型實現為 3x3 卷積或 Depthwise 卷積。可以看到,這兩種實現均取得了顯著的性能增益。
本文認為,采用卷積網絡作為 TTT 內部模型,能夠實現局部信息和全局信息的完美融合。具體而言,在 TTT 的內部訓練過程中,key 和 value 的全局信息被壓縮到卷積核的權重中。在 query 通過內部模型前向傳播時,它很自然地同時得到局部信息(卷積核感受野)和全局信息(包含全局信息的卷積核權重),因此完美契合視覺任務。
ViT3:視覺測試時訓練模型
基于以上探索,本文提出純 TTT 架構 Vision Test-Time Training (ViT3)模型。具體而言,對于內部訓練,采用一次全批量梯度下降、學習率 1.0、點積損失函數;對于內部模型,采用簡化的 SwiGLU 和 Depthwise 卷積。作為 TTT 架構的一個 baseline,ViT3 可以應用于各類視覺 Transformer 框架,完成各類任務。
本文在分類、檢測、分割、生成等代表性視覺任務上驗證了 ViT3 的性能。實驗結果表明,ViT3 能夠超越各類先進的線性復雜度設計,如線性注意力、視覺 Mamba 模型。這一結果驗證了 ViT3 設計的有效性,同時凸顯了測試時訓練模型在高效、可擴展的線性復雜度視覺序列建模領域的突出潛力。
以下是部分實驗結果,詳細結果請見論文原文:

圖 7:ImagNet-1K 圖片分類實驗

圖 8:高分辨率下游任務 —— 物體檢測

圖 9:圖片生成實驗

圖 10:推理 Throughput 與顯存隨輸入分辨率變化
總結與展望
設計高效、高表達能力的序列建模范式一直是計算機視覺領域的核心課題。
本文在視覺領域系統性地探索了測試時訓練(Test-Time Training,TTT)這一新興序列建模范式的設計空間。
通過大量實驗,本文將觀察與分析總結為六條設計見解,為構建高效的視覺 TTT 模型明確了設計準則,并討論了一些未來的研究方向。
基于上述研究,本文提出了視覺測試時訓練模型(ViT3),一種面向視覺序列建模的純 TTT 架構。ViT3 在多項視覺任務上取得了具有競爭力的結果,為線性復雜度模型提供一個有力的基線。
我們期望本文的研究結論與分析能夠幫助視覺 TTT 模型的后續研究工作。
投稿或尋求報道:liyazhou@jiqizhixin.com
本專題其他文章