Source-linked AI summary
DeepViT: Towards Deeper Vision Transformer
Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Zihang Jiang, Qibin Hou, Jiashi Feng
TL;DR
ViT performance saturates as models become deeper because attention maps become overly similar in later layers. The paper proposes Re-attention, which regenerates attention maps through learnable information exchange among heads, and reports consistent improvements for deep ViTs. The method maintains increasing performance with depth while adding minimal computation and memory overhead.
Problem
ViT performance saturates as depth rises because later attention maps become overly similar, limiting the evolution of representations.
Method
Re-attention learnably exchanges information among attention heads to regenerate attention maps before combining them with value features.
Results
Re-attention enables training ViTs with 32 transformer blocks and consistent performance improvements on ImageNet-1k without pretraining on larger datasets.
Takeaways & Limitations
Re-attention maintains increasing ViT performance as depth increases with minimal computation and memory overhead.
Takeaways & Limitations
The depth-scaling study fixes the hidden dimension and number of attention heads, although similar behavior was also observed when varying hidden dimension.
Abstract
from arXiv · showhide
Vision transformers (ViTs) have been successfully applied in image classification tasks recently. In this paper, we show that, unlike convolution neural networks (CNNs)that can be improved by stacking more convolutional layers, the performance of ViTs saturate fast when scaled to be deeper. More specifically, we empirically observe that such scaling difficulty is caused by the attention collapse issue: as the transformer goes deeper, the attention maps gradually become similar and even much the same after certain layers. In other words, the feature maps tend to be identical in the top layers of deep ViT models. This fact demonstrates that in deeper layers of ViTs, the self-attention mechanism fails to learn effective concepts for representation learning and hinders the model from getting expected performance gain. Based on above observation, we propose a simple yet effective method, named Re-attention, to re-generate the attention maps to increase their diversity at different layers with negligible computation and memory cost. The pro-posed method makes it feasible to train deeper ViT models with consistent performance improvements via minor modification to existing ViT models. Notably, when training a deep ViT model with 32 transformer blocks, the Top-1 classification accuracy can be improved by 1.6% on ImageNet. Code is publicly available at https://github.com/zhoudaquan/dvit_repo.
1. Introduction
ViTs do not reliably gain performance from greater depth, unlike deep CNNs, because attention maps become overly similar in later layers. Re-attention addresses this collapse by regenerating diverse attention maps and supports deeper ViTs with consistent improvements.
- Deep CNN progress has been driven by stacking many layers, motivating investigation of whether ViTs can similarly benefit from greater depth.
- ViT performance saturates as depth increases, with a 32-block model performing worse than the 24-block model on ImageNet.
- Attention collapse occurs when attention maps become overly similar after certain layers, causing representations to stop evolving.
- Re-attention exchanges information among attention heads to regenerate attention maps and address attention collapse.
- Replacing MHSA with Re-attention enables training ViTs with 32 transformer blocks and consistent performance improvements without extra augmentation or regularization.
2. Related Work
Prior work established transformers as important architectures for language and vision, while deep CNNs showed the value of effective depth scaling. This paper studies why ViT performance saturates with depth and proposes a self-attention modification to address it.
- Transformers replaced recurrence and convolutions with self-attention in machine translation before becoming dominant across NLP tasks.
- CNN research treated increasing depth as an effective route to performance improvement, while developing architecture strategies to make deep models trainable.
- DeepViT replaces each ViT self-attention layer with Re-attention to address attention collapse and enable deeper ViT training.
3. Revisiting Vision Transformer
This section reviews ViT architecture and examines how self-attention maps and features evolve across transformer depth. It finds that deeper blocks develop highly similar attention maps and features, helping explain performance saturation.
- ViT architecture: A ViT combines patch embeddings, stacked transformer blocks, and a classification layer; each block contains MHSA and an MLP.Image patches are encoded as tokens, processed with a class token, and passed through the transformer stack.
- Multi-head self-attention: MHSA computes attention from queries and keys, then uses it to retrieve and combine value embeddings for subsequent feature processing.For each head, the attention map is formed from Q and K, while V supplies the embeddings being combined.
- Depth scalability: The study varies ViT depth from 12 to 32 blocks while fixing the hidden dimension and number of heads to 384 and 12, respectively.This setup investigates whether increasing depth improves image-classification performance under controlled architectural settings.
- Attention-map analysis: Cross-layer attention similarity is measured with cosine-similarity matrices that compare token-specific attention contributions between layers.A similarity of one means a token plays exactly the same self-attention role in both layers.
- Attention collapse: In a 32-block ViT, more than 90% of attention maps are similar after block 17, while the fraction of similar blocks rises with model depth.Attention maps from different heads within the same block remain sufficiently diverse, unlike maps across deeper blocks.
- Feature degeneration: Feature similarity becomes high and features stop evolving after block 20, while Re-attention substantially reduces cross-layer feature-map similarity.The observed correspondence between attention and feature similarity links attention collapse with degeneration of deeper representations.
4. Re-attention for Deep ViT
The section identifies attention collapse as a major obstacle to scaling ViTs deeper and examines higher embedding dimensions and Re-attention as solutions. Re-attention regenerates attention maps through learnable cross-head communication, improving diversity with negligible overhead.
- Attention collapse is identified as a major obstacle to scaling ViTs deeper.
- 4.1. Self-Attention in Higher Dimension Space: Increasing embedding dimension reduces similar attention maps and alleviates attention collapse, with corresponding performance gains.
- 4.1. Self-Attention in Higher Dimension Space: Higher embedding dimensions substantially increase computation and model size, while performance improvements diminish and overfitting risk can increase.
- 4.2. Re-attention: Re-attention regenerates attention maps by dynamically aggregating information across attention heads using a learnable transformation matrix Θ.
- 4.2. Re-attention: Re-attention exchanges complementary information among heads to improve attention-map diversity with negligible computational overhead.
5. Experiments
Experiments show that deeper vanilla ViTs suffer attention collapse, while Re-attention preserves cross-layer diversity and improves deep-model performance with negligible overhead.
- Experimental scope: The experiments evaluate attention collapse, Re-attention, and comparisons with temperature tuning, attention dropping, and state-of-the-art models.Ablations and visualizations examine why performance saturates as ViTs become deeper.
- Attention collapse: For a 32-block ViT, sharing Q and K values across the last 15 blocks causes negligible performance degradation.The result indicates that many top blocks contribute little additional accuracy because their attention maps are highly similar.
- Re-attention results: Replacing self-attention with Re-attention removes similar blocks and yields consistently increasing performance as transformer depth increases.The gain is especially significant for the 32-block model, without extra datasets or augmentation policies.
- Attention visualization: Original self-attention maps become nearly uniform and lose diversity in deeper blocks, whereas Re-attention produces more diverse maps.The visualization contrasts local shallow-block relationships with progressively collapsed attention in the baseline.
- Baseline comparisons: Adding a SoftMax temperature reduces attention-map similarity only slightly, while learnable temperature improves ImageNet performance marginally.Attention dropping reduces map similarity but does not substantially reduce feature similarity.
- State-of-the-art comparison: DeepViT-L outperforms T2T-ViT and DeiT by 0.4 points with a smaller model size of 55M versus 64M and 86M.The comparison uses DeepViT variants built with Re-attention and reports ImageNet top-1 accuracy.
6. Conclusion
The paper identifies attention collapse as a depth-related problem in ViTs and proposes Re-attention to address it with minimal overhead. The method maintains increasing performance as ViT depth grows.
- Conclusion: The paper finds that attention maps become similar as ViTs deepen, causing attention collapse, and introduces Re-attention to address it.Re-attention is designed with minimum computation and memory overhead.
- Conclusion: Re-attention enables increasing ViT performance when network depth increases.The conclusion presents this as the practical consequence of the proposed mechanism.
A. Experiment Implementation Details
The implementation details define attention reuse through a learned transformation and normalization, alongside a cosine-similarity regularizer and comparable model sizing.
- Attention reuse: Forcing the top 15 blocks of a 32-block ViT to share one attention map causes negligible ImageNet accuracy degradation.This supports the claim that those blocks do not contribute to accuracy improvement.
- Attention reuse: The reused attention output applies Norm(Θ⊤Aunique)V, where Θ models small inter-layer variance and Norm is batch normalization.Aunique is taken from the last unique block and used by subsequent blocks.
- Training objective: Training uses cross-entropy loss plus cosine similarity between adjacent attention maps for layer-diversity regularization.The number of bottom blocks regularized is set to 4, 8, and 12 for 16-, 24-, and 32-block ViTs.
B. DeepViT architecture design
The architecture study applies Re-attention only to upper blocks, where attention maps collapse, and selects the number of such blocks using similarity and accuracy patterns.
- Design motivation: Because attention maps become similar mainly in top blocks, applying Re-attention to every block is unnecessary.The architecture search therefore focuses on upper transformer blocks.
- Block selection: Accuracy changes with the number of Re-attention blocks, and peaks when that number matches the number of similar attention-map blocks.This criterion yields 5 Re-attention blocks for DeepViT-S and 12 for DeepViT-L.
C. Impacts of hyper-parameters
DeepViT performance can improve further when trained with optimized training hyper-parameters rather than the main paper’s recipe. The appendix compares training recipes, including DeiT-style settings without repeated augmentation.
- The appendix evaluates DeepViT-S and DeepViT-L using the same training recipes as DeiT, except without repeated augmentation.
- Optimized training hyper-parameters further improve DeepViT performance.The appendix reports this improvement relative to the main-paper training setup.
- Table 8 reports DeepViT models trained with different training recipes, with a marker identifying DeiT-recipe training.
D. Similarity calculation
The paper measures attention-map similarity across layers using cosine similarity and summarizes local block similarity with an 80% threshold.
- Cross-layer similarity is defined as the ratio of similar vector pairs to all vector pairs between two attention maps.The similarity compares attention maps from layers p and q.
- A hyper-parameter τ determines whether vectors are considered similar.
- A block is counted as similar when its attention map has over 80% similarity with the adjacent block’s map.
- Block similarity for a ViT with B blocks is measured as the ratio of similar blocks to the total number of blocks.
- The paper selects 20.5 as a visualization threshold.