[Paper Review] DeepViT: Towards Deeper Vision Transformer

Daquan Zhou, Bingyi Kang+6 more|arXiv (Cornell University)|Mar 22, 2021

Advanced Neural Network Applications46 references348 citations

TL;DR

This paper analyzes why Vision Transformers saturate with depth and introduces Re-attention to regenerate diverse attention maps, enabling stable training of very deep ViTs (e.g., 32 blocks) and achieving improved ImageNet accuracy without extra data.

ABSTRACT

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.

Motivation & Objective

• Investigate why scaling Vision Transformers in depth saturates performance.
• Identify the cause of attention collapse in deep ViTs.
• Propose a lightweight mechanism (Re-attention) to diversify attention across layers.
• Demonstrate that deeper ViTs trained from scratch can improve accuracy on ImageNet-1k.

Proposed method

• Empirical study of ViT depth scaling on ImageNet to observe attention map similarity across layers.
• Define and quantify attention collapse via cross-layer attention similarity.
• Introduce Re-attention as a learnable head-to-head transformation that exchanges information among attention heads.
• Replace MHSA with Re-attention in ViT blocks to form DeepViT architectures.
• Compare DeepViT against SOTA CNNs and ViTs on ImageNet-1k without extra data or augmentation.
• Provide ablations on embedding dimension vs. depth and alternative attention sharpening baselines.

Experimental results

Research questions

• RQ1 Can ViTs meaningfully benefit from deeper architectures like CNNs, or do they saturate in performance?
• RQ2 What causes attention maps to become similar across layers in deep ViTs?
• RQ3 Can a lightweight mechanism reuse information across attention heads to restore diversity and enable deeper ViTs?
• RQ4 Do DeepViT models trained from scratch on ImageNet-1k outperform existing SOTA models with similar compute?

Key findings

• Directly increasing ViT depth yields saturating or even reduced accuracy on ImageNet when using standard MHSA.
• Attention maps in deeper layers become highly similar (attention collapse), correlating with stagnating feature evolution.
• Re-attention, which linearly mixes attention maps across heads with a learnable matrix, preserves diversity and eliminates cross-layer attention collapse.
• DeepViT models (32 blocks) achieve consistent accuracy gains and, on ImageNet-1k, outperform baseline ViTs and several CNN/ViT SOTA models without extra data or training tricks.
• Replacing MHSA with Re-attention yields 0 similar-block attention maps and improves Top-1 accuracy by up to 1.6 percentage points for 32-block models.
• DeepViT-S and DeepViT-L attain competitive or superior accuracy with fewer parameters than several contemporaries (e.g., DeepViT-L reaches 82.2–83.1% depending on variant).