TokenMixup: Efficient Attention-guided Token-level Data Augmentation for Transformers
Abstract
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 × 15 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–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.
Vision Transformers.
Originating from the natural language processing field, the Transformer [10] model enjoyed tremendous success in numerous computer vision applications, e.g. , image classification, object detection [11, 12], and human-object interaction detection [13, 14] et cetera .
Specifically, image classification is one of the most fundamental computer vision tasks.
Starting from the ViT [15], many works have attempted to assimilate the convolutional operation with transformer blocks [16–21].
Other variants have effectively incorporated the convolutional neural network into the transformer to achieve superior performance [22–26].
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.
3.1 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
[Equation]
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.
3.2 Attention-guided Saliency Detection & Optimal Assignment
Several works [19, 27–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 𝑖 + ℓ ( ℓ ≥ 0 ) as
[Equation]
such that Φ ( 𝑖 ) = 1 𝐻 (cid:205)
𝐻ℎ = 1 Φ ( 𝑖 ) ℎ 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
[Equation]
where 𝑡 = 1 , 2 , . .
. , 𝑛 denotes the token index.
Input
𝑋 ∈ IR b × n × d , 𝑌 ∈ IR b × c , 𝜏 , 𝜌
Output d
IR b c s.t.
𝑖 = 1 , 2 , . .
𝑏 ⊲ Eq.
⊲ Eq.
(3) s.t.
𝑖 = 1 ,
− ˜ 𝑆 ( 𝑖 ) 𝑡 − 𝜌, 0 ) s.t.
𝑖 = 1 , 2 , . .
. , 𝑏 ′ and 𝑗 = 1 , 2 , . .
𝑏 ⊲ Eq.
(6) s.t. 𝑚
= 1 , 2 , . .
′ and 𝑚 ) ∈ { 1 , 2 , . .
. , 𝑏 } s.t.
𝑖 = 1 , 2 , . .
𝑏 ′
⊲ Eq. ⊲ Eq.
(11) s.t.
𝑖 = 1 , 2 , . .
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 𝑃 𝑖,𝑗 △ = 𝑆
[Equation] − 𝑆 𝑖 , (4) 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
[Equation]
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
[Equation]
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
[Equation]
Then, by mixing ˜ 𝑋 ( 𝑖 ) and 𝑋 ( 𝜎 ∗ ( 𝑖 )) , our objective is optimized.
By incorporating sample difficulty assessment (section 3.1) into (5), our final objective is written as
[Equation]
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).
3.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
[Equation]
where ˜ 𝑆 ( 𝑖 ) 𝑡 is the saliency score of ˜ 𝑋 ( 𝑖 ) 𝑡 .
Then, the mask is employed as
[Equation]
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
[Equation]
[Equation]
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.
3.4 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
[Equation]
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
[Equation]
where 𝑊 𝑞 , 𝑊 𝑘 , 𝑊
𝑣 denote the projection layers for query, key, and value, respectively.
See Figurefor visualization, and the supplement for pseudocode.
3.5 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–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.
Experiments 4.1 Preliminaries
Baseline.
In experiments on CIFAR [33], we used Compact Convolution Transformer (CCT) [22] as our baseline.
For ImageNet-1k [34] experiments, we used the vanilla ViT-B/16 [15] as baseline.
Then, our augmentation methods are evaluated on three representative image classification datasets:
CIFAR-10, CIFAR-100, and ImageNet-1K.
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 A 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.
4.2 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
[Table]
[Table]
ours)
[Table]
huge margin of 0.85%.
In the case of CIFAR-10, we achieved higher accuracy of 97.57% and 97.78% with HTM and VTM, respectively.
Considering the highly saturated performance on this dataset, a 0.30% improvement over the baseline is not trivial.
Performance on ImageNet-1K is evaluated on the vanilla ViT [15] model.
The officially reported
Top-1 accuracy is 81.2%, which is achieved by using weights pre-trained on ImageNet-21K and fine-tuning on ImageNet-1K. By fine-tuning the model further by employing our augmentation methods, we achieve Top-1 accuracy of 82.37% with Horizontal TokenMixup, 82.30% with Vertical
TokenMixup, and 82.32% by using both.
Also note that HTM marginally increases parameter size by only 0.7M, while VTM does not require any additional parameters.
See the supplement for further experiments on other baselines.
[Table]
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?
Q3 .
What properties does TokenMixup have?
5.1 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.
5.2 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.
[Table]
[Table]
5.3 What TokenMixup Learns
Here
[Table]
we answer our research question Q3 by analyzing what our mixup methods are learning.
We demonstrate a curriculum learning perspective in terms of HTM
[Table]
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.
Acknowledgments and Disclosure of Funding
This work was partly supported by ICT Creative Consilience program (IITP-2022-2020-0-01819) supervised by the IITP; the National Supercomputing Center with supercomputing resources including technical support (KSC-2022-CRE-0100) and Kakao Brain Corporation