TransMix: Attend to Mix for Vision Transformers (2111.09833v1)

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.

• The paper introduces TransMix, a novel method that uses Vision Transformer attention maps to inform label mixing during data augmentation, addressing limitations of traditional mixup approaches.
• TransMix leverages existing attention maps to derive label weights based on the attention paid to mixed regions, requiring no additional parameters or computational resources.
• Experiments show TransMix consistently boosts ViT accuracy on ImageNet and improves transferability and robustness in downstream computer vision tasks.

An Analytical Overview of "TransMix: Attend to Mix for Vision Transformers"

The paper "TransMix: Attend to Mix for Vision Transformers" introduces a novel approach to data augmentation specifically designed for Vision Transformers (ViTs). The primary emphasis of the work is on mitigating the issues arising from traditional mixup-based augmentation methods, which often assume a linear correspondence between input interpolation and label mixing—a premise that may not hold for datasets where images have an uneven distribution of salient objects.

Key Contributions

The authors propose a method called TransMix, which leverages the inherent attention maps generated by Vision Transformers to inform label mixing. This approach aims to improve label assignment by weighing labels according to the attention paid by the model to different image regions, thus addressing the problem of mismatch between input space and label space. This technique enhances model generalization without increasing model complexity or computational overhead.

Methodology

1. Vision Transformers and Mixup-based Augmentation:
   • Vision Transformers, influenced by their success in NLP, have made notable inroads in computer vision tasks. The architecture, however, is prone to overfitting.
   • Mixup-based augmentations (e.g., Mixup, CutMix) have been effective at improving generalization. These techniques, however, uniformly interpolate between labels based on pixel mixing, potentially ignoring context or importance.
2. TransMix Approach:
   • TransMix assigns label weights based on an attention map derived from ViTs. Specifically, the sum of the attention weights of mixed regions informs the label mixing coefficient, thereby offering a more context-aware label interpolation.
   • By using an existing attention map, this method requires no additional parameters or computational resources, which is a significant advantage in scalable deployment.

Experimental Findings

The experimentation conducted comprehensively demonstrates the effectiveness of TransMix across diverse transformer models such as DeiT, XCiT, and PVT. Key findings include:

• Performance on ImageNet: Across various ViT architectures, TransMix consistently boosts top-1 accuracy by up to 0.9%, compared to traditional augmentation baselines.
• Transferability: TransMix-pretrained models exhibit superior performance in downstream tasks like image segmentation, object detection, and instance segmentation. These results suggest that models trained with TransMix on ImageNet generalize better to other vision tasks.
• Robustness: The paper evaluates model robustness under occlusion and adversarial example setups. TransMix-enhanced models demonstrate improved resilience compared to conventional training regimes.

Implications and Future Work

The implications of TransMix are both practical and theoretical. Practically, it provides a simple yet effective augmentation technique that can be immediately applied within existing ViT pipelines, optimizing their performance without additional training costs or architectural changes. Theoretically, it explores the underutilized potential of attention maps for label adjustment, paving the way for future research on adaptive label mixing methodologies.

Possible directions for future work include adapting TransMix for transformer variants without class tokens, where integrating class attention might require architectural changes. Further exploration into the synergy between self-attention mechanisms and data augmentation could yield additional insights into enhancing model robustness and effectiveness across varied tasks.

Overall, this paper offers valuable insights into the improvement of data augmentation techniques for vision transformers, laying the groundwork for subsequent advances in the field of computer vision.