Vision-TTT: Efficient and Expressive Visual Representation Learning with Test-Time Training

TL;DR

This work introduces a new linear-time sequence modeling method Test-Time Training (TTT) into vision and proposes Vision-TTT, which compresses the visual token sequence in a novel self-supervised learning manner and effectively extends vanilla TTT to model 2D visual correlations with global receptive fields.

Abstract

Learning efficient and expressive visual representation has long been the pursuit of computer vision research. While Vision Transformers (ViTs) gradually replace traditional Convolutional Neural Networks (CNNs) as more scalable vision learners, their applications are plagued by the quadratic complexity of the self-attention mechanism. To address the challenge, we introduce a new linear-time sequence modeling method Test-Time Training (TTT) into vision and propose Vision-TTT, which compresses the visual token sequence in a novel self-supervised learning manner. By incorporating bidirectional scan strategy and the Conv2d module, Vision-TTT effectively extends vanilla TTT to model 2D visual correlations with global receptive fields. Extensive experiments show that \texttt{Vittt-T/S/B} achieve 77.3%,81.2%,82.5% Top-1 accuracy on ImageNet classification and also greatly outperform their counterparts on downstream tasks. At 1280x1280 resolution, \texttt{Vittt-T} reduces FLOPs by 79.4% and runs 4.38x faster with 88.9% less memory than DeiT-T. These results demonstrate the expressiveness and efficiency of Vision-TTT as a strong candidate for the next-generation generic visual backbone.

Figures (8)

• Figure 1: Performance and efficiency comparison between DeiT deit, Vim vim and our Vittt model. Results show that Vittt not only achieves superior performance on ImageNet classification and downstream detection and segmentation tasks, but also is more computation and memory efficient in dealing with high-resolution images.
• Figure 2: Basics of TTT. The core idea is to update the hidden state $\mathrm{W}$ with steps of self-supervised gradient descent and then forward the output to the next layer.
• Figure 3: Gradient Magnitude Map (illustrated by Eq. \ref{['ttt_Gt']}). Vision-TTT employs gradients as the explicit indicators to quantify token semantics.
• Figure 4: Overall architecture of Vision-TTT. It consists of three stages: Patchification, Vision-TTT Encoder, and Task Adapters to learn representation for vision tasks.
• Figure 5: Vittt Block. It evolves from vanilla TTT with two 2D designs: the long-term bidirectional scan strategy and the short-term Conv2d token aggregation module.