[Paper Review] Improve Vision Transformers Training by Suppressing Over-smoothing

Chengyue Gong, Dilin Wang+3 more|arXiv (Cornell University)|Apr 26, 2021

Advanced Neural Network Applications4 references22 citations

TL;DR

This paper proposes training stabilization techniques for vision transformers without architectural modifications, addressing over-smoothing in self-attention layers that causes feature collapse and performance degradation. By introducing diversity-promoting loss functions and a patch classification head for Cutmix, the method enables deeper and wider vision transformers to achieve 85.0% top-1 accuracy on ImageNet without extra parameters or teacher models.

ABSTRACT

Introducing the transformer structure into computer vision tasks holds the promise of yielding a better speed-accuracy trade-off than traditional convolution networks. However, directly training vanilla transformers on vision tasks has been shown to yield unstable and sub-optimal results. As a result, recent works propose to modify transformer structures by incorporating convolutional layers to improve the performance on vision tasks. This work investigates how to stabilize the training of vision transformers \emph{without} special structure modification. We observe that the instability of transformer training on vision tasks can be attributed to the over-smoothing problem, that the self-attention layers tend to map the different patches from the input image into a similar latent representation, hence yielding the loss of information and degeneration of performance, especially when the number of layers is large. We then propose a number of techniques to alleviate this problem, including introducing additional loss functions to encourage diversity, prevent loss of information, and discriminate different patches by additional patch classification loss for Cutmix. We show that our proposed techniques stabilize the training and allow us to train wider and deeper vision transformers, achieving 85.0\% top-1 accuracy on ImageNet validation set without introducing extra teachers or additional convolution layers. Our code will be made publicly available at this https URL .

Motivation & Objective

• To address the instability and sub-optimal performance of vanilla vision transformers during training.
• To identify over-smoothing in self-attention layers as the root cause of feature collapse and performance degradation in deep transformers.
• To improve training stability and model capacity without modifying the transformer architecture or adding convolutional layers.
• To enable training of deeper and wider vision transformers by preserving feature diversity and information.
• To achieve state-of-the-art accuracy on ImageNet using only loss-based regularization techniques.

Proposed method

• Introduce an auxiliary patch classification head to enhance feature discrimination during Cutmix augmentation.
• Apply a diversity-promoting loss function to encourage distinct representations across image patches.
• Use a contrastive-like loss to prevent feature collapse by preserving inter-patch differences.
• Integrate the patch classification loss with Cutmix data augmentation to improve generalization.
• Train vision transformers with standard self-attention mechanisms while stabilizing training through these loss components.
• Avoid architectural changes such as residual connections or convolutional inductive biases.

Experimental results

Research questions

• RQ1 What causes training instability in vision transformers when using only self-attention mechanisms?
• RQ2 How does over-smoothing in self-attention layers affect feature representation and model accuracy?
• RQ3 Can training stability be improved without architectural modifications to vision transformers?
• RQ4 To what extent can deeper and wider vision transformers be trained using only loss-based regularization?
• RQ5 Can high accuracy be achieved on ImageNet without additional parameters or knowledge distillation?

Key findings

• The proposed method stabilizes training of vision transformers without architectural modifications, enabling reliable training of deeper and wider models.
• The model achieves 85.0% top-1 accuracy on the ImageNet validation set, matching state-of-the-art results without using knowledge distillation or extra parameters.
• The addition of a patch classification head and diversity loss significantly reduces over-smoothing and improves feature discrimination.
• The method maintains performance gains even with increased depth and width, demonstrating robustness to architectural scaling.
• The approach outperforms vanilla vision transformers and matches models that rely on residual connections or convolutional inductive biases.
• The code is publicly released, enabling reproducibility and further research in training-stable vision transformers.