TransMix: Attend to Mix for Vision Transformers

Abstract

Mixup-based augmentation has been found to be effective for generalizing models during training, especially for Vision Transformers (ViTs) since they can easily overfit. However, previous mixup-based methods have an underlying prior knowledge that the linearly interpolated ratio of targets should be kept the same as the ratio proposed in input interpolation. This may lead to a strange phenomenon that sometimes there is no valid object in the mixed image due to the random process in augmentation but there is still response in the label space. To bridge such gap between the input and label spaces, we propose TransMix, which mixes labels based on the attention maps of Vision Transformers. The confidence of the label will be larger if the corresponding input image is weighted higher by the attention map. TransMix is embarrassingly simple and can be implemented in just a few lines of code without introducing any extra parameters and FLOPs to ViT-based models. Experimental results show that our method can consistently improve various ViT-based models at scales on ImageNet classification. After pre-trained with TransMix on ImageNet, the ViT-based models also demonstrate better transferability to semantic segmentation, object detection and instance segmentation. TransMix also exhibits to be more robust when evaluating on 4 different benchmarks. Code will be made publicly available at https://github.com/Beckschen/TransMix.

arXiv commentCode will be made publicly available at https://github.com/Beckschen/TransMix

References in corpus (16)

• TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
• BEiT: BERT Pre-Training of Image Transformers
• MMDetection: Open MMLab Detection Toolbox and Benchmark
• Manifold Mixup: Better Representations by Interpolating Hidden States
• Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• Intriguing Properties of Vision Transformers
• XCiT: Cross-Covariance Image Transformers
• Do Vision Transformers See Like Convolutional Neural Networks?
• Are Transformers More Robust Than CNNs?
• Puzzle Mix: Exploiting Saliency and Local Statistics for Optimal Mixup
• Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
• CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification
• Towards Robust Vision Transformer
• Learning to Sample the Most Useful Training Patches from Images
• TS-CAM: Token Semantic Coupled Attention Map for Weakly Supervised Object Localization