Source-linked AI summary
Going deeper with Image Transformers
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, Hervé Jégou
TL;DR
Image transformers lacked evidence that increasing depth would improve ImageNet performance because deeper models were difficult to optimize. The paper introduces LayerScale and CaiT's class-attention design, achieving strong results without additional training data, including 86.5% ImageNet accuracy with lower reported complexity than a competing model.
Problem
The work addresses the limited understanding of optimization for deeper image transformers, whose larger depths had performed poorly or failed to converge effectively on ImageNet.
Method
The paper combines LayerScale, which scales each residual branch per channel from a small initialization, with CaiT class-attention layers that separate patch processing from class-embedding extraction.
Results
86.5% top-1 accuracy on ImageNet1k-val is reported with 329B FLOPs and 356M parameters versus 377B FLOPs and 438M parameters for the best competing model, while state-of-the-art results are achieved on ImageNet-Real and ImageNet-V2 matched frequency without additional training data.
Takeaways & Limitations
The results show that deeper transformer-based image classifiers can be trained on ImageNet alone and can compete with leading convolutional networks on accuracy-complexity trade-offs.
Takeaways & Limitations
The class-attention design assumes that copying information from the class embedding back to patch embeddings provides no benefit, an assumption supported by a fixed-12-layer comparison.
Abstract
from arXiv · showhide
Transformers have been recently adapted for large scale image classification, achieving high scores shaking up the long supremacy of convolutional neural networks. However the optimization of image transformers has been little studied so far. In this work, we build and optimize deeper transformer networks for image classification. In particular, we investigate the interplay of architecture and optimization of such dedicated transformers. We make two transformers architecture changes that significantly improve the accuracy of deep transformers. This leads us to produce models whose performance does not saturate early with more depth, for instance we obtain 86.5% top-1 accuracy on Imagenet when training with no external data, we thus attain the current SOTA with less FLOPs and parameters. Moreover, our best model establishes the new state of the art on Imagenet with Reassessed labels and Imagenet-V2 / match frequency, in the setting with no additional training data. We share our code and models.
1 Introduction
The paper studies why deeper image transformers are difficult to optimize and proposes LayerScale and class-attention to improve training and accuracy. These changes enable competitive or state-of-the-art results with favorable complexity and transfer performance.
- Motivation: Deeper image transformers are difficult to optimize, motivating analysis of the interaction between architecture and optimization.Residual architectures are easier to train because they provide clear forward and backward paths, while existing image-transformer training showed limited evidence that depth helps.
- LayerScale: LayerScale adds a learnable diagonal transformation after each residual block, initialized near zero to improve the training dynamics of deeper networks.The method progressively incorporates residual-branch contributions while preserving the identity-like initial function.
- CaiT: Class-attention separates patch self-attention from class-embedding extraction, avoiding the contradictory objectives assigned to the class embedding in ViT.The class-attention layers compile patch embeddings into a single class embedding for a linear classifier.
- Results: 86.5% top-1 accuracy on ImageNet1k-val is achieved with 329B FLOPs and 356M parameters, compared with 377B FLOPs and 438M parameters for the best competing model.The result is reported as on par with the state of the art while using fewer FLOPs and parameters.
- Results: The best CaiT models establish state-of-the-art results on ImageNet-Real and ImageNet-V2 matched frequency without additional training data.The paper also reports competitive transfer-learning results and overall better performance than EfficientNet-B7 across datasets.
2 Deeper image transformers with LayerScale
This section addresses unstable or saturating optimization in deeper image transformers by combining residual-branch scaling with training choices that preserve convergence. LayerScale uses small, learnable per-channel weights to keep initial residual contributions small and enable deeper high-capacity models to benefit from depth.
- Deeper ViT architectures lacked evidence of benefiting from depth on ImageNet-only training, motivating work on optimization stability.
- Pre-normalization and warmup are retained because removing them made Fixup and T-Fixup training unstable, while amended variants converge with DeiT models.
- LayerScale adds a learnable diagonal matrix after each residual block, providing per-channel rather than scalar weighting of residual outputs.
- Small initial diagonal values keep residual branches’ initial contribution small, encouraging training near the identity function before progressively integrating additional parameters.The initialization is ε = 0.1 through depth 18, ε = 10^-5 at depth 24, and ε = 10^-6 for deeper networks.
- LayerScale preserves expressive power while offering more optimization degrees of freedom than a single learnable scalar.
3 Specializing layers for class attention
CaiT separates patch self-attention from class-attention so different layers can specialize in processing patches and summarizing them for classification. Its class-attention stage updates only the class embedding and makes the added layers cheaper in memory and computation.
- CaiT separates self-attention between patches from class-attention that extracts processed patch information into a classifier-ready class vector.
- The design removes the contradictory requirement that shared weights both guide patch attention and summarize information for classification.
- Architecture: The self-attention stage is ViT-like but excludes the class embedding, while the class-attention stage compiles patch embeddings into CLS for a linear classifier.
- Architecture: During class-attention processing, only the class embedding is updated; the patch embeddings remain frozen.
- Two class-attention and feed-forward blocks are sufficient to cap performance in preliminary experiments, denoted as 12+2 for twelve self-attention/feed-forward and two class-attention/feed-forward blocks.
- Complexity: CaiT’s extra class-attention layers change attention complexity from quadratic to linear in the number of patches and use feed-forward matrix-vector multiplications.
4 Experiments
Experiments address the instability and early saturation of deeper image transformers through depth-aware optimization and architectural analysis. LayerScale stabilizes residual contributions, while class-attention improves class-embedding processing.
- Training deeper transformers: DeiT training fails to converge properly above 18 layers without depth-related hyperparameter adjustments.
- Training deeper transformers: LayerScale and adapted normalization methods enable deeper models to converge without saturating as early.
- Hyperparameter analysis: Stochastic-depth drop rates must be adapted to network depth, while the selected formula requires further adjustment for different working dimensionalities.
- LayerScale analysis: LayerScale makes residual-branch contributions more uniform across layers in a 36-block transformer.
- Class-attention analysis: Class-attention layers separate patch self-attention from class-vector extraction, and adding two class-attention blocks improves performance at fixed self-attention depth.
4.3 Our CaiT models
CaiT combines ViT with LayerScale and a two-stage class-attention architecture, then evaluates model capacity, resolution, distillation, and transfer performance. The resulting models achieve strong accuracy and efficiency across ImageNet settings and transfer datasets.
- Model design: CaiT incorporates LayerScale and class-attention layers into ViT, with depth and working dimensionality controlling model capacity.
- Training and evaluation variants: Higher-resolution fine-tuning boosts performance without changing parameter count but increases computation, whereas hard distillation boosts accuracy without changing parameters or speed.
- ImageNet results: 86.5% top-1 accuracy is achieved by CaiT-M48↑448Υ on ImageNet1k-val, improving over DeiT’s 85.2%.
- Transfer learning: CaiT models achieve overall better transfer-learning performance than EfficientNet-B7 across the evaluated datasets.
4.5 Ablation
Ablations show that CaiT’s components and training choices contribute complementarily to performance. Resolution, training duration, crop ratio, head count, and fine-tuning affect the accuracy–efficiency trade-off.
- DeiT to CaiT: CaiT is complementary with LayerScale and improves performance without significantly increasing FLOPs in the stepwise DeiT-to-CaiT ablation.
- Resolution and training: Fine-tuning at higher resolution improves performance while saving computation during training compared with training from scratch.
- Attention heads: Eight self-attention heads provide a compromise between accuracy and speed because increasing heads fragments computation and lowers effective hardware throughput.
- Evaluation crop ratio: A crop ratio of 1.0 improves performance for transformer models relative to the conventional 0.875 setting.
- Training schedule: Increasing training from 300 to 400 epochs improves CaiT-S-36, while extending training to 500 epochs changes performance insignificantly.
5 Visualizations
CaiT’s class-attention maps separate object-focused and more globally contextual processing. Visualizations aggregate attention across heads to show salient regions and occasional classification failures.
- The first class-attention layer focuses on the main object driving the classification decision, with heads attending to identical or complementary object parts.
- Figure 6 displays four attention heads for each of the two class-attention layers in an XXS model.
- The second class-attention layer appears to emphasize image context or the image globally rather than a narrowly localized object.
- Attention visualizations average the heads’ maps from the first class-attention layer, upsample them, and modulate a normalized grayscale image.
- The saliency examples show that CaiT can focus on distinct regions, such as a racket and tennis ball, while also producing failures such as classifying a church top as a flagpole.
6 Related work
Related work frames deep image-transformer performance as an architecture-and-optimization problem. Prior approaches modify receptive fields, combine CNNs and transformers, or address training through initialization, staging, losses, added components, and regularization.
- CNNs have been the standard for image classification, while attention-based approaches investigate longer-range interactions within or beyond convolutional receptive fields.
- Transformer-based vision models evolved from pixel-level approaches to patch-based architectures, with performance relative to CNNs changing as the field progressed.
- Encoder-decoder transformers originated in machine translation and later became influential in language modeling and speech recognition.
- Training deeper architectures may require changes to initialization, multi-stage training, depth-specific losses, architectural components, or regularization.
7 Conclusion
The paper shows how to train deeper transformer-based image-classification networks using only ImageNet and introduces CaiT as an encoder/decoder-inspired architecture. It positions transformers as competitive with leading CNNs when accuracy and complexity are considered together.
- The work demonstrates training deeper transformer-based image-classification networks on ImageNet alone.
- CaiT is introduced as a simple architecture designed in the spirit of encoder/decoder architectures.
- Transformer models are presented as a competitive alternative to the best CNNs under accuracy–complexity trade-offs.
Appendix
The supplement reports architectural variations that informed the design of class-attention layers and LayerScale.
- Supplementary experiments explored variations of the proposed architecture to guide the design of class-attention layers and LayerScale.
A Variations on LayerScale init
The section examines LayerScale initialization and whether its learned scaling factors, rather than merely their final values, support deeper transformer training. Experiments indicate that learnable evolution benefits the deepest models.
- Initialization choices: LayerScale uses per-channel scaling factors initialized at small values to stabilize training of deeper models.The section compares zero, random, and small-constant initialization strategies.
- Retraining: Retraining with fixed scaling factors obtained from LayerScale converges, but does not match the performance of learnable weighting factors.The control uses fixed factors while other training conditions remain identical.
- Retraining: The lower results from fixed-weight retraining suggest that parameter evolution during training benefits the deepest models.The fixed-weight control is only slightly better than the baseline with adjusted stochastic-depth drop rate.
B Design of the class-attention stage
The class-attention stage separates patch processing from extracting a class representation, while experiments assess key composition, LayerScale, and distillation choices. The reported alternatives produce results close to the adopted design, except that the distillation token offers no observed advantage here.
- Key composition: 83.31% top-1 accuracy on ImageNet1k-val is obtained without including the class token in class-attention keys, versus 83.44% with the adopted choice.The reported +0.13% difference is considered likely not significant.
- LayerScale: 83.36% top-1 accuracy on ImageNet1k-val is obtained without LayerScale in class-attention blocks, versus 83.44% with LayerScale.The reported +0.08% difference is not considered significant enough to establish a clear advantage.
- Distillation: A distillation token provides no observed advantage over hard distillation when used with class-attention layers.The paper therefore considers only hard distillation, which replaces the label with an average involving the teacher prediction.