{'fraction_non_alphanumeric': 0.036595782942948345, 'fraction_numerical': 0.008047475610304366, 'mean_word_length': 4.920681394729837, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 1, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
TokenMixup: Efficient Attention-guided Token-level Data Augmentation for Transformers

Mixup is a commonly adopted data augmentation technique for image classification. Recent advances in mixup methods primarily focus on mixing based on saliency. However, many saliency detectors require intense computation and are especially burdensome for parameter-heavy transformer models. To this end, we propose TokenMixup, an efficient attention-guided token-level data augmentation method that aims to maximize the saliency of a mixed set of tokens. TokenMixup provides x15 faster saliency-aware data augmentation compared to gradient-based methods. Moreover, we introduce a variant of TokenMixup which mixes tokens within a single instance, thereby enabling multi-scale feature augmentation. Experiments show that our methods significantly improve the baseline models' performance on CIFAR and ImageNet-1K, while being more efficient than previous methods. We also reach state-of-the-art performance on CIFAR-100 among from-scratch transformer models. Code is available at https://github.com/mlvlab/TokenMixup.


Introduction
Various data augmentation methods have been proposed for computer vision tasks. One of the most successful augmentation methods is mixup [1], which is commonly applied to image classification. It attempts to augment the input by taking a convex combination of two random instances, and by reassigning the ground truth label correspondingly. Building on mixup, many subsequent works [2][3][4][5][6][7][8][9] have been introduced. They focus on mixing instances in a more meaningful way (e.g., saliencyaware mixup, manifold level mixup, submodular diversity maximization etc.), while a great majority attempts to perform mixup based on saliency. These saliency-aware mixup methods often entail gradient computation. To detect salient regions, the input is forward-propagated once and back-propagated to extract the gradient map. However, such a mechanism tends to be computationally heavy. This is especially pertinent to transformer-based models [10], as they generally retain high parameter volume. Keeping this in mind, we aim to present a saliency-aware mixup method favorable for transformers. Specifically, we regard self-attention as the inherent saliency detector, which serves as an efficient means for augmentation.
We accordingly propose TokenMixup, an efficient attention-guided token-level mixup method. Its objective is to mix intermediate token sets so that the saliency level of a batch is maximized. To detect salient tokens, we take advantage of the attention map to achieve a ×15 speed-up compared to the gradient-based saliency estimator, without compromising model performance. The batch saliency is maximized by optimally matching instance pairs, accomplished by an algorithm that exactly solves the optimization problem. We also introduce ScoreNet, a simple module that measures the difficulty of an instance. Based on ScoreNet, TokenMixup is applied selectively with respect to sample difficulty, which can be viewed as a type of curriculum learning method. To the best of our knowledge, this is the first attempt in the literature to discuss mixup from a curriculum learning perspective. Furthermore, we present a novel mixup approach that combines intermediate tokens across transformer layers within a single instance. We name this Vertical TokenMixup (VTM) since it vertically mixes tokens to provide rich multi-scale information to the attention layer.
Experiments on image classification datasets imply that our methods are both effective and efficient. Also, an advantage of TokenMixup is that it can be used in parallel with other mixup methods. In combination with those methods, we improve our baselines' performance by significant margins on CIFAR-10, CIFAR-100, and ImageNet-1K. Especially for CIFAR-100, we achieve a new state-ofthe-art performance among transformer-based models that do not use pre-trained models.
Then, our contributions are fourfold: • We propose TokenMixup, an efficient attention-guided token-level mixup for transformers. By incorporating the self-attention map, we achieve ×15 faster saliency-aware data augmentation compared to gradient-based methods. • We incorporate curriculum learning into mixup, to adaptively perform TokenMixup in an intermediate layer based on the confidence score of ScoreNet. • We also studied a variant, Vertical TokenMixup, which performs mixup with a single sample and enables multi-scale feature augmentation. • We achieve state-of-the-art performance on CIFAR-100 among transformer models trained from scratch. We also gain significant improvements over baseline models for image classification.

Related Works
Mixup. Input Mixup [1] is a data augmentation method widely adopted for image classification. The classification model is trained with a convex combination of input images and labels. A special case of Input Mixup is CutMix [2], which can be seen as a pixel-wise mixup method with binary masks. Recent advances in these mixup methods mainly focus on appropriately utilizing saliency information, and on mixing in image feature levels.
The core motivation of saliency-based mixup is that salient regions should be preserved when mixed, to retain a sufficient amount of information and learn more natural feature representations. SaliencyMix [5] adopts various saliency detectors to directly extract salient regions. Puzzle Mix [3] attempts to mix images using saliency information while retaining local statistics. Co-Mixup [4] maximizes gradient-based saliency while encouraging supermodular diversity of the mixed images. SuperMix [9] takes advantage of supervised signals to mix input images based on saliency.
On the other hand, Manifold Mixup [6] has provided theoretical grounds on the advantages of mixing images in the higher level of features. Manifold mixing helps the model learn flatter representation and smoother decision boundaries. Moreover, it is noted that it enhances robustness to adversarial attacks. Recently, AlignMix [7] was proposed to further align geometric properties of objects by mixing instances in the feature space. Furthermore, StyleMix [8] proposed to manipulate content and style information disparately to generate more abundant and robust samples.

Methods
Here, we introduce TokenMixup, a simple but effective mixup method for transformer models. The goal of our method is to augment intermediate tokens while maximizing the saliency level. This can be accomplished in three major steps: 1) Sample difficulty assessment, 2) Attention-guided saliency detection & optimal assignment, 3) Token-level mixup. The following subsections provide detailed descriptions of the steps.

Sample Difficulty Assessment
We first assess the difficulty of the input to adaptively decide whether to apply augmentation. To evaluate sample difficulty, a measuring function F needs to be defined. We utilize a parameterized module, ScoreNet, which is a simple MLP that predicts the target value ( ) based on the intermediate tokens ( ) , where is the index of the sample in a batch (See the supplement for figure). Given the ScoreNet output, the difficulty score is evaluated with the prediction loss computed as If the score is greater than a threshold , the sample is deemed sufficiently hard and no mixup is performed. On the other hand, the set of easy samples with F ( ( ) , ( ) ) < denoted as˜will be augmented by Token-level mixup after pairing with other samples in a mini-batch, as illustrated in Figure 1. Note that ScoreNet is trained simultaneously with the main model and training ScoreNet to assess sample difficulty at the intermediate layer can be viewed as an auxiliary task.

Attention-guided Saliency Detection & Optimal Assignment
Several works [19,[27][28][29] have discussed the imbalanced information of tokens. Due to this imbalance, random token mixing would potentially cause significant information loss and meaningless token replacements (See the supplement for relevant analysis). Thus, we aim to mix tokens based on saliency.
Attention-guided saliency detection. Instead of the computationally heavy gradient-based saliency detectors [3,4], we take advantage of the attention map, an inherent saliency approximator within transformer modules. Specifically, we can infer the saliency of the tokens in the ℎ layer via Attention Rollout [30], which computes the attention imposed on the tokens from layer to where is the number of heads in the multi-head attention layer, and Φ ( ) ℎ is the ℎ ℎ attention head of layer . In order to retrieve , the samples need to be propagated for the subsequent ℓ layers with stop-gradient. Since each additional forward propagation is an overhead, we approximate Attention Rollout by setting ℓ = 0. That is, we only use the attention map of the following layer as the saliency estimator. We find this a sufficient approximation of the full Attention Rollout. See section 5.1 for analysis on the soundness of the approximated saliency map. Then, the saliency score can be computed as
Optimal assignment. Based on the estimated saliency of tokens, we aim to maximize the overall saliency level by optimally assigning a different mixup target for each instance. So, we first define , ∈ IR n as the saliency difference between a random instance pair ( , ), which is computed as where , = 1, 2, . . . , and ∈ IR n is the token saliency map from the ℎ instance. Then, we represent our objective as an optimization problem written as where M is the set of all possible batch permutations of the instances, and ∈ M refers to an arbitrary permutation. is the threshold hyperparameter that controls the minimum saliency gain required for a token to be mixed, and is a binary decision vector for the ℎ instance. i.e., token of is replaced with token of ( ) when , = 1, and preserved when , = 0 ( = 1, 2, . . . , ).
The above optimization problem can be exactly solved by utilizing the Hungarian Matching algorithm [31]. The matching algorithm requires the score matrix, ∈ IR b×b , which is computed as where each item in , refers to the maximum saliency gain resulting from mixing˜( ) and ( ) . The actual maximum gain will be reached when = 0, but we set a positive threshold to control the minimum saliency gain required for each token to be replaced. Then, by applying the Hungarian algorithm, we can find the optimal batch permutation * such that Then, by mixing˜( ) and ( * ( )) , our objective is optimized.
By incorporating sample difficulty assessment (section 3.1) into (5), our final objective is written as where F is the difficulty measurer, and is the difficulty threshold. 1 denotes the indicator function, while 1 is a one vector. If the sample is hard, i.e., F ( ( ) , ( ) ) ≥ , the constraint becomes 1 ≤ 0, enforcing all elements of to be 0. In that case, none of the tokens will be mixed. On the other hand, if the sample is easy, i.e., F ( ( ) , ( ) ) < , the constraint is 1 ≤ which is satisfied by any random vector. Then, the objective becomes identical to (5). are brought up to form the key and value tokens. By utilizing a cross-attention mechanism, the dimension is preserved while being able to aggregate multi-scale information for token augmentation. The figure demonstrates the case when VTM is applied to layer 3.

Token-level Mixup
Now, the remaining question is how the tokens of˜( ) and ( * ( )) are going to be mixed. As defined in objective (5), we use a hard replacement of tokens with respect to decision vector . Then, mixing is performed by replacing token index from the original instance , with token from the paired instance = * ( ), so that ( ) , > is satisfied. In other words, only the token indices that increase saliency level greater than are mixed. This is accomplished by defining an appropriate binary mask ∈ IB n (t = 1, 2, . . . , n) containing where˜( ) is the saliency score of˜( ) . Then, the mask is employed as where operator denotes element-wise multiplication and˜( ) , ( ) refer to the paired instances. We also reassign labels based on the replaced tokens' saliency score. The new mixed label is computed as˜( Consequently, relatively unimportant tokens will be replaced with salient tokens from another instance, and its label will be adjusted with respect to the change in overall saliency level. See Figure 1 for visualization, and Algorithm 1 for overall computation process. More detailed pseudocode is provided in the supplement.

Vertical TokenMixup
Mixing multiple instances is not the only way to mixup. We also present a simple variant of TokenMixup, which provides token-level augmentation within a single sample by utilizing the tokens from previous layers. Similar to (5), we can define a new objective for Vertical TokenMixup as where L is the set of previous layer indices, and ( ) returns an arbitrary layer index for the ℎ instance. This objective can be optimized by a similar scheme as in Algorithm 1. However, considering that the tokens from each layer with an identical index tend to contain similar information, saliency difference matrix , ( ) = ( ) − (as in (4)) will most likely be constant across indices. This may lead to meaningless mixup without much increase in saliency level. Therefore, we take a different approach to mix tokens vertically.
The simplest method would be to concatenate all tokens from previous layers and apply self-attention. But due to the quadratic complexity of self-attention, naive concatenation may induce excessive computation. To reduce overhead, we selectively pool the most salient tokens from each previous layer and concatenate them to the original token set. Also, to preserve input dimension, we adopt the cross-attention mechanism. If 1 is the original token set and 2 is the concatenated token set, vertical cross-attention is expressed as where , , denote the projection layers for query, key, and value, respectively. See Figure 2 for visualization, and the supplement for pseudocode.

Discussions
TokenMixup as Curriculum Learning. A general framework for curriculum learning consists of two main components: Difficulty Measurer and Training Scheduler [32]. TokenMixup satisfies these conditions, as ScoreNet measures the difficulty of each input, and augmenting the tokens based on input difficulty naturally schedules training. In the early training phase where the model parameter is not optimized, most instances will be evaluated difficult and no TokenMixup is performed. As the model converges, on the other hand, many samples will become easy, triggering augmentation to make the task more challenging overall. Also, TokenMixup is applied discriminatively based on individual sample difficulty which enables more intricate curriculum scheduling. See section 5.3 for empirical demonstration. These properties will be especially useful when other mixup methods [1][2][3][4] are used in parallel. For instance, if input mixup is already applied, further augmentation may be unnecessary. In such a case, ScoreNet will regard the sample sufficiently difficult, and no further mixup will be performed.
Experimental setup. We evaluate TokenMixup and its variant, Vertical TokenMixup (VTM), on CIFAR and ImageNet with different settings. To avoid confusion, we denote the original TokenMixup as Horizontal TokenMixup (HTM) in the following sections. For CIFAR experiments, we adopt the 1500-epoch version of CCT. We modified the learning rate scheduler and positional embedding type to achieve better performance than original papers (denoted * in Table 1). Other experiment settings follow [22], and all experiments on CIFAR datasets were conducted on a single RTX A6000 GPU. For ImageNet-1K experiments, we used the ViT-B/16 model pre-trained and fine-tuned on ImageNet-21k and ImageNet-1k, where we took advantage of the officially released pre-trained weights. Experiments for Horizontal TokenMixup were conducted on a single NVIDIA A100 GPU, and 4 RTX 3090 GPUs were used in parallel for Vertical TokenMixup. Other experiment settings are reported in the supplement.

Experimental Results
First, we demonstrate the image classification performance on CIFAR-10 and CIFAR-100 when TokenMixup is applied to CCT. From Table 1, we can see that applying Horizontal TokenMixup and Vertical TokenMixup simultaneously improves its strict CIFAR-100 baseline, by a significant margin of 0.70%. HTM and VTM each surpasses the baseline performance by a margin of 0.69% and 0.67% as well. Also, we outperform the previous state-of-the-art model (among from-scratch models) by a state-of-the-art models that do not use pretrained models for initialization. Top-1 validation accuracy is reported for each model, and baseline models denoted with * reports retrained performance. Also, the number in the parenthesis indicates the number of epochs used for training, and the parameter numbers are retrieved from the CIFAR-100 models. Best performances are highlighted in yellow.

Analysis
Here, we analyze TokenMixup to answer the following research questions: Q1. Is the Attention an appropriate saliency approximation? Q2. What are the effects of each component in TokenMixup?

Soundness of Saliency Estimation
In TokenMixup, we estimate the saliency using the attention map from the subsequent layer. The soundness of this approximation would be one of the main questions regarding saliency detection. So, we provide relevant analyses on the appropriateness of our saliency approximation method to answer the research question Q1. Qualitative examples. Here, we provide a qualitative analysis on the performance of each saliency detector. In Figure 3, we can observe that there is no significant difference between the full attention rollout and the 1-step approximation, i.e., the subsequent attention map. Slight discrepancy in the sharpness of the distribution can be observed, but it does not greatly affect our method that utilizes saliency threshold . Moreover, attention-based methods seem to be more accurate than the gradientbased detector in certain cases. See the supplement for additional qualitative examples, and also the quantified comparison on the sharpness measures of the two methods. Efficiency comparison. We compare the efficiency of three saliency detectors: gradient-based, attention-based, and random baseline. For the gradient-based detector, we follow the method used in Co-Mixup [4], and random baseline refers to a random selection of tokens. In Table 3, we report the average latency of saliency detection per iteration in our CIFAR-100 experiment setting with a batch size of 128. Compared to the gradient-based method, attention-based saliency detection with 1-step rollout achieves approximately ×15 speed-up while demonstrating even better performance.

Component Analysis
In this section, we analyze the effect of each component in Horizontal TokenMixup and Vertical TokenMixup so as to answer Q2. In Table 4, we provide ablation studies on components of HTM: Sample difficulty assessment (ScoreNet), Optimal assignment, Token-level mixup, and Saliency-based label reassignment. By ablating ScoreNet, all instances are mixed regardless of sample difficulty; note, ScoreNet still receives supervision signals. By getting rid of optimal assignment, mixup pairs are not assigned via Hungarian matching; pairs are randomly selected. Token-level mixup and label reassignment is relevant to the method by which tokens are selected. Then, by ablating mixup, we randomly select tokens to replace, while label reassignment ablation derives new labels based on the number of tokens that have been replaced. We also conducted sensitivity tests on key hyperparameters , , and , but tables and discussions were moved to the supplement due to spatial constraints.

What TokenMixup Learns
Here, we answer our research question Q3 by analyzing what our mixup methods are learning. We demonstrate a curriculum learning perspective in terms of HTM, and compare our method with a direct mixup alternative, manifold mixup [6]. A Curriculum Learning perspective. In Figure 4 we plot the trend in the number of samples per batch TokenMixup is activated based on ScoreNet. In the first few epochs (regime (A)), the number of mixed samples is relatively low, as there is less need for augmentation in the early phase of learning. But as the learnable parameters warm up, the number increases abruptly. In regime (B), the average number of mixed tokens slowly increases as model performance steadily improves. In the last learning phase (regime (C)), mixup frequency converges.
Comparison with Manifold Mixup. TokenMixup is applied to the intermediate tokens, similarly to manifold mixup [6] which is an input mixup [1] method applied to intermediate features. By using manifold mixup, transformer tokens are mixed via linear interpolation, instead of the hard replacement of tokens as in TokenMixup. In Table 5, our methods are compared with manifold mixup, along with various combinations of original mixup methods (i.e., input mixup and cutmix). The best setting in the CCT paper adopted input mixup and cutmix. By applying HTM and VTM on top of that setting, we achieved state-of-the-art performance. On the other hand, by applying manifold mixup, we observed severe deterioration in accuracy. We conjecture that soft mixing of transformer tokens leads to over-smoothing of features, which may be inadequate for self-attention layers. See the supplement for visual aid.

Conclusion
We proposed TokenMixup for transformers, which adopts the attention map in place of computationally heavy gradient-based saliency detectors. By adopting the attention as the saliency map, we achieve ×15 speed-up compared to the gradient-based method with even better performance. We also introduce a novel perspective of curriculum learning to mixup, which enables adaptive augmentation based on sample difficulty and training schedule. To add diversity, Vertical TokenMixup is introduced, which mixes tokens from different layers for multi-scale feature augmentation within a single sample.
With TokenMixup, we achieve state-of-the-art performance on CIFAR-100, and improve baselines by significant margins on CIFAR-10 and ImageNet-1K.