Source-linked AI summary
Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, Shuicheng Yan
TL;DR
ViT’s simple tokenization and redundant backbone limit local-structure modeling and feature richness when training from scratch on ImageNet. T2T-ViT progressively aggregates neighboring tokens and uses a deep-narrow backbone, achieving stronger reported performance with lower complexity than vanilla ViT and favorable CNN comparisons.
Problem
ViT performs worse than similar-sized CNNs on ImageNet from scratch because fixed tokenization misses local structure and its backbone is redundant.
Method
T2T-ViT progressively restructures and aggregates neighboring tokens, then uses a deep-narrow vision-transformer backbone.
Results
81.5% top-1 accuracy is achieved with 21.5M parameters and 4.8G MACs, versus 78.1% for ViT with 48.6M parameters and 10.1G MACs on ImageNet.
Takeaways & Limitations
T2T-ViT outperforms ResNets and matches MobileNets at similar model sizes when trained from scratch on ImageNet.
Abstract
from arXiv · showhide
Transformers, which are popular for language modeling, have been explored for solving vision tasks recently, e.g., the Vision Transformer (ViT) for image classification. The ViT model splits each image into a sequence of tokens with fixed length and then applies multiple Transformer layers to model their global relation for classification. However, ViT achieves inferior performance to CNNs when trained from scratch on a midsize dataset like ImageNet. We find it is because: 1) the simple tokenization of input images fails to model the important local structure such as edges and lines among neighboring pixels, leading to low training sample efficiency; 2) the redundant attention backbone design of ViT leads to limited feature richness for fixed computation budgets and limited training samples. To overcome such limitations, we propose a new Tokens-To-Token Vision Transformer (T2T-ViT), which incorporates 1) a layer-wise Tokens-to-Token (T2T) transformation to progressively structurize the image to tokens by recursively aggregating neighboring Tokens into one Token (Tokens-to-Token), such that local structure represented by surrounding tokens can be modeled and tokens length can be reduced; 2) an efficient backbone with a deep-narrow structure for vision transformer motivated by CNN architecture design after empirical study. Notably, T2T-ViT reduces the parameter count and MACs of vanilla ViT by half, while achieving more than 3.0\% improvement when trained from scratch on ImageNet. It also outperforms ResNets and achieves comparable performance with MobileNets by directly training on ImageNet. For example, T2T-ViT with comparable size to ResNet50 (21.5M parameters) can achieve 83.3\% top1 accuracy in image resolution 384$\times$384 on ImageNet. (Code: https://github.com/yitu-opensource/T2T-ViT)
1. Introduction
ViT’s fixed patch tokenization and inefficient attention backbone limit local-structure modeling and feature richness on ImageNet when trained from scratch. T2T-ViT addresses these issues with progressive tokenization and a deep-narrow backbone, improving efficiency and performance.
- Motivation: ViT underperforms similar-sized CNNs on ImageNet when trained from scratch because its tokenization misses local structure and its backbone contains redundancy.The paper links these limitations to greater data requirements, limited feature richness, and training difficulty.
- Motivation: ResNet features progressively capture edges, lines, and textures, whereas ViT poorly models structure and contains channels with zero values.The comparison uses feature visualizations of ResNet50 and ViT-L/16 trained on ImageNet.
- Proposed approach: T2T progressively aggregates neighboring tokens after reconstructing them as an image and applying overlapping soft splits.This process embeds surrounding-patch structure into tokens while iteratively reducing token length.
- Proposed approach: Deep-narrow architecture engineering reduces parameter count and MACs while preserving performance and improving feature richness for ViT.Among the explored CNN-inspired designs, deep-narrow structure performed best in the reported experiments.
- Results: 81.5% top-1 accuracy is achieved by a 21.5M-parameter, 4.8G-MAC T2T-ViT, versus 78.1% for ViT with 48.6M parameters and 10.1G MACs.The T2T-ViT result also exceeds the reported 76%-79% range for similarly sized ResNet50, while lite variants are comparable to MobileNets.
- Results: T2T-ViT is presented as outperforming CNNs at different complexities on ImageNet without JFT-300M pretraining.The contribution claim covers comparisons with CNNs under direct ImageNet training.
2. Related Work
Vision Transformers apply self-attention to model global dependencies in vision, while related CNN work explores attention mechanisms and local or patchwise alternatives. T2T extends token processing by repeatedly restructuring and aggregating neighboring tokens.
- Vision Transformers: Vision Transformers rely entirely on self-attention to draw global dependencies between inputs and outputs.A transformer layer typically combines multi-head self-attention, an MLP block, layer normalization, and residual connections.
- Self-attention in CNNs: Self-attention has been incorporated into CNNs for channel attention, long-range dependencies, local patches, and patchwise operations related to convolution.The cited CNN approaches span both global and localized attention mechanisms.
- T2T process: The T2T process restructures tokens as an image, applies overlapping splitting, and concatenates neighboring tokens into shorter token sequences.The illustrated example concatenates four input tokens into one output token.
3. Tokens-to-Token ViT
T2T-ViT combines progressive Tokens-to-Token tokenization with a deep-narrow transformer backbone. The design models local structure, progressively reduces token length, and improves efficiency through smaller channel dimensions and greater depth.
- Tokens-to-Token module: T2T-ViT uses a layer-wise T2T module to model local image structure and progressively reduce token length.The module repeatedly performs re-structurization and soft split operations.
- Tokens-to-Token module: Each T2T process applies self-attention and an MLP, reshapes tokens into an image, then performs overlapping soft splits.Reshape reorganizes tokens into spatial dimensions, while overlapping patches aggregate surrounding local information.
- Tokens-to-Token module: Overlapping soft splits concatenate each k × k × c patch into one token, correlating neighboring patches and producing tokens for the next T2T process.The output token length is reduced iteratively as the process repeats.
- T2T-ViT backbone: The backbone adopts a deep-narrow structure with smaller channel and hidden dimensions but more layers to reduce redundancy and improve feature richness.The design uses small channel number and hidden dimension d with more layers b.
- Tokens-to-Token module: The T2T module uses three soft splits and two re-structurizations, reducing a 224 × 224 input to 14 × 14 tokens.The soft-split patch sizes are P = [7, 3, 3], with overlaps S = [3, 1, 1].
- Model variants: The architecture includes T2T-ViT models sized comparably to ResNet50/101/152 and lite models sized comparably to MobileNetV1/V2.The lite models reduce layer depth, hidden dimension, and MLP ratio without special efficient-convolution designs.
- T2T-ViT backbone: T2T-ViT-14 uses 14 backbone transformer layers and 384 hidden dimensions, whereas ViT-B/16 uses 12 layers and 768 hidden dimensions.The paper states that ViT-B/16 is 3x larger than T2T-ViT-14 in parameters and MACs.
4. Experiments
Experiments show that T2T-ViT improves ImageNet performance over comparable ViT and ResNet models while using an efficient architecture. Ablations attribute gains to progressive tokenization and a deep-narrow backbone, while other CNN-inspired structures have mixed effects.
- Experimental setup: T2T-ViT is evaluated against comparable CNNs and ViTs on ImageNet, with additional studies of backbone designs and component ablations.The experiments include training from scratch, architecture comparisons, and ablations of the T2T module and deep-narrow structure.
- T2T-ViT vs. ViT: 81.5% top-1 accuracy is achieved by T2T-ViT-14 with 21.5M parameters and 4.8G MACs, versus ViT-S/16 at 78.1% with 48.6M parameters and 10.1G MACs.The comparison is based on training from scratch on ImageNet.
- T2T-ViT vs. ResNet: 1.4%-2.7% performance gains over comparable ResNets are reported for T2T-ViT models on ImageNet.For example, T2T-ViT-14 uses 21.5M parameters and 4.8G MACs and achieves 81.5% accuracy, compared with ResNet50 at 25.5M parameters and 4.3G MACs.
- T2T-ViT vs. MobileNets: 76.5% top-1 accuracy is achieved by the 6.9M-parameter T2T-ViT-12, exceeding MobileNetV2 1.4x by 0.9%.The T2T-ViT lite models have comparable model sizes but larger MACs because of dense Transformer operations.
- CNN-inspired designs: SE blocks improve both ViT and T2T-ViT, whereas dense connections can hurt performance and many attention heads have few effects while increasing GPU memory.Ghost operation reduces parameters and MACs but decreases accuracy by 2.0% on T2T-ViT.
- Ablation study: 2.0%-2.2% ImageNet improvement comes from the T2T module at similar model size and MACs, while replacing it with convolution layers is 0.5%-1.0% worse.The T2T module is reported to model both global relations and image structure, whereas early convolution layers model structure information.
- Ablation study: 2.7% top-1 accuracy is lost when the deep-narrow backbone is changed to a shallow-wide structure with similar model size and MACs.The authors also report that deep-narrow design reduces parameters and MACs with nearly no performance degradation.
5. Conclusion
The paper proposes T2T-ViT for training vision Transformers from scratch on ImageNet. It combines progressive token aggregation with a deep-narrow backbone and reports performance superior to ResNets and comparable to MobileNets at similar model sizes.
- Conclusion: T2T-ViT progressively tokenizes images and structurally aggregates tokens to model image structure while using a deep-narrow architecture.The paper presents this combination as overcoming limitations of ViT.
- Conclusion: T2T-ViT achieves superior performance to ResNets and comparable performance to MobileNets with similar model size when trained from scratch on ImageNet.The authors position the result as supporting further development of Transformer-based vision models.