PaperlayerView exact revision on arXiv

Source-linked AI summary

TransMix: Attend to Mix for Vision Transformers

Jie-Neng Chen, Shuyang Sun, Ju He, Philip Torr, Alan Yuille, Song Bai

arXiv:2111.09833v1cs.CV

TL;DR

Mixup-based augmentation assumes that label proportions should match input proportions, although pixels may differ in significance. TransMix uses Vision Transformer attention to weight mixed labels, improving ImageNet classification and transferring to downstream tasks and robustness benchmarks.

  • Problem

    Mixup methods use the same proportion for input and label interpolation, although image regions can contribute unequally to the label and ViTs can overfit.

    Source: §1 · pp. 1–2

  • Method

    TransMix uses multi-head class attention over image patches to recalculate mixed-label weights according to the attention assigned to each input.

    Source: §3.2 · p. 3

  • Results

    TransMix raises ImageNet top-1 accuracy by 0.9% for both DeiT-S and XCiT-L, transfers to downstream vision tasks, and improves robustness on four benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    TransMix uses attention already generated by ViTs, requiring no extra parameters and minimal computation overhead while showing effectiveness across ten benchmarks.

    Source: §1 · p. 2, §5 · p. 8

    TransMix relies on class attention and may not work well with backbones lacking a class token; it also requires attention maps spatially aligned with the input.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

ViTs can overfit when training data is insufficient, and conventional mixup assumes label proportions should match input proportions even when pixels differ in significance. TransMix uses ViT attention to assign mixed-label weights and improves ImageNet results while supporting downstream transfer and robustness.

1. Introduction

  • ViTs can be hard to optimize and overfit when training data is insufficient, while mixup-based augmentation can help them generalize.Mixup assumes linearly interpolated inputs should have targets interpolated at the same ratio, although background and salient pixels may contribute unequally to labels.Source: §1 · p. 1
  • TransMix sets the weight of each mixed label from the sum of attention-map values over that image’s region, rather than using the input mixing ratio.It uses attention maps naturally generated in ViTs, with no extra parameters and minimal computation overhead.Source: Figure 1 · p. 1, §1 · p. 2
  • TransMix raises ImageNet top-1 accuracy by 0.9% for both DeiT-S and XCiT-L, and Figure 2 reports steady gains across the listed ViT models.The largest XCiT-L gains the most among XCiT scales.Source: §1 · p. 2, Figure 2 · p. 2
  • ImageNet pretraining with TransMix transfers to object detection, instance segmentation, semantic segmentation, and weakly supervised object segmentation or localization.The authors also report improved robustness on four benchmarks.Source: §1 · p. 2

2. Related Work

Related work situates TransMix among Vision Transformers, mixup augmentation, and data-adaptive loss weighting. Its distinction from saliency-based mixup methods is to reassign labels using native Transformer attention rather than manipulate salient input regions.

Vision Transformers

  • Vision Transformers adapt Transformers to image recognition by tokenizing and flattening images, and have been applied to classification, detection, and segmentation.DeiT adds a training recipe and knowledge distillation to strengthen ViT.Source: §2 · p. 2

Mixup and its variants

  • TransMix differs from saliency-based mixup methods by reassigning labels instead of cropping salient input regions, using Transformer attention without extra parameters.The paper reports better ImageNet results than those methods; prior approaches may use extra parameters to extract salient regions.Source: §2 · p. 2

Data-adaptive loss weight assignment

  • TransMix’s attentional label reassignment is related to data-adaptive loss weighting, which prior work has also explored for instance segmentation and image demosaicing.Source: §2 · p. 2

3. TransMix

The section introduces CutMix and self-attention, then defines TransMix’s label assignment using class attention over image patches. A short plug-and-play implementation uses downsampled masks to weight labels dynamically.

3.1. Setup and Background

  • CutMix combines two input-label pairs into an augmented sample using a binary mask, with the mixed-label proportion tied to the cropped area ratio.The mask indicates which image pixels are removed and filled from the other image.Source: §3.1 · pp. 2–3
  • Self-attention computes an attention map from queries and keys, then uses it to form a weighted sum of token features.Multi-head attention applies learned projections across multiple heads.Source: §3.1 · p. 3

3.2. TransMix

  • TransMix uses multi-head class attention: the class token queries image-patch tokens, and their attention indicates which patches are useful to the classifier.The attention heads are averaged, and the resulting attention is available from the last Transformer block without architecture modification.Source: §3.2 · p. 3
  • TransMix retains CutMix’s input mixture but recalculates the label weight from attention summed over the downsampled mask, dynamically favoring the image more attended by the model.Nearest-neighbor downsampling maps the original mask to the patch grid.Source: §3.2 · p. 3, §3.1 · p. 3

3.3. Pseudo-code

  • The pseudo-code presents TransMix as a plug-and-play method implemented in a few lines of code.Source: §3.3 · p. 3

4. Experiments

The experiments assess TransMix on ImageNet classification, downstream transfer, robustness, attention, Swin generalization, and comparisons with mixup variants. Across these settings, the reported findings include classification gains, downstream improvements, and robustness benefits.

4.1. ImageNet Classification

  • TransMix improves ImageNet top-1 accuracy across the evaluated DeiT, PVT, CaiT, and XCiT models, including 0.9% gains for DeiT-S and XCiT-L.The reported gains for base and large variants range from 0.6% to 0.9%.Source: §4.1 · p. 4, Table 1 · p. 4

4.2. Transfer to Downstream Tasks

  • TransMix-pretrained models improve over vanilla-pretrained baselines on semantic segmentation, object detection, and instance segmentation.On Pascal Context, DeiT-S-Linear and DeiT-S-Segmenter improve by 0.6% and 0.9% mIoU; on COCO, the PVT-backed detector improves by 0.5% box AP and 0.6% mask AP.Source: §4.2 · p. 5, Table 2 · p. 5, Table 3 · p. 5

4.3. Robustness Analysis

  • TransMix-trained DeiT-S performs better than vanilla DeiT-S across reported occlusion, spatial-shuffling, natural-adversarial, and out-of-distribution robustness evaluations.For shuffling, average accuracy is 62.8% for TransMix-DeiT-S versus 58.4% for DeiT-S; on ImageNet-A, TransMix is superior on all reported metrics, and on ImageNet-O it improves AUPR by 1%. For occlusion, it outperforms vanilla DeiT-S at all levels, especially above 0.7 information loss.Source: §4.3 · pp. 5–6, Table 4 · p. 6

4.4. Mutual Effect of TransMix and Attention

  • Using external parameter-frozen attention maps contradicts the hypothesis that better attention necessarily improves TransMix, while the authors attribute the outcome to TransMix’s dynamic attention.With external fully supervised DeiT-S attention, top-1 accuracy falls from 80.7% to 80.4%, remaining above the vanilla model’s 79.8%. The attention evaluation also examines weakly supervised segmentation and localization using class-level training labels.Source: §4.4 · p. 6

4.5. Generalizability Study

  • A classification-attention block makes TransMix applicable to Swin, whose patch-only average pooling does not expose class attention, and TransMix-CA-Swin-T reaches 81.8% top-1 accuracy.The compared Swin-T, CA-Swin-T, and TransMix-CA-Swin-T models have 28.3M parameters; both CA variants use 7% fewer FLOPs than baseline Swin-T.Source: §4.5 · p. 7, Table 7 · p. 7

4.6. Comparison with State-of-the-art Mixup Variants

  • TransMix significantly outperforms the compared mixup variants, with a 2.1% performance advancement, highest training throughput, and no parameter overhead.The comparison uses DeiT training recipes; the supplementary attention-map ablation reports that the default is best. Visualizations show attention can correct area-based label assignments and increase label weight for discriminative fine-grained attributes.Source: §4.6 · p. 8, Table 8 · p. 8, Figure 5 · p. 8

5. Conclusion

TransMix uses ViT attention to guide Mixup label assignment and reports gains across image classification, transfer, robustness, and generalizability. The paper also identifies limitations involving backbone attention and spatial alignment.

5. Conclusion

  • TransMix assigns Mixup labels with attentional guidance, using Transformer attention to set confidence for the mixed target.The paper reports 0.9% higher ImageNet top-1 accuracy for both DeiT-S and XCiT-L, and experiments across 10 benchmarks.Source: §5 · p. 8
  • TransMix can struggle with backbones lacking a class token and may not suit deformable Transformers whose attention is not spatially aligned with the input.The paper notes that the class-token limitation can be mitigated with architecture modification and suggests calibrating attention to input locations for deformable models.Source: §5 · p. 8

A.1. Implementation Details of Compared Mixup Variants

The compared Mixup variants are implemented on DeiT-S, with the baseline excluding CutMix; the variants use different saliency or gradient procedures.

A.1. Implementation Details of Compared Mixup Variants

  • The comparison applies Mixup variants originally implemented mainly for ResNet-50 to DeiT-S.The baseline in Table 8 is the default DeiT-S framework without CutMix.Source: §A.1 · p. 9
  • Attentive-CutMix uses a ResNet-50 model to extract saliency maps alongside its classification backbone.Source: §A.1 · p. 9
  • SaliencyMix extracts saliency with OpenCV, whereas Puzzle-Mix detects object saliency by computing network gradients with two model passes.The supplied implementation description names the OpenCV saliency extraction function for SaliencyMix.Source: §A.1 · p. 9

B. Additional Results

Additional results report mixup-variant comparisons on CNNs and ViTs and examine attention-generation choices for TransMix.

B. Additional Results

  • For DeiT-S, attention from block d = 12 performs best among the tested attention-generation choices.The ablation also evaluates attention rollout, which multiplies attention matrices across all 12 blocks.Source: B. Additional Results · p. 9, Table 9 · p. 9
  • Using the last Transformer block's attention avoids exhaustively selecting a block and is compatible with ViT variants.For PVT's hierarchical blocks, using early-scale attention requires an additional downsampling step.Source: B. Additional Results · p. 9
  • SaliencyMix and Puzzle-Mix improve over CutMix by at most 0.2% on ResNet-50, but show no advancement over CutMix on DeiT-S.Table 10 compares state-of-the-art mixup variants on ImageNet-1k with ViT or CNN backbones; listed models use 300 training epochs.Source: B. Additional Results · p. 9, Table 10 · p. 9
Loading 2111.09833v1…