Vision Transformers with Patch Diversification

TL;DR

This paper introduces a novel training approach for vision transformers that promotes diversity among patch representations, leading to more stable training and improved performance on various vision tasks.

Contribution

It proposes new loss functions to explicitly encourage patch diversity, stabilizing training without modifying the transformer architecture.

Findings

Stabilizes training of wider and deeper vision transformers.

Enhances transfer learning performance on downstream tasks.

Improves state-of-the-art results in semantic segmentation.

Abstract

Vision transformer has demonstrated promising performance on challenging computer vision tasks. However, directly training the vision transformers may yield unstable and sub-optimal results. Recent works propose to improve the performance of the vision transformers by modifying the transformer structures, e.g., incorporating convolution layers. In contrast, we investigate an orthogonal approach to stabilize the vision transformer training without modifying the networks. We observe the instability of the training can be attributed to the significant similarity across the extracted patch representations. More specifically, for deep vision transformers, the self-attention blocks tend to map different patches into similar latent representations, yielding information loss and performance degradation. To alleviate this problem, in this work, we introduce novel loss functions in vision transformer training to explicitly encourage diversity across patch representations for more discriminative feature extraction. We empirically show that our proposed techniques stabilize the training and allow us to train wider and deeper vision transformers. We further show the diversified features significantly benefit the downstream tasks in transfer learning. For semantic segmentation, we enhance the state-of-the-art (SOTA) results on Cityscapes and ADE20k. Our code is available at https://github.com/ChengyueGongR/PatchVisionTransformer.