Sparse Token Transformer with Attention Back Tracking
Authors: Heejun Lee
Minki Kang
Youngwan Lee
Sung Ju Hwang
Venue: ICLR 2023 poster (2023)
Keywords: Token Pruning
Sparse Token
Attention Back-tracking
BERT
Vision Transformer
DynamicViT

Abstract: Despite the success of Transformers in various applications from text, vision, and speech domains, they are yet to become standard architectures for mobile and edge device applications due to their heavy memory and computational requirements. While there exist many different approaches to reduce the complexities of the Transformers, such as the pruning of the weights/attentions/tokens, quantization, and distillation, we focus on token pruning, which reduces not only the complexity of the attention operations, but also the linear layers, which have non-negligible computational costs. However, previous token pruning approaches often remove tokens during the feed-forward stage without consideration of their impact on later layers' attentions, which has a potential risk of dropping out important tokens for the given task. To tackle this issue, we propose an attention back-tracking method that tracks the importance of each attention in a Transformer architecture from the outputs to the inputs, to preserve the tokens that have a large impact on the final predictions. We experimentally validate the effectiveness of the method on both NLP and CV benchmarks, using Transformer architectures for both domains, and the results show that the proposed attention back-tracking allows the model to better retain the full models' performance even at high sparsity rates, significantly outperforming all baselines. Qualitative analysis of the examples further shows that our method does preserve semantically meaningful tokens.

Published as a conference paper at ICLR 2023
SPARSE TOKEN TRANSFORMERS WITH ATTENTION BACK TRACKING
Heejun Lee Minki Kang 1.3 Youngwan Leel:4 Sung Ju Hwang'
KAIST DeepAuto.ai"; AITRICSS ETRI" {aini, 2 ac.kr yw.leeBetri.
jjwwang82@kiist.ac.. ac.rr
ABSTRACT
Despite the success of Transformers in various applications from text, vision, and speech domains, they are yet to become standard architectures for mobile and edge device applications due to their heavy memory and computational requirements. While there exist many different approaches to reduce the complexities of the Transformers, such as the pruning of the weights/atenntions/toenn quantization, and distillation, we focus on token pruning, which reduces not only the complexity of the attention operations, but also the linear layers, which have non-negligible computational costs. However, previous token pruning approaches often remove tokens during the feed-forward stage without consideration of their impact on later layers' attentions, which has potential risk of dropping out important tokens for the given task. To tackle this issue, we propose an attention back-tracking method that tracks the importance of each attention in a Transformer architecture from the outputs to the inputs, to preserve the tokens that have large impact on the final predictions. We experimentally validate the effectiveness of the method on both NLP and CV benchmarks, using Transformer architectures for both domains, and the results show that the proposed attention back- bockktracking allows the model to better retain the full models' performance even at high sparsity rates, significantly outperforming all baselines. Qualitative analysis of the examples further shows that our method does preserve semantically meaningful tokens.
INTRODUCTION
Transformers have achieved huge success in various application domains such natural language processing (NLP) and computer vision (CV), obtaining state-of-the-ar performances on variety of tasks, and are now considered the de facto standard architectures for number of domains. However, Transformers require over few tera-flops per entry (Devlin et al., 2019; Dosovitskiy et al., 2021b) to compute, which is orders of magnitude larger than the computational cost for previous CNN and RNN architectures (Tan & Le, 2019; Bahdanau et al., 2015). To reduce such computational burdens of Transformmr models, previous works explored model compression methods such as distillation (Wang et al., 2020; Jiao et al., 2020), quantization (Han et al., 2015; Frankle & Carbin, 2018), and pruning (Zaheer et al., 2020; Guo et al., 2019; Kim et al., 2022; Rao et al., 2021; Wang et al., 2021). Pruning approaches for Transformers mostly aim to remove unnecessary model weights (Han et al., 2015; Frankle & Carbin, 2018), or the attentions (Zaheer et al., 2020), which could achieve a linear reduction in complexity. Token pruning can reduce the complexity of attentions and fully connected layers simultaneously by removing less relevant tokens for the target task (Kim et al., 2022: Goyal al., 2020; Kong et al., 2021). How can we decide which tokens prune then? Previous studies either compute the importance score of each input token as the average attention scores (Goyal et al., 2020; Kim al., 2022) (Figure 1). or learn the importance score of each token with an additional neural network (NN) at each layer (Rao et al., 2021; Kong et al., 2021).
However, previous token pruning methods have the following problem: they prune the tokens in the input sequence without explicitly evaluating the importance of each token on the final sequence representation and prediction tasks. This is because all existing works (Kim et al., 2022; Wang
'DeepAuto.ai, Scoul, South Korea

---

Published as a conference paper at ICLR 2023
Required Output Tokens
(Tpp( Threshold)
Layer 1
of Layer
Recursive Taken Pruning
Recursive Token Pruning
TruiLay Tinee
Layer 1
I
Prunes
Token Pruning
Moon
(Concrete)
A prroh
Da
Sort
(ART)
Token Pruning
65 B.
0.5
A
Layer!
Important Inppu Takens of Layer
Accumulate
prak
PrunsAAper
Endd
Approdvet
Feed-Forward
(CLS)
Steve
Love
Cat
Run ApproxNet For ABT Token Pruning
[CLS]
Steve
Luve
Cat
Token Pruning
Attention Back-Tracking Token Pruning
Image
Image
Figure I: Concepts. Comparison of token pruning method between feed-forward and attention back-tracking. The thickness of arrows represents the weight of attention probability. The left side is token pruning in feed-forward pass. The right side is token pruning with attention back-tracking.
et al., 2021; Goyal et al., 2020; Rao et al., 2021; Kong et al., 2021) compute the importance score of an input token at each attention layer while doing the forward pass from the input to the final output (feed-forward token pruning. Figure 1). Thus, they may prune out important tokens at earlier layers that are important for the final representation as well as the task loss. For example, in Figure the token Love in the layer | + 1 is pruned by the feed-for forward method, although the token has a high attention probability in the representation at the final layer 1 + 2 that is used by the classification task.
To tackle such a limitation of conventional forwafo token pruning methods, we propose an Attention back-tracking method for computing the importance score of each input token by their importance on the required output tokens (e.g.: sequence representation token, in last layer) and the task performance. As illustrated in Figure (right), we take backward pass from the last layers' token representations to input tokens, pruning the tokens by their importance score of each input token by accumulating its attention score. By doing so, we are able better select and keep the important tokens that will minimally affect the output and the prediction. We name our novel token pruning method as Sparse Token Transformer with Attention Back-Trackin
(STTABT).
However, one challenge here is that such backtracking requires us to know the attention score of each token before fow-rdin the input sequence to the model. To handle the issue, we introduce a lightweight attention approximation network trained with knowledge distillation called ApproxNet. Moreover, to actually prune the tokens, we need to decide which tokens to be retained based on the importance scores. The top-k method, by design, et al., 2020), but requires sorting the tokens by their importance scores, which is not differentiable. The thresholding method, which prunes the tokens with scores under the threshold value (Kim et al., 2022) does not require sorting, but setting the threshold can be tricky since it is difficult to know how many tokens will be remained in advance with the given threshold. To remedy such issues in both methods, we propose a learnable and smooth threshold function named Concrete masking, inspired by Concrete Dropout (Gal et al., 2017). Specifically, we jointly train the threshold value for each layer and importance score which is computed by attention back-tracking, with the task objective to find the thresholds that can minimize the task loss.
select predefined number of tokens (Goyal
Our method is generally applicable to Transformer encoders for any domains, and we validate our method on text and image classification tasks. For text classification, we validated the proposed token pruning methods on GLUE (Wang et al., 2019) benchmark with BERT (Devlin et al., 2019). For image classification, we validate it on ImageNet-IK IK (Deng et al., 2009) benchmark with DeiT (Touvron et al., 2021). The experimental results show that our model works surprisingly well even at very high pruning rates, for which baselines show significant performance degeneration. For example, on GLUE benchmark task (QQP), our method obtains 45.54% increased token sparsity (17.15% to 9 9.34%) compared to the bascline with the same accuracy. Moreover, our method obtains very high token sparsity (18.77 7%), keeping only half of the tokens retained by Dynamic ViT (37.5%), with only 0.8% accuracy loss on ImageNet-1K 1K benchmark. Our contributions can be summarized as follows:
We propose novel token pruning method for Transformers based attention back-tracking, which considers the importance of each token on the final representations as well as the task loss based on the approximated attention from the distilled model.
We propose Concrete masking, which automatically learns the pruning rate for each layer with Concrete dropout. (Gal et al., 2017).
2

---

Published as a conference paper at ICLR 2023
Our token pruning is generally applicable to any Transformers, and we validated it trained Transformers from both NLP (BERT) and CV (DeiT) domains, on GLUE and ImageNet benchmark, whose results show that achieves significantly better recuracy-compputtion tradeoff. (Kim et al., 2022; Rao et al., 2021).
2 RELATED WORK
Attention Pruning for Transformers Since the literature the general neural network compres- sion techniques such as quantization and weight pruning is vast, we discuss compression techniques specifically targeting Transformers (Vaswani et al., 2017). Due the heavy computational cost of the self-attention, several recent works have focused to reduce its time complexity from O(N²) to about (LV) given the N length of the sequence by pruning the attentions the self-attention layers. BigBird (Zaheer et al., 2020) is one such method, and it utilizes set of sparse attentions, that are both local and random, while using few global attentions, to deal with longer sequences. Similarly, Star- Transformer (Guo et al., 2019) leverages the concept of a virtual relay node and reduces the number of fully-connected attentions with the radial and ring connections. However, these prior works do not consider the content of the input sequence, as they make use of static sparsity patterns. To remedy this, some recent works aim to dynamically sparsify attentions of the Transformer considering the input data. Specifically, Reformer (Kitaev et al., 2020) divides the tokens into several groups through Locally Sensitive Hashing (LSH) and calculates attentions by sparsely connecting them. Routing Transformer (Roy al., 2021) uses the k-means clustering algorithm to cluster input tokens.
Token Pruning for Transformers However, in most datasets used in Natural Language Processing (NLP), such as text classification (Wang et al., 2019), input sequences with with excessively long lengths are quite rare, as the length of most sequences is limited to 256 tokens (Kim al., 2022). Moreover, fully connected layers require significant computations that are non-negligible, although they have been relatively overlooked. Therefore, some recent works (Kim et al., 2022; Goyal et al., 2020) focus on reducing the number of tokens by selecting the most important tokens for the given task, to reduce the computational cost of both linear and self-attention layers. Notably, POWER-BERT (Goyal et al., 2020) progressively eliminates tokens based on their importance scores, computed from the attentions at each layer the Transformer encoder, and Learned Token Pruning (LTP) (Kim et al., 2022) aims learn to threshold the importance scores prune the token. However, both methods prune the tokens during the feed-forward pass, assuming that important input tokens also important output tokens. Therefore, they may prune out tokens that are important for the output vector at earlier layers. Compared to them, our method STTABT can consider the impact of each token the output vector and the prediction task by attention back-tracking, thanks to novel lightweight attention approximation network.
SPARSE TOKEN TRANSFORMER WITH ATTENTION BACK TRACKING
We start by briefly describing the Transformer model that our method focuses Then, we propose two components for our Sparse Token Transformer with Attention Back Tracking (STTABT): Ar- tention Back Tracking (ABT) which obtains the importance score for each token with an attention approximation network and Concrete masking for dynamically learning the token-wise threshold for each input. STTABT is post-hoc token pruning method and assumes that the target Transformer encoder has previously been optimized on the training dataset for the target task. Our goal is to lower the inference cost of the given pretrained Transformer by pruning the tokens at each layer.
3.1 TRANSFORMER
For concepts used in later sections, we define the self-attention layer - the fundamental element of the Transformer architecture (Vaswani et al., 2017)- in this section. Let us denote the attention score and probability at 1-th Transformer layer as follows:
Q'K.T dah
A' score
A prob = softmax(A* score).
(I)
3

---

Published as a conference paper at ICLR 2023
where Q' IR fxde is the query matrix, K' E Rixdd is the key matrix, de is the dimension of query and key matrix, and E [1, L] is the layer index.
3.2 ATTENTION BACK-TRACKING (ABT)
We now describe the details of our method, Attention Back Tracking (ABT). The fundamental idea of ABT is to prune the tokens based on their importance scores that measure their impacts on the final sequence representation for task prediction, by back-tracking the attentions. Specifically, as shown in Figure| (right), we start with the tokens with higher attention probabilities from the last layer, and backtrack their accumulated attention probabílities consisting of T, layers with regards to the token for the sequence classification task (specifically, [CLS] token). However, to apply such a technique into practice, we must rerdfforward the input sequence through all of the Transformer layers in order to compute attention scores beforehand. Thus a straightforward implementation of the technique could be quite inefficient because will require us to perform forward pass of the model twice for each input.
the first layer of the Transformer encoder
Attention Approximation Network To alleviate the cost, we use an attention approximation network (ApproxNet), distilled Transformer encoder that approximates the attention probabilities of the original network with a significantly less number of parameters. We perform knowledge distillation based on the distillation loss from TinyBERT (Jiao et al., 2020) build the ApproxNet. Specifically, we leverage embedding based distillation loss Lemba, hidden state based distillation loss Chidn. attention- based distillation loss Lattn and prediction-layy the distillation. Especially, for the attention-based distillation, we use the KL divergence loss from MiniL. (Wang et al., 2020) instead of the mean squared error loss as follows:
distillation loss Cpred for
Lattn HTT TIEEE
(A'roob (log Apros log A' A'pros)) broo)))a.s
(2)
HT (=1 H=
where A Apros E HxTXT denotes 1-th layer attention probability of the original network (teacher), A prob RHxTTTT denotes I-th layer attention probability of the ApproxNet (student), H is the number of heads for multi-head self-attention layers, T is the length of input sequence, and L is the number of layers in the Transformer network.
Algorithm 1: Update token mask from the output token indices def tpddat
Layer
ApproxNet
( (Aprobs 44+1 8 1+1 k)
Layeri
Layentvi
// Reduce attention heads and gather required output tokens' attention probabilities with previous importance score.
Tokenn
Prun
A + gather (mean (Arro dim-0), gl+1, dim-0)
2
3 a + mean (diag(s" dilag(s'+1) +*)+***=+++ A+ p), dim=0)
A
+ unique (cat (441 topk (a, )))
[0]
[1.0]
+ gather (a, 8
Selected by Required Output Token Index
+ s² sum (8')
Weighted Average (p)
return '
Top-k 110
volue
Figure 2: Computation flow of update token mask of whole Trans- former encoder layers with attention back-tracking by Algorithm
[0,2]
00663, 0.37]
Token Pruning with Attention Back-tracking on Approximated Attention Then, using the approximated attentions from the ApproxNet, can prune the tokens at each layer of the main network. In Algorithm 1, we define the pruning function Updat e Tokeeeakkkfoaaaaaaaaaaaa Layer, which updates the mask each Transformer layer for token pruning. The input arguments of the pruning function are the approximated attention probability Arroos of (-th layer from the ApproxNet, the token mask +1+1 E 2 and the importance scores of the tokens [+1 E IR from the upper layer / 1, where is the first layer of the Transformer model and 0++1 is the number of selected tokens in (I + 1) layer. jt is initialized with indices of required tokens, which is [CLS] token in the classification task. s Rer is initialized as the one vector 1/oT where or is the number of output
4

---

Published as a conference paper at ICLR 2023
tokens, which is generally a single token. p (0, C R determines how much importance scores should be propagated from the upper layer. If = 0, the importance scores from the upper layer will be completely ignored when pruning the tokens in (-th layer. & E Z determines how many tokens should be selected in case of using top-k selection method, and we could easily change the number of tokens retain at each layer, by adjusting k. The function returns the index of remaining input tokens at I-th layer and their importance scores given those arguments.
3.3 CONCRETE MASKING
The top-k tokens at each layer of the Transformer model may still be selected using the previous scheme's static value of k, but in order to regulate the pruning based on the input, mechanism for dynamically altering the pruning ratio is also necessary. Therefore, previous work, such LTP (Kim al., 2022) proposed to learn layer-wise threshold for token pruning. In detail, during fine-tuning, LTP keeps tokens that have larger scores than the threshold, which is jointly optimized with the network parameters. However, it can be challenging to control the number of tokens left at each layer when the importance score is relatively sparse and the majority of score values are close to zero or very small values. To address the problem, suggest novel dynamic thresholding method named as Concrete Masking, inspired by the learnable dropout function idea in Concrete dropout (Gal et al., 2017).
Let us consider single row of one of multi-head attention score matrix a² = Ascoreh, R in this section for brevity, where is the head index and is the row index. We further assume that the attention scores follow normal distribution. We transform a, into the uniform distribution where ar,i E [0, 1] so that the importance score becomes uniformly distributed. Formally,
a' Ho
a
à'
Å'
(3)
Ja
a'p
where @ is the cumulative distribution function of the standard normal distribution, Paa Et= d
is the mean of a', and aa V Pa) / is the standard deviation of the a'. Then, we apply above transformation to all the rows of the predicted attention score matrix from ApproxNet A! score acquire A, where all of rows uniformly distributed. Similar to the line 3 of Algorithm 1, we then propagate the transformed importance score a as follows:
T
(pddiag(s +11) + +1-1 p) (iag(m++))At
b'
b'
s max b +c)
(4)
mm
where the token mask m I+1 E (0, 1)Th C Rº is the Concrete mask vector. Then, the updated Concrete token mask m' is defined as follows (Gal al., 2017):
(log - log(1 - pot) + log s¹ log(1 - ))
m' = max m H+1 sigmoid
(5)
where sigmoid is the sigmoid function, F is the temperature hyperparameter, and pæ = sigmoid (0') where of is learnable parameter and initialized with 0,it- During training, the model learn the proper value of 0' with the soft mask function applied to each layer. After training for reproducing change the soft mask into hard mask, then fine-tune
the base accuracy with sufficient epochs, the parameters to fit the hard mask as in LTP (Kim et al., 2022). For hard masking, we threshold the soft Concrete token mask m' using set threshold, usually 0.5.
To train our Concrete masking mechanism, we need regularization terms in addition to the task loss. First of all, we add the regularization term Lp, which prevents 6' from deviating too far from the initialized value and another regularization term Lmask so that the average token retention ratio stay around Pr as follows:
E10
(111)) TL
L
T
2
- Ap
(0'
),
Lmask Amask
m
Po
(6)
1-1
1=1
t=1
where Ao and Amask hyperparameters.
5

---

Published as a conference paper at ICLR 2023
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP
To show the general applicability of our method, STTABT, we validate it on both Natural Language Processing (NLP) and vision tasks. In order to evaluate our method, we primarily compare the performance of our method against those of the baselines at the same token retention ratios, which is metric that measures how many tokens are retrained after pruning. We use the average token retention ratio values across all the layers of the target Transformers. To accurately measure the efficiency of our method, even considering the overhead for computing the ApproxNet for attention approximation, we also report the performance as a function of FL.OPs in Section 5.
4.1.1 PRUNING TRANSFORMER BBASED TRAINED LANGUAGE MODELS
Dataset and metric. We use nine datasets from GLUE (Wang et al., 2019) benchmark for the text classification and the BERT (Devlin et al., 2019) the base model. For evaluation metric, we use the metrics from the original paper (Wang et al., 2019). Baselines. We compare STTABT against two notable token pruning baselines- the manual top-k method and Learned Token Pruning (LTP) (Kim et al., 2022). The manual top-k method is the feed-forward token pruning which keeps the k tokens with highest importance scores for each /-th layer. LTP is a baseline with learnable threshold, which requires to fine-tune the whole parameters to learn the appropriate threshold value for each input. We match the training settings of STTABT and LTP, then follow the hyperparameter search strategies described in (Kim et al., 2022). For the comparison against LTP, we use STTABT with Concrete masking strategy ($3.3), while use top-k pruning with our method ($3.2) when comparing against the manual top-k baseline.
4.1.2 PRUNING VISION TRANSFORMERS
Dataset and metric. We use ImageNet- 1k (Deng et al., 2009) for the image classification experiment, using the Data-efficient Image Transformer (DeiT) (Dosovitskiy et al., 2021a; Touvron et al., 2021) as the base Vision Transformer model. Specifically, use DeiT-small as the base architecture for all experiments on the image classification. Baselines. We use Dynamic ViT (Rao et al., 2021) 1) as the primary baseline. We follow the experimental setups, including data preprocessing and training pipeline of Dynamic for all experiments on the image classification task.
4.2 EXPERIMENTAL RESULTS ON NLP TASKS
Attention back-tracking. We measure the performance of our attention back-tracking (ABT) method with the ApproxNet ($3.2). We experiment with different k in Algorithm I to vary the average token retention ratios. The experimental results in Figure show that the our method significantly outperforms the manual top-k with the same average token retention ratio. Furthermore, we find that using the ApproxNet (blue lines) obtains similar performance to using the attention from the full network (green lines). These results support that our attention backtracking strategy can identify more important tokens than the feed-forward pruning methods do, which compute the importance scores and perform pruning with forward passes.
Concrete masking. In Figure 3, we further present the performance evaluations with or without the Concrete masking, examine its effect. One of the advantages of our model with Concrete masking is that the user can flexibly adjust the token retention ratio by changing 0init value. In contrast, the user cannot easily adjust the token retention ratio of LTP since the method does not support the control of the token and relies on regularization on the mask, which is the token retention ratio with a regularization scaling factor. For the lower token retention ratio of LTP, we must increase the scaling factor, then the regularization term becomes significantly bigger than task loss some point. As a result, the LTP token retention ratio converges to zero or an extremely low ratio which is not appropriate to deal with task loss. For instance, in QNLI, MRPC, CoLA, and SST-2 of Figure 3, LTP baseline cannot cover the low token retention ratio, whereas our method can.
At relatively high token retention ratios around 30-50%. both STTABT with concrete masking and LTP work well, although our method still outperforms LTP in most cases. However, at lower than 30% token retention ratios, our method with Concrete masking show significantly superior performance

---

Published as a conference paper at ICLR 2023
MNU
QNL
MRPC
50
75
20
50
60
10
20
40
50
60
70
Auenge Thben
180)
Average Thien Retettion (8))
Auerage Takrnn Retertion Ratio
CoLA
SST-2
QQP
50
20
10
20
60
50
60
70
60
60
20
40
Avanage Takan Ratendon Raloo 156)
Average Iokien Rotention Ratio (96)
Asarage lakon Retertion Ratio (9))
WILLI
S150
RTE
0.40
Pearsin Corraerion
62
55
30
60
20
23
Avtrage Token Retendoon Radio 196)
Average Tiken Retention Radio (9))
Aucrage Yoken Retention Ratio (95)
STIART (Approx, AlL.)
STTART [Concrate, with
LTP [Best valid)
BERT
STIABT (Actual ALL]
STTAB (Concrete, wio train)
Manual top
Figure 3: Experimental results with varying average token retention ratio for our method STTABT (by Nep with factor 4) and baselines on GLUE and BERT buse- Our results with top-k pruning (blue, green) are compared against the manual top-k method (black dot). In contrast, results with Concrete masking (red, orange) are compared to the LTP (grey dot).
over LTP, better preserving the accuracy of the full model (Please see the results MNLI, QNLI, MRPC, CoLA, SST-2, QQP, and RTE). To examine which tokens are pruned with our method and LTP, we further visualize the pruning mask at each layer of BERT base in Figure 6. We observe that LTP keeps most of the important tokens, but unimportant tokens as well such as dots in Figure 6. On the other hand, our method with Concrete masking prunes useless tokens such dots much earlier layers than LTP, yielding higher sparsity in total.
The Concrete masking method have significantly higher accuracy in of cases in GLUE bench- mark than ABT without Concrete masking and baselines. The Concrete masking method works especially well at high pruning rates, while have competitive performance baselines at low pruning rates too.
STTAB (Cancrete)
[ayer]
layers
lapers
layers
Layer?
layer7
layera
layera
lapers
layer9
Layer 10
Layer 10
Layor 11
Layer 11
Laver 12
Laver 12
Figure 4: Token pruning visualization between our Concrete masking and LTP layer by layer. Layer (
means the input token mask the layer 1. Yellow color indicates the remained input tokens in each layer. The input texts are randomly selected from The Stanford Sentiment Treebank (SST-2) (Wang al., 2019). Similar accuracy of trained LTP and our Concrete masking models used (es 89%).
7

---

Published as a conference paper at ICLR 2023
ImageNet-1k
Original
Layer
Layer
Layer
Original
Layer
Layer
Layer
SITABT
Accurscy
DynemicvIT
START IConcrece, with traiy
-6- STTABT. (Concrete, maint
DynamiciT
Deliuse
STTABT
SITTATT (Apprex Al.)
STRE (Actual Alt)
-K-
Kanual lpp.
DynamicvIY
20
40
60
80
Average Tokon Ratention Ratio (8))
Figure 5: Top-1 accuracy plot vary- Figure 6: The visualization of token pruning with STTABT ing the average token retention (Concrete) and DynamicViT. Models with similar accuracy tio for our method and baselines on are used (Ours @4: 74.1%, Dynamic VIT: 74.9%). ImageNet-Ik.
Table 2: Relative FLOPs of various overhead sources of our methods. The FLOPs results are calculated with sequence length of 128. We randomly generate token retention ratios for each layer for simulation. Fused MulAdd for matmul is not considered in FLOPs calculation.
BERT
ApproxNet
Manual Top-k
Attention Back- BrackTracking
Concrete Masking
Reduce Factor
Token Retention Ratio
Retention Ratio
Token Retention
4
8
16
0.2
0.4
1.0
0.2
0.4
1.0
0.2
0.4
1.0
FLOPs Relative
22.4G
1.5G
436.9M
140.6M
559.5K
1.IM
2.8M
597.KK
1.2M
2.8M
9.2M
18.5M
46.IM
100%
6.86%
1.95%
0.63%
0.00%
0.00%
0.01%
0.00%
0.01%
0.01%
0.04%
0.08%
0.21%
4.3
EXPERIMENTAL RESULTS ON COMPUTER VISION TASKS
The experimental results show that our Attention back-tracking (ABT) performs competitively to Dynamic ViT even without Concrete masking, on computer vision tasks, and our full model with Concrete masking significantly outperforms Dynamic ViT 5, especially at lower token retention ratios. For example, our method with Concrete masking achieves 74.1%[[5... 5.8%p) accuracy with 18.7% token retention ratio, but Dynamic requires 37.5% token retention ratio with 74.9% 5.0%)) accuracy. Moreover, since our ABT requires training ApproxNet only once regardless of the target token retention ratio, has huge advantage in training cost than Dynamic The Dynamic ViT used twice more tokens than our Concrete masking, but the accuracy improvement is only 0.8%. Even considering attention approximation overhead, our model performs better than Dynamic ViT since the overhead is only about 1.95 6.86% of the base model, as shown in Table 2.
We visualize the pruned tokens at different layers the vision Transformers in Figure 6. The visual- ization shows that our model with Concrete masking concentrates the important tokens in earlier layer than Dynamic VIT. Moreover, our model preserves more important tokens while Dynamic attends larger number of tokens, including unimportant ones. For instance, Dynamic selects tokens with other attributes at Layer 4, while our model keeps only the tokens related to Chihuahua from the beginning of the model, thanks to the attention back-tracking in Figure 6.
COMPUTATIONAL EFFICIENCY OF STTABT
5.1 FLOPS COMPARISION
For the GLUE benchmark, we compare the attention back- tracking (ABT) method and the manual top-k method with respect to FLOPs. In Figure 7a, compare the performance with respect to FLOPs over the test set. We find that the ABT performs better than the manual top-k baseline on most datasets.
Table 1: Approximation emor of ApproxNet depending on initializing methods.
Random
Distilled
Dataset
Factor
Avenue
MSK
D.0122
0.0122
0.0013
0.0018
GLUE
Dis.
1.6430
1.6450
0.2019
0.2981
The ABT results look shifted toward the right than the Figure 3, because the computational cost includes overheads such as computing the mask update function (Algorithm 1) and the ApproxNet. We examine the relative overheads of each additional component in Table 2. We that ApproxNet
8

---

Published as a conference paper at ICLR 2023
Average GLUE
ImageNet-1k
80
belfs
100
SPVT--DIIT
DeiT-S/320
78
DelT-S/288
Top-1 Accuracy (%)
HVT:S-
99
De/7-9/256
Relative Metric
76
98
STTABT@rR (Approx. Art.]
74
ST (Actual All.)
STTABT@(4 (Approx. AIt.]
S (Actual Att.)
72
DelTT
DynamicviT
96
Manual Tap-k
STTABT (Concrete) De
S (Concrete) D
Relative FLOPs
GFLOPS
(a)
(b)
Figure 7: (a) FLOPs comparison between attention back-tracking with top-k and manual top-k baseline on BERT buse. (b) FLOPs comparison between our method with concrete masking baselines on VIT. plotted models except DeiT-T are based from DeiT.mad for fair comparison. STTABIS f means the attention back-tracking by ApproxNet with factor æ. DeiT-S/////ee means the DeiT smalt with hidden size (Kong et al., 2021).
is the major bottleneck, as the ApproxNet with factor 4 requires about 6.86% of the BERT base FLOPs, which is significantly larger than the manual top-k (0.01%). However, the overhead is more convincing with higher reduction factor of 8 (1.95%) than 4. Therefore in Figure 7a, the ABT results with factor 8 shows better result than 4.
We further compare the performance of STTABT with different reduction factors of the NetroxNet. Table I shows that the ApproxNet performs accurately approximations of the actual attention scores with a lower reduction factor, thus results in better preformance. However, when considering the accuracy- -ecurac-efficienny tradeoff, the performance of the model with the reduction factor of are often worse than the one with the reduction factor of 8.
In Figure 7b, we show the accuracy over FLOPPs comparison of our model against baselines on ImageNet-Ik (Deng et al., 2009), using DeiT Dynamic VIT (Rao et al., 2021) used as the primary baseline, DeiTsmat (Touvron et al., 202 I), SPVIT (Kong et al., 2021), IA-RED² (Pan et al., 2021a), S ViTE (Chen et al., 2021), and HVT (Pan et al., 2021b). We observe that the fine-tuned DeiT model pruned by our Concrete masking significantly outperforms the fine-tuned DeiT tiny model with the same FLOPs trained from scratch. DynamicViT with DeiT smalt. on the other hand, results in a model with much lower performance than DeiTring- This is good news for practitioners, this suggests that may not need to train Transformers at different scales from scratch, to deploy the model to systems or applications with different computational budgets.
the base model to prune. The baselines include
CONCLUSIONS
In this paper, propose sparse token Transformer architecture based novel token pruning method, Sparse Token Transformer with Attention Back Tracking (STTABT), with two novel ideas Attention Back-Tracing (ABT) and Concrete masking. In particular, the ABT method enables our pruning method to back-track the important tokens for classification from the topmost layer to the lower layers. while approximating the attentions using a lightweight approximation network (ApproxNet). Then, we propose a novel thresholding method named as Concrete masking, which adapts the idea from Concrete dropout learn the appropriate token-wise pruning threshold given the desired amount of tokens to be retained. We evaluate the performance of our method on classification tasks from NLP and CV domains, by pruning tokens from the trained Transformer encoder- based trade-off
language models and Vision Transformers. Our method shows a better uccuracy-efficient and yields Transformers with significantly higher accuracy for the same ratio of tokens retained, or at the same FLOPs even when considering the computational overhead from the ApproxNet. Qualitative analysis of the tokens retained by STTABT shows that it preserves relatively more important tokens while pruning out less relevant ones, compared to previous methods.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We will introduce code construction and data collection in this section for reproducibility.
Model Implementation. First we construct our model with PyTorch (Paszke al., 2019) and Huggingface (Wolf et al., 2020). We modify bert-baas--ncased from Huggingface Models (Wolf al., 2020). ApproxNet. We create ApproxNet ($3.2, trainer/ /glue base by) with the pretrained models. The approximation result is shown in Section A. ABT. We implement Algorithm on top of bert-base-uncase
(Devlin et al., 2019) (model s / s spa reerteken -.pee py). LTP, Concrete masking. Then, we modify forward path to support LTP (Kim et al., 2022) and fix AgproxSparseBertModee to implement Concrete masking ($3.3).
NLP training. The training for NLP tasks, we implemented our own training loop wwoine//concrete.rraae
and hyper-parameter tuning settings inher- its LTP hyper-parameter settings (Kim et al., 2022) Shain//tp..que.ploo
main/concrest glue plot. We use the GLUE dataset from Huggingface Dataset (Wolf al., 2020).
Vision training. We start from DynamicViT implementation (Rao et al., 2021) including data loader, training loop, training loss, and evaluation codes (main/vit__cncrren wew.eeeennddddddpem Therefore, we did not apply hyper parameter grid search on ViT pipeline. And we remove hard token mask training stage, and we use larger Amask 100 than NLP task, because of lack of padding. Also, we did not use prediction distillation loss for training ApproxNet for ViT. The training dataset ImageNet- (ILSVRC2012) (Deng al.. 2009). We replace the Transformer encoder in ViT model (Dosovitskiy et al., 2021b) into our model code. Our PyTorch model is modified to apply ViT changes from Transformer encoder (Dosovitskiy et al., 2021b)
ACKNOWLEDGMENTS
This work was supported by DeepAuto.ai. This work was partly supported by Institute of Information & communications Technology Plamning & Evaluation (ITP) grant funded by the Korea govemmment (MSIT) (No. RS-2022-001882338 Development of Large Korean Language Model Technology for Efficient Pre-training), (No.2022-0-001224 Development of Artificial Intelligence Technology for Self-Improving Competency-Awart Learning Capabilities), and (No.2019-0-0007) Artificial Intelligence Graduate School Program (KAIST)). This work was supported by the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF-2018RIASAION 10599221...
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: /ttp::://aaaav...a http/aaxvvvvvo.../... aas//aa////a
1
Tianlong Chen, Yu Cheng. Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chas- ing sparsity in vision transformers: zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 19974-19988, 2021. URL htt.ttpss://rooceediiie newrips wtcppaper//22///ee a61f27ab22 2 wwweee55ff0eeeeeeem 55522 Abstrs-bssrttt..... html.
end-to-end exploration. In Mare' Aurelio Ran-
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. 2009 IEEE Conference Computer Vision and Pattern Recognition, pp.
2009. ISSN 1063-6919.
doi: 10. 20.100//vv::200055206633
2, 6, 9, 10
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of
10

---

Published as a conference paper at ICLR 2023
the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis. MN, USA, June 2-7, 2019. Volume (Long and Short Papers), pp. 4171--188 Association for Computational Linguisties, 2019. URL httts:t: /ttp::://do..oooo 1tt/:/11.1111133
n1 9-1423. 1, 2, 6, 10, 14
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL ptttpp:://ppezeeeeee
net forum? idd
YicbF FNbTF. 6
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b. URL https: https:://ppppeeveeeeeee net forum?id= YicbF CNTTy. 1, 10
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Training pruned neural networks. CoRR, abs//803,mmm 880300363, 2018 URL http:://aaaxiii org/ otg/abs/////ooooo.//.././/...be/ 1803.800300 03635. 1
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Gamett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neu- ral Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 3581-3590, 2017. URL tttt:p::/ppocceedinge neuringggnnnnuusssseeeeeee...i neurips cc/ http/pppr/22111 / 84ddfb34126ff
Aostt7--bbsrrat..... html. 2,5
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan T. Chakaravarthy, Yogish Sabharwal, and Ashish Verma. Power-bert: Accelerating BERT inference via progressive word- vector elimination. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 3690-3699. PMLR, 2020. URL http: http::httocceeeeeeeeeeee pttp://eccceee.e.eeeeee ailr.press/vi99 goyal20a. html. 1, 3
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pp. 1315-133 1325. Association for Computational Linguistics, 2019. URL https: / http//o..ooo owro/ioorr/ii 10. 1 (8653vvvll
1133. 1,3
Song Han, Jeff Pool, John Tran, and William J. Dally. Learning both weights and connections for efficient neural networks. CORR, Ahs//1506.02266. 2015. URL http: / /atxiv.iivvvv.////v///.//aavai org / /bs/ 1506.020
Zihang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Xiaojie Jin, Anran Wang, and Jiashi Feng. Token labeling: Training 85.4% top-1 accuracy vision transformer with 56m parameters on imagenet.
CoRR, abs/2104. 10858, 2021. URL https: : arxiv. org/ abs////////// / 22024444 10858. 18
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling BERT for natural language understanding. In Trevor Cohn, Yulan He, and Yang Liu (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Findings of ACL, pp. 4163-4174. Association for Computational I Linguistics, 2020. URL. https: http//ooooooooo///////oooooo wtt..o/.......///..... 14
2020. findings eeaiinggseenllp 372.
Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Aidong Zhang and Huzefa Rangwala (eds.), KDD '22: The 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Washington, DC, USA, August 14. 18, 2022, pp. 784-794444779 ACM, 2022. URL https: /ddoi. org/10 100.0.115/33300 /33111155/3345 3539260300 60. 1, 2, 3, 6, 10, 15
11

---

Published as a conference paper at ICLR 2023
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020. URL https:://peeneeeeeeeeee net wwt./orunffffietoooooooo. forum?id-
rkgNKkHtvR. 3
Zhenglun Kong, Peiyan Dong, Xiaolong Ma, Xin Meng, Wei Niu, Mengshu Sun, Bin Ren, Minghai Qin, Hao Tang, and Yanzhi Wang. Spvit: Enabling faster vision transformers via soft token pruning. CoRR, abs/2112.13899 2021. URL https::/p/ /arxivi wwt.vvsrrraaaaoooo 2212 1 13890. 1,
Yanghao Li, Chao- Wu, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik, and Christoph Feichtenhofer. Mvitv2: Improved multiscale vision transformers for classification and detection. In IEEE/CVE Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 4794-4804. IEEE, 2022. URL https: //doi. wwg.....gg... /10. 1109/ 20010//VVEEE CCCCCCVP22222// 2022. 00476. 18
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10- 17, 2021, pp. 9992-10002 IEEE, 2021. URL https::/ddoo http://r.g/11.//.......... 10111000019////C4B22 2021 , 00986. 18
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.net, 2017. URL https: / //openrevi ew net/ foormmiddBBy joaddd-yyyyydxxxx.oo 72udxe. 14
Bowen Pan, Rameswar Panda, Yifan Jiang, Zhangyang Wang, Rogerio Feris, and Aude Oliva. Ia- red 2: Interpretability-aaarr redundancy reduction for vision transformers. Advances in Neural Information Processing Systems, 01--24888249111 2021a. 9
Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai. Scalable visual transformers with hierarchical pooling. CoRR, abs/2103.10619 2021b. URL htttps://axxiv.ooo abs/ 2103.10619,
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Detito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- -Bc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019, URL http: http:/aaappaaeeeaa/eeeeeeeeeeee www.ine..rr.ooooo cctpaper/////pccccceeecccc 9015-pytorch-ar ieeee-immeraiiieee style Mateeehigghe pdf. 10
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsich. Dynam- ievit: Efficient vision transformers with dynamic token sparsification. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 13937-13949, 2021. URL https: http://oooeeen...cooo/ 747d3443e319 (tt.eee1111127777100000031103000 2tt.:eaabbbeea--aaaaaattet.
www.diinnnnnnnuusrnoeeeemmuu nwwrips..p.... CC/ paper /022 /ash// Abstract html. 1, 2, 3, 6, 9, 10, 15, 18
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68, 2021. URL
https:// //ops::/doi.. http://oo.oogg/////// wwogoogg11o. 10. 1162/tacll 00353.
Mingxing Tan and Quoc V. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th Inter- national Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 6105-61 (005-6116 6114 PMLR, 2019. URL /att:://proceeeee ngs. ml mtrr.press/99. tan1 9a httll
12

---

Published as a conference paper at ICLR 2023
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 2021. URL http::/pproeeeeeeeee ngs. mlr.
htmm. 2, 9,
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need, In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. newrips cc/ hace/pape//20 2tterr/2011// /hash/
1999-6008 2017. URL https: 3f5ee243547 www.eeeeedoooeeeeeeeeoe1oc Abstraabbbbbtbtta 3
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman, GLUE: A multi-task benchmark and analysis platform for natural language understanding, In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL
net ferumoorrmmcc iid=
xJ4km2R5t 2, 3, 7, 17, 18
Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In IEEE International Symposium Tlihh-Performnce Computer Architecture, HPCA 2021, Seoul, South Korea, February 27 March 3, 2021, pp. 97-110. IEEE, 2021. URL https: / /wo.ooioooo orgg..rr///e 10. 110111199/H11 (A099/PCA515444 2021 00018.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Anmual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: 3 3555e24454776 91 httttps://prooeeediiee neurriiiinnnnnnsnnoeoeee....i neurips cC/ httpppapp//22//aa///// html. 1,4,
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-an natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Pro- cessing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Computational Linguistics. URL https: www. ac Lweb. org/ hnthology///o/oooooooooooooooogt 2020 mmnlp-demos 6. 10
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang. and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL hwwph:ttt//tto /proceed ngs neurips www.c//c..............oeee..eeee cc/paper/ 2.0 0/ wwwash//aaaeeaaaaa 244
Abstt ract. html. 1,3
13

---

Published as a conference paper at ICLR 2023
APPENDIX
A EXPERIMENTAL DETAILS
A.1 DETAILS ON ATTENTION APPROXIMATION NETWORK
Init.
Random
Interpolate
Distilled
Dataset
Factor
2
2
4
MSE
0.0087
0.0087
0.0086
0.0086
0.0013
0.0018
MNLI
KL Div.
1.7503
1.7520
1.7174
1.7097
0.2393
0.3423
QNLI
MSE 0.0052
0.0052
0,0052
0.0052
0.0009
0.0012
KL Div.
1.8026
1.8078
1.7826
1,7750
0.2666
0.3866
MSE 0.0052
0.0052
0,0052
0.0052
0.0007
0.0011
MRPC
Div.
1.9680
1.9751
1,9484
1.9406
0.2248
0.3740
MSE 0.0339
0.0338
0.0337
0.0334
0.0032
0.0042
CoLA
KL Div.
1.0845
1.0810
1.0709
1.0629
0.1223
0.1644
MSE 0.0167
0.0167
0,0166
0,0165
0.0019
0.0025
SST-2
KL Div.
1.4762
1.4744
1,4622
1.4528
0.1899
0.2606
MSE 0.0099
0.0099
0.0099
0.0098
0,0014
0.0020
QQP
KL Div.
1.5960
1,6020
1,5756
1.5679
0.2269
0.3149
MSE 0.0070
0.0070
0,0070
0.0069
0,0005
0.0009
WNLI
KL. Div.
1.5982
1.5977
1.5752
1,5688
0.1375
0.221
MSE
0.0097
0.0097
0.0096
0.0096
0.0009
0.0012
STSB
KI. Div.
1.5230
1.5266
1,4991
1.4915
0.1561
0.2258
MSE
0,0048
0.0048
0,0048
0.0048
0.0007
0.0010
RTE
Div.
1.9883
1.9881
1.9678
1,9593
0.2533
0.3930
MSE
0.0112
0.0112
0.0112
0.0111
0,0013
0.0018
Average GLUE
KL Div.
1.6430
1.6450
1.6221
1.6143
0,2019
0.2981
Table A.1: The ApproxNet performance results over GLUE henchmark with various initialization setting.
In this section, we provide additional details on the attention approximation networks (ApproxNet). We train the ApproxNet following the instructions in Section 3.2. As we mentioned, the architecture of the ApproxNet is identical with the original network such as BERT (Devlin et al., 2019) and Vision Transformer (ViT) (Touvron et al., 2021), except for the dimension of hidden and intermediate hidden unit size. We downsize the architecture by reducing the hidden unit dimension with different reducing factors. We use 4 and 8 as the reducing factors all experiments of the paper. We then train the ApproxNet with the knowledge distillation (Wang et al., 2020; Jiao et al., 2020). Specifically, we use the embedding based distillation loss Lembbd hidden state based distillation loss Lxidns attention- based distillation loss Lattn. and prediction-layer distillation loss Lyred train the ApproxNet for BERT. For Vision Transformer, we do not use the prediction-layer distillation loss.
In Table I, we present the approximation error between the attention probability in the original network and the Net.oxNet. We observe that the distillation helps the ApproxNet to approximate almost identical attention probability of the original network. We also observe that the smaller reduce factor leads to lower attention approximation error. However, we also find that the ApproxNct with a larger reducing factor also accurately approximate the attention probability of the original network.
To boost the training speed for ApproxNet on downstream tasks, we first distill the pre-trained Transformer encoders prior the distillation of the fine-tuned network on each task. This procedure is not necessary if the downstream task is large enough or there is no downstream task such as ImageNet-Ik, kk. For instance, we directly train ApproxNet by setting the fine-tuned VIT the teacher network with ImageNet-Ik Ik dataset. However, for experiments on GLUE benchmark, we first train the ApproxNet by setting the original pre-tr BERT as the teacher network with the Wikitext- 103 dataset (Merity et al., 2017) for 200 epochs. Then, we finetune the ApproxNet by setting the fine-tuned BERT on each task as the teacher network with each task dataset for 30~-50 epochs.
14

---

Published as a conference paper at ICLR 2023
B ADDITIONAL EXPERIMENTAL RESULTS
MNLI
QNLI
MRPC
86
84
92
80
Accurazz/ (N2
Alcurracy (50
60
76
CoLA
SST-2
QQP
55
10
Accurasy (59)
90
(15)
85
86
40
2.5
5.0
7.5
10.0
12.5
15.0
GFLOPs
GFLOPS
GFLDPs
WNLI
STSB
RTE
60
Pearzon Carolation
0.88
Ricuracy 167
Accuriiyy
58
0.87
70
0.86
14
10
GFLOPs
GFLOPs
GFLDPs
STTABT (Approx. Att.)
STTABT@M (Approx Att.
Manual Top-k
BERTa<ae
STTAST (Artual
Figure B.I: FLOPs comparison between attention back-tracking with top-k and manual top-k baseline on all of datasets in GLUE benchmark. STTABT for means the attention back-tracking by ApproxNet with factor z.
B. ADDITIONAL RESULTS ON GLUE EXPERIMENTS WITH FLOPS
In Figure B.I, we present evaluation results on GLUE benchmark with regards FLOPs. We average every data point across 9 datasets in Figure B.1 to plot Figure 7a in the main paper.
B.2 ADDITIONAL. VISUALLIZATION EXAMPLES
In Figure B.2 and B.3, we visualize the additional visualization of the token pruning on each layer with STTABT (Conerete) and baselines (Kim al., 2022; Rao al., 2021). We further confirm that our method prunes more tokens throughout all of transformer layers than baselines in both sentence and image classification tasks.
15

---

Published as a conference paper at ICLR 2023
mma
SmoRT
START
Figure B.2: More visualization examples of the token pruning on each layer with STTABT (Concrete) and Dynamic VIT on ImageNet-I1 1k dataset.
16

---

Published as a conference paper at ICLR 2023
STRAT FONCRE
STAAT Coree
Ure
Upel
(ree
Figure B.3: More visualization examples of the token pruning on each layer with STTABT (Conerete), LTP, and Manual Top-k on SST-2 dataset.
C ABLATION STUDY
C.1 PROPAGATION FACTOR p
CoLA
0.50
MeBiC
0.05
0.40
ABT
Concrete Maskinggt4
Concrete Maskinggrs
Figure C.4: Ablation study about propagation factor p in Algorithm 1,
We perform an ablation study on the importance score propagation factor p. which is introduced in Algorithm 1. We use CoL dataset in GLUE benchmark (Wang et al., 2019) for this ablation study.
17

---

Published as a conference paper at ICLR 2023
In Figure C.4, we observe that Attention Back-Tracking method with value between the range of 0.0-0.5 shows the similar task performance. However, we also observe that using value as 0.1 performs the best.
C.2 CONCRETE Amask AND Xp
CoLA (Metric/Token Retention Ratio)
CoLA (Metric)
1e-05
Anast
0.001
0.1
0.1
15
0,001
0.1
1e:05
0,001
0.1
Ap
Figure C.5: Ablation study on Amask. Ap used Concrete masking training ($3.3). Yellow indicates relatively large value while dark blue indicates relatively small value. (Left) Heatmap of task performance divided by token retention ratio. (Right) Heatmap of task performance.
We perform ablation study two balancing terms Amasks Ap used in Concrete masking training, which is introduced in Section 3.3. We use CoLA dataset in GLUE benchmark (Wang et al., 2019) for this ablation study. In Figure C.5, we could not observe any clear pattern from varying Ap and Amask- Therefore, it is better to adjust both values individually depending on the task. For instance, for tasks in GLUE benchmark, we use low Ap around le 3 in GLUE task training. On the other hand, we observe the model with concrete masking tend fail in keeping the target token retention ration in the image classification task with VIT. In this case, we use a high Amask around 100 to keep the token retention ratio close to the desired target value.
observe 0' does change lot during the training. Therefore,
D EXTRA BASELINES ON VIT EXPERIMENT
D.1 TRADE-OFF BETWEEN TOP-1 ACCURACY AND TOKEN RETENTION RATIO
Technically, it is possible apply proposed method any transformer-based architecture. To validate this, we perform additional experiments on image classification task with another VIT architecture (Jiang et al., 2021). In Figure D.6, additionally plot results of STTABT (Concrete masking) on LVViT (Jiang et al., 2021). Please note that Figure D.6 is the extended version of Figure where we add experimental results STTABT (Concrete) with LVVIT and the original LVVIT (Jiang et al., 2021). Concrete masking on LVViT results in accuracy drop of 3.68%p when drop token retention ratio from 76% to 18%, while Conerete masking DeiT suffers accuracy drop of 5.53%p in the same range of token retention ratio. This result shows our Concrete masking method preserves the performance of the original model by preserving important tokens better.
D.2 TRADE-OFF BETWEEN FLOPS AND TOKEN RETENTION RATIO
We additionally plot the results of STTABT (Concrete masking) with LVVIT small (Jiang al., 2021) and Vyn (Rao et al., 2021) with LVViT amail in Figure D.7. Figure D.7 is an extended version of Figure 7b. When we use LV ViT as the backbone model, Concrete masking outperforms DynamicVil The additional baselines include MViTv2 (Li et al., 2022), SwinT (Liu et al., 2021), and LVVIT (Jiang al., 2021).
amail
18

---

Published as a conference paper at ICLR 2023
ImageNet-1k
82.5
LVVIT mal
80.0
Del
Tap-1 Accuracy (%)
77.5
75.0
STTABT (Concrete, train)
72.5
STTABT oma (Concrete, train)
DynamicVIT
DeiT
70.0
STTABT (Concrete) LWiTsmall
LVVT
67.5
STTABT (Approx. Att.)
STTABT (Actual Att.)
Manual Top-k
20
40
80
Average Token Retention Ratio (%)
Figure D.6: Top-1 accuracy plot with results of STTABT (Concrete) with LVVIT.
ImageNet-1k
84
LVVIT
MViTv2-T
82
tys
Swin-T
80
Top-1 Accuracy (%)
DerT-S
SP R
DeIT-5/320
DelT-S/288
78
HVT-S-1
DeiT S/256
76
74
DynamicViT
STTABT@14 (Concrete) DeiT-S
72
Delt-T
STTABT@R (Concrete) DeiT-S
STTABT (Concrete) LVViTcmal
GFLOPs
Figure D.7: Results of STTABT and Dyn using LVViT small as the backbone network. DyViT-LV-S means Dynamic VIT with LVViT smalt. Note that this figure is the extended version of Figure 7b.
19


Reviews:
[{"review_id": "vayM-WcR7D", "reviewer": "Reviewer_7U3T", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a novel token pruning method based on attention back tracking for efficient transformer inference. The goal (and the difference from the previous token pruning methods) is to mitigate the mistakes of removing important tokens at lower layers. An attention approximation network (ApproxNet) is trained by distillation objectives and the attention probability is used to calculate token masks and significance scores at all transformer layers. The authors also apply Concrete Masking for dynamic thresholding. They have shown the effectiveness of STTABT on NLP and CV datasets, outperforming previous methods having higher sparsity while keeping the accuracy.", "summary_of_the_review": "I enjoyed the paper because the proposed methods are well supported by reasonable motivations (e.g., token pruning decision based on the final prediction and learnable dynamic thresholding). Moreover, STTABT achieves good performance, meaning the practical usefulness of this paper.", "strength_and_weaknesses": "STTABT shows better efficiency-accuracy compared to other token pruning baselines. The paper is well-motivated. STTABT is successful in two modalities, NLP and CV, indicating its generality to any modality that uses a transformer. \n\nAlthough attention back tracking utilizes information from the later layers, it is somewhat heuristic and also induces additional costs (forward pass of ApproxNet and additional computations). The authors should explain these costs.\n\nIn terms of computational efficiency, latency instead of FLOPs might be more important for practitioners.\n\nAn ablation study and intrinsic evaluation of each component (ABT and Concrete Masking) are necessary.\n\nAlthough the authors provided representative token pruning methods (PoWER-BERT, LTP, and DynamicVIT), there are other relevant works that might be more advanced than the included ones in NLP and CV. It would be great if a more comprehensive literature survey and comparison with them were included. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "a4hFJT2p45s", "reviewer": "Reviewer_qjDB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on token pruning for Transformer models. The authors observe that previous token pruning approaches do not consider the impact of a token on later layers’ attentions. Therefore, they propose an attention back-tracking method that tracks the importance of each attention from the outputs to the inputs, based on which they can preserve tokens with a large impact on the final predictions. The method includes an attention approximation network (ApproxNet) that learns to approximate the attention probabilities and a Concrete masking mechanism to learn the thresholds for token pruning. The authors provided results on ImageNet and GLUE to validate the proposed method.\n", "summary_of_the_review": "The idea is novel and makes sense to me. My concern is mainly about the clarity and evaluation of the method (reporting FLOPs numerically, more baselines and types of ViT architectures). I vote for borderline rejection and would like to see the rebuttal to make the final decision.\n", "strength_and_weaknesses": "The idea of attention backtracking makes sense as it is possible that some tokens are seemingly not important at the current layer but may be important at later layers. The proposed method also demonstrates good performance compared to some baseline methods.\n\nI have some concerns about the clarity of the method and empirical evaluation.\n\nHow is ApproxNet trained? Is it trained along with the main Transformer model or separately trained after the main model training is done?\n\nWhat does recursive pruning mean in Figure 1? I understand the main idea of the caption but find the figure hard to interpret.\n\nWhat does DeiT-S/256, DeiT-S/288 mean in Figure 7(b)? Are 256/288 represent different resolutions? If so, why do they have lower accuracy than the original DeiT-S, which runs at resolution 224x224?\n\nI found it’s a bit misleading to mostly use token retention ratio as the main metric for computation due to the additional cost of ApproxNet. It would be great to report FLOPs consistently for all methods and indicate how much is consumed by ApproxNet numerically.\n\nHow does the proposed method perform compared to other types of Vision Transformers, e.g., Swin Transformer or Multiscale Vision Transformers? If the proposed method is generic enough, it would be great to demonstrate results on other types of transformers too.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "V_LhhS6And9", "reviewer": "Reviewer_yJGa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an algorithm to remove a number of tokens during inference passes for Transformer models. The goal is to save time with respect to a full pass (i.e. no token dropping) while sacrificing as little performance as possible. The solution contains two main ingredients: an attention approximation network and the application of concrete masking.\n\nThe core idea is that we should be able to identify the important tokens to keep by looking at the attention scores from the whole network, as they suggest how much some tokens affect others. In particular, the paper claims it is essential to start from the latest layers and recursively propagate the scores backwards. Accordingly, methods that operate on a single forward pass can't identify the right tokens (as they \"must\" make decisions at layer L before seeing the scores of subsequent layers L+1, L+2, ...). The proposal for dealing with this is to train a small auxiliary network that --given a pre-trained \"main\" network-- learns to predict the intermediate attention scores for every layer directly from the inputs via distillation. This way, given a new input at inference, we can apply this network (ApproxNet) first, and use these predictions as proxy to do the token dropping.\n\nThe second idea tackles the problem of thresholding. How many tokens should we drop? A simple approach is to set some fixed value K, and always keep K. I guess this can be easily extended to K_i for layer i (as long as there's some monotonicity). The paper claims it may be hard to set K as it may be suboptimal --and probably K* is input dependent-- so they propose a method that drops all tokens whose score is below some learnable threshold per layer.\n\nThe paper then presents an extensive set of experiments, both for language and vision models.\n\nFigure 3 shows that using STTABT with ApproxNet outperforms forward pass manual top-K, suggesting the former is able to better pick tokens. Figure 3 also shows that approximating the true \"future\" attention values with the ApproxNet doesn't lead to almost any performance loss wrt the real attention values. Concrete masking seems to be helpful when keeping very few tokens. However, this figure does not reflect the overhead of running ApproxNet. In other words, simpler forward-pass approaches can process more tokens at the same cost as STTABT processing fewer. Figure 7 tries to highlight this fact --while it uses FLOPs rather than time.\n\nFigure 4 is cool and informative but maybe you could add the manual top-k choices too; Figure 6 is a bit hard to see (too small).\n\nFigure 5 shows STTABT beats DynamicVIT for pretty much every retention ratio.", "summary_of_the_review": "This paper proposes a way to reduce the number of processed tokens at inference by a pre-trained Transformer model. It trains a small network to predict attention scores, and learns a per-layer threshold to decide which and how many tokens to drop. A number of language and vision experiments suggest the method works well especially when we want to keep very few tokens. Alternative approaches that tend to be very efficient (like merging) aren't mentioned or discussed.\n\nWhile the algorithm is a bit intricate, I think there's value in showing that using attention scores to select tokens can indeed work well.\n", "strength_and_weaknesses": "In terms of strengths, the problem tackled by the paper is both relevant and important these days. Also, the paper provides experiments both for language and vision models, something I really appreciate. It also contains a number of interesting ablations.\n\nIn terms of weaknesses, the proposed algorithm is quite convoluted and it can't be applied directly during training. After training a model, it involves further training of new components (ApproxNet) that will be required during inference. I think this will most likely hinder its use in real-world practical setups.\n\nA few questions:\n\n- What's the size for VIT? is it Base?\n\n- What's the impact of hyper parameter p (defined at the end of page 4)? Can we see performance as a function of p? \n\n- At least for Computer Vision, there are very competitive algorithms [1, 2] that merge tokens (rather than dropping) for Transformers. Accordingly, they are end-to-end differentiable --so no top-K involved anywhere-- and the cost savings are applied and realized also during training (which leads to massive overall savings). It seems both can reduce the number of tokens to just 8 (even from say 256) in the middle of the network, thus matching performance while saving 40-50% of the training and inference costs. It would be nice to discuss the merits of merging (probably way more practical nowadays) and somewhat compare performance wrt these type of algorithms.\n\n[1] = Learning to Merge Tokens in Vision Transformers, Renggli et al.\n[2] = TokenLearner: Adaptive Space-Time Tokenization for Videos, Ryoo et al.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ECMmBIfw0bw", "reviewer": "Reviewer_ibLQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies a novel token pruning/sparsification technique via backtracking the importance from the final layer to the first layer. In order to control the computation cost, the token importance is approximated via a forward procedure through a lightweight distilled counterpart (i.e., ApproxNet), and then a smoothed threshold function (i.e., Concrete masking) is adaptively learned to balance the sparsity level and total importance to preserve. The authors verify the efficiency of the proposed method on both GLUE with BERT base and ImageNet-1k classification with DeiT.", "summary_of_the_review": "This paper considers a sparse token Transformer architecture based on a novel token pruning technique. The proposed technique contains two components: attention backtracing with an approximated model and concrete masking via a learnable threshold function. The authors validate the efficiency of the proposed method in both CV and NLP tasks. In particular, for the low token retention ratio regime, the proposed method beat the benchmarks by a significant margin.\n\nMinor issues:\n\n1. Influence of the ApproxNet. In the proposed method, the token importance is approximated via the ApproxNet. Thus the quality ApproxNet may play a critical role. I suggest authors add a section to directly discuss the influence of the accuracy of ApproxNet.\n\n2. The proposed method still uses Top-k selection, which breaks the linear computation cost. I'm wondering if we can remove it or replace it with some other thresholding policies to maintain the linear cost.\n\n\nI currently tend to accept this paper according to the method's novelty and good numerical performance and I'm willing to change my evaluation after the rebuttal.\n\n\n", "strength_and_weaknesses": "Strength:\n1. The idea of using approximated network during the token pruning seems new in literature.\n2. Promising accuracy and computation trade-off in the low token retention ratio regime.\n\nWeaknesses:\n1. The readability of the paper can be improved. For example, the right part of Figure 1 gives too much detailed information. Without referring to the contents of the later sections, it could be hard to serve as an illustration picture in the introduction section. Figure 2 is also a little bit confusing. \n2. The proposed method requires a reliable lightweight distilled counterpart (ApproxNet) of the model to prune. When ApproxNet is not a good approximation of the original model, the proposed method may fail.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster