TokenMix: Rethinking Image Mixing for Data Augmentation in Vision Transformers

To enhance the localization ability of CNNs, CutMix [37] proposed to mix pairs of samples with a random rectangular binary mask. Let $x\in\mathbb{R}^{H\times W\times C}$ and $y$ denote a training image and its label, respectively. Given a pair of training samples $(x_{a},y_{a})$ and $(x_{b},y_{b})$, CutMix generates a new training sample $(\tilde{x},\tilde{y})$ as follows:

where $M\in{\{0,1\}}^{H\times W}$ denotes the rectangular mask that decides where to drop out and fill in the contents of the two images, $\odot$ denotes element-wise multiplication, and $\lambda$ is sampled from a beta distribution Beta$(\alpha,\alpha)$. The binary mask $M$ is a randomly sampled rectangle, which guarantees $\frac{\sum M}{HW}=\lambda$. Similar to Mixup [39], CutMix assigns a mixed target for the generated image as a linear combination of $y_{a}$ and $y_{b}$.

We argue that the region-level mixing in CutMix might not be suitable for transformer-based architectures. As CNNs are primarily designed to encode local image contents, training with CutMix effectively prevents CNNs from over-focusing on local context. However, transformer-based architectures might be less benefited from CutMix as all of its layers have global receptive fields. In addition, the label of the mixed image is a linear combination of $y_{a}$ and $y_{b}$ with mixing ratio $\lambda$ being estimated only according to the size of the mask, which might be inappropriate in many cases as shown in Figure 8 (b). Although there were recent methods on attempting to improve CutMix by choosing the salience regions to maximize the saliency in the mixed images [17, 18, 34, 31], the salient areas might not correctly correspond to the target class [2], and the label noise problem is still serious.

In this paper, we propose TokenMix to mix a pair of images to generate a mixed image and learning target. We generate the mask $M$ at token-level to encourage better learning of long-range dependency and assign the target score of the mixed image according to the content-based neural activation maps of the two mixing images, which follows the general spirit of distillation to create more robust targets.

Figure 2 shows an overview of our proposed TokenMix. We first partition the input image $x$ into non-overlapping patches $x^{p}\in\mathbb{R}^{\frac{H}{P}\times\frac{W}{P}\times(P^{2}\cdot C)}$, which are then linearly projected to visual tokens. We then generate a random mask $M_{t}\in\mathbb{R}^{\frac{H}{P}\times\frac{W}{P}}$ at token-level according to the mask-out ratio $\lambda$. The mixed new training sample $(\tilde{x}^{p},\tilde{y})$ is created as follows:

where $\mathfrak{S}$ indicates the set of all tokens, $\odot$ denotes element-wise multiplication, $M_{ti}$ denotes the $i$-th token of the mask $M_{t}$, $A_{ai}$ and $A_{bi}$ are the $i$-th token of the spatially normalized neural activation maps of $x_{a}$ and $x_{b}$ respectively. The neural activation maps are generated with pre-trained networks’ last layer before the classification head [16, 38].

Instead of masking a whole rectangular area, we partition the mask area into multiple separated parts. For each part, we randomly choose the number of masked tokens and the aspect ratio [1, 37]. We set the minimum number of tokens to 14 and log-uniformly sample the aspect ratio in the range of $[0.3,\frac{1}{0.3}]$. We repeatedly mask a part of the image until the total number of masked tokens reaches the pre-defined ratio $\lambda\frac{HW}{P^{2}}$. Instead of sampling $\lambda$ from a beta distribution, we set $\lambda$ to 0.5 unless otherwise specified. Our intuition is that the distributed masking regions are easier to recognize compared with masking a whole rectangular area. For investigation, we also introduce a uniformly random version, where each masking part is only a single token. While totally random mixing is harmful to the performance of CNNs, we show that transformers are still benefited from the simplified version.

To solve the issue of inaccurate target scores generated by CutMix, we propose to set the target score with the content-based neural activation maps of the two mixing images, generated by a pre-trained teacher network. Our intuition is that not all regions correspond to the foreground object. Concretely, the regions with rich semantic information would have a bigger impact on the target score than other regions. Inspired by the distillation technique that sets the target score of an image by a teacher network, we extend the design to set the target score by combining a teacher network’s neural activation maps of the two mixing images. As shown in Figure 2, the target scores of two mixing regions are calculated as the summation of spatially normalized neural activation maps within the mask for $x_{a}$ or outside the mask for $x_{b}$. We then combine the two target scores as the final target of the mixed image.

Compared to previous arts [37, 34, 18, 17], our proposed TokenMix has two main advantages: 1) We explicitly encourage the transformer to better encode long-range dependency to correctly classify the image with the other image mixed inside. We show that our approach can lead to consistent accuracy gain when used in various vision transformers, and also enhances the occlusion robustness of the transformers. 2) The target label of the mixed image that is generated with content-based neural activation map is more robust than those of previous approaches, which takes advantage of the distillation technique. Besides, we show that our approach promotes transformers to better localize the discriminate regions, with attention weights.

We use ImageNet-1K [10] dataset to demonstrate the effectiveness of our method. The dataset contains 1.2 million images for training and 50K for validation. The top-1 accuracy is reported as the evaluation metric.

We also use ADE20K [42] to verify the transferability of our TokenMix pre-trained models. ADE20K is a widely-used semantic segmentation dataset, covering 150 semantic categories. The dataset has 25K images in total, with 20K for training, 2K for validation, and another 3K for testing.

We evaluate our method on several recent vision transformer architectures, including DeiT [29], CaiT [30], PVT [35] and Swin Transformer [35]. We also test TokenMix on ResNet [13], which is representative of convolution models, as comparison. We follow the training recipe of DeiT [29]. The batch size is set to 1024. We use AdamW [19, 22] as the optimizer and set the learning rate as 0.001 with 5 warm-up epochs. The learning rate is decayed following a cosine scheduler down to $10^{-6}$. Without other specification, we train the models for 300 epochs. Rand Augment [8] and Mixup [39] are both used by default. Following [29], we switch TokenMix and Mixup with the probability of 0.5. For training architectures with smaller model sizes, e.g., DeiT-T [29], PVT-T [35], or CaiT-XXS [30], we sample the $\lambda$ in Equation 2 from a beta distribution Beta$(1.0,1.0)$. We use binary cross-entropy (BCE) loss instead of the typical cross-entropy (CE) loss by default following [36, 2], as the mixed images are more likely to contain multiple labels. To generate the neural activation maps, we defautly use NFNet-F6 [3] following [16].

For transferring to the ADE20K dataset, we follow the setting in BEiT [1], and fine-tune for 160K steps with Adam [19] optimizer. The detailed hyperparameters are described in supplementary materials.

We report the results on ImageNet-1K dataset with our TokenMix. As shown in Table 1, TokenMix consistently improves CutMix on various transformer-based architectures, i.e., DeiT [29], PVT [35], CaiT [30], and Swin Transformer [21]. Specifically, TokenMix outperforms CutMix [37] by +1% for DeiT, across DeiT-T to DeiT-B. We also improve popular hierarchical transformer architectures Swin-T and PVT-T for +0.4% and +0.5%, respectively. All the results demonstrate the effectiveness and generalization of the proposed TokenMix.

Our proposed TokenMix consists of two parts, i.e., token-level mixing and label refinement. We decouple the two parts and then compare them with the previous methods by fixing one part. In Table 5.1, we compare TokenMix to ReLabel [38] and TokenLabeling [16] with the same data augmentation method. The two methods utilize pixel-level supervision, but our TokenMix summarizes neural activations to create image-level target scores and is, therefore, more robust to individual pixel-level errors. Note that we use the same teacher network, i.e., NFNet-F6, to generate the offline targets. As shown in Table 5.1, TokenMix outperforms both ReLabel (+0.5%) and TokenLabeling (+0.3%) with the same training cost. We further compare TokenMix to previous mixing-based augmentation methods in Table 5.1. For a more fair comparison, we only use the labels from ImageNet. As shown in Table 5.1, TokenMix have performance advantages compared to other approaches. We see that the methods that introduce more foreground regions fail to improve CutMix on Vision Transformer. In contrast, our proposed TokenMix improves CutMix for +0.5% accuracy.

Pre-training on ImageNet-1K then finetuning to the downstream tasks is a common practice for many visual recognition tasks. It is important to verify whether the better pre-train with TokenMix can boost the performance on the downstream task. To do so, we transfer our TokenMix pre-trained models to the semantic segmentation task and compare them with regular pre-train. Note that TokenMix does not introduce extra computation overhead in the transfer stage. As shown in Table 4, we find that better pre-train from TokenMix consistently improves the segmentation performance on the ADE20K dataset. Notably, we improve DeiT-T for +0.7% mIoU, DeiT-S for +2.2% mIoU, DeiT-B for +0.5% mIoU. We notice that the performance gap becomes even larger (e.g. +1.1% mIoU for DeiT-T) when using multi-scale testing. All the results demonstrate the transferability of our TokenMix pre-trained models.

Besides the performance gains, we find that our proposed TokenMix improves transformers to be robust to occlusion, and focus more on the foreground area. All the visualization and analysis are conducted on DeiT-S.

We visualize some examples in ImageNet-1K (Figure 8) to show the disagreement between the salient areas and the foreground regions of the target class.