Original: S PARSE T OKEN T RANSFORMERS WITH A TTENTION B ACK T RACKING
Revised: S PARSE T OKEN T RANSFORMERS WITH A TTENTION B ACK T RACKING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: While there exist many different approaches to reduce the complexities of the
Revised: While there exist many different approaches to reduce the complexities of the

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Qualitative analysis of the examples further shows that our method does preserve semantically meaningful tokens.
Revised: Qualitative analysis of the examples further shows that our method does preserve semantically meaningful tokens.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformers have achieved huge success in various application domains such as natural language processing (NLP) and computer vision (CV), obtaining state-of-the-art performances on a variety of tasks, and are now considered the de facto standard architectures for a number of domains.
Revised: Transformers have achieved huge success in various application domains such as natural language processing (NLP) and computer vision (CV), obtaining state-of-the-art performances on a variety of tasks, and are now considered the de facto standard architectures for a number of domains.

Original: However,
Revised: However,

Original: Transformers require over a 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
Revised: Transformers require over a 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

Original: RNN architectures (Tan & Le, 2019; Bahdanau et al., 2015).
Revised: RNN architectures (Tan & Le, 2019; Bahdanau et al., 2015).

Original: To reduce such computational burdens of Transformer 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).
Revised: To reduce such computational burdens of Transformer 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).

Original: 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.
Revised: 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.

Original: 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 et al., 2020; Kong et al., 2021).
Revised: 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 et al., 2020; Kong et al., 2021).

Original: How can we decide which tokens to prune then?
Revised: How can we decide which tokens to prune then?

Original: Previous studies either compute the importance score of each input token as the average attention scores (Goyal et al., 2020;
Revised: Previous studies either compute the importance score of each input token as the average attention scores (Goyal et al., 2020;

Original: Kim et al., 2022)
Revised: Kim et al., 2022)

Original: (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).
Revised: (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).

Original: 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 .
Revised: 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 .

Original: This is because all existing works (Kim et al., 2022; Wang 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 1, the [Figure]
Revised: This is because all existing works (Kim et al., 2022; Wang 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

Original: Back

Original: [Figure]
Revised: [Figure]

Original: Tracking Token Pruning
Revised: Tracking Token Pruning

Original: [Figure]
Revised: [Figure]

Original: [ C L S ]
Revised: [ C L S ]

Original: S t e v e L o v e C 
Revised: S t e v e L o v e C a t

Original: [Figure]
Revised: [Figure]

Revised: u e r y

Revised: [Figure]

Revised: (feed-forward token pruning, Figure 1).

Revised: Thus, they may prune out important tokens at earlier layers that are important for the final representation as well as the task loss.

Original:  token Love in the layer l + 1 is pruned by the feed-forward method, although the token has a high attention probability in the representation at the final layer l + 2 that is used by the classification task.
Revised: For example, in Figure 1, the token Love in the layer l + 1 is pruned by the feed-forward method, although the token has a high attention probability in the representation at the final layer l + 2 that is used by the classification task.

Original: To tackle such a limitation of conventional feed-forward token pruning methods, we propose an
Revised: To tackle such a limitation of conventional feed-forward token pruning methods, we propose an

Original: 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.
Revised: 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.

Original: As illustrated in Figure 1 (right), we take a 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.
Revised: As illustrated in Figure 1 (right), we take a 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.

Original: By doing so, we are able to better select and keep the important tokens that will minimally affect the output and the prediction.
Revised: By doing so, we are able to better select and keep the important tokens that will minimally affect the output and the prediction.

Original: We name our novel token pruning method as S parse T oken T ransformer with A ttention B ack- T racking ( STTABT ).
Revised: We name our novel token pruning method as S parse T oken T ransformer with A ttention B ack- T racking ( STTABT ).

Original: However, one challenge here is that such backtracking requires us to know the attention score of each token before feed-forwarding the input sequence to the model.
Revised: However, one challenge here is that such backtracking requires us to know the attention score of each token before feed-forwarding the input sequence to the model.

Original: To handle the issue, we introduce a lightweight attention approximation network trained with knowledge distillation called ApproxNet .
Revised: To handle the issue, we introduce a lightweight attention approximation network trained with knowledge distillation called ApproxNet .

Original: Moreover, to actually prune the tokens, we need to decide which tokens to be retained based on the importance scores.
Revised: Moreover, to actually prune the tokens, we need to decide which tokens to be retained based on the importance scores.

Original: The top-k method, by design, can select a predefined number of tokens (Goyal et al., 2020), but requires sorting the tokens by their importance scores, which is not differentiable.
Revised: The top-k method, by design, can select a predefined number of tokens (Goyal et al., 2020), but requires sorting the tokens by their importance scores, which is not differentiable.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: Our method is generally applicable to Transformer encoders for any domains, and we validate our method on text and image classification tasks.
Revised: Our method is generally applicable to Transformer encoders for any domains, and we validate our method on text and image classification tasks.

Original: For text classification, we validated the proposed token pruning methods on GLUE (Wang et al., 2019) benchmark with BERT (Devlin et al., 2019).
Revised: For text classification, we validated the proposed token pruning methods on GLUE (Wang et al., 2019) benchmark with BERT (Devlin et al., 2019).

Original: For image classification, we validate it on ImageNet-1K (Deng et al., 2009) benchmark with DeiT (Touvron et al., 2021).
Revised: For image classification, we validate it on ImageNet-1K (Deng et al., 2009) benchmark with DeiT (Touvron et al., 2021).

Original: The experimental results show that our model works surprisingly well even at very high pruning rates, for which baselines show significant performance degeneration.
Revised: The experimental results show that our model works surprisingly well even at very high pruning rates, for which baselines show significant performance degeneration.

Original: For example, on a GLUE benchmark task (QQP), our method obtains 45.54% increased token sparsity ( 17 . 15% to 9 .
Revised: For example, on a GLUE benchmark task (QQP), our method obtains 45.54% increased token sparsity ( 17 . 15% to 9 .

Original: compared to the baseline with the same accuracy.
Revised: compared to the baseline with the same accuracy.

Original: Moreover, our method obtains very high token sparsity ( 18 . 7% ), keeping only half of the tokens retained by DynamicViT ( 37 . 5% ), with only 0 .
Revised: Moreover, our method obtains very high token sparsity ( 18 . 7% ), keeping only half of the tokens retained by DynamicViT ( 37 . 5% ), with only 0 .

Original: 8% accuracy loss on ImageNet-1K benchmark.
Revised: 8% accuracy loss on ImageNet-1K benchmark.

Original: Our contributions can be summarized as follows:
Revised: Our contributions can be summarized as follows:

Original: • We propose a novel token pruning method for Transformers based on 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.
Revised: • We propose a novel token pruning method for Transformers based on 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.

Original: • We propose Concrete masking, which automatically learns the pruning rate for each layer with Concrete dropout.
Revised: • We propose Concrete masking, which automatically learns the pruning rate for each layer with Concrete dropout.

Original: (Gal et al., 2017).
Revised: (Gal et al., 2017).

Original: • Our token pruning is generally applicable to any Transformers, and we validated it on trained Transformers from both NLP (BERT) and CV (DeiT) domains, on GLUE and ImageNet benchmark, whose results show that it achieves a significantly better accuracy-computation tradeoff.
Revised: • Our token pruning is generally applicable to any Transformers, and we validated it on trained Transformers from both NLP (BERT) and CV (DeiT) domains, on GLUE and ImageNet benchmark, whose results show that it achieves a significantly better accuracy-computation tradeoff.

Original: (Kim et al., 2022; Rao et al., 2021).
Revised: (Kim et al., 2022; Rao et al., 2021).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Attention Pruning for Transformers Since the literature on the general neural network compression techniques such as quantization and weight pruning is vast, we discuss compression techniques specifically targeting Transformers (Vaswani et al., 2017).
Revised: Attention Pruning for Transformers Since the literature on the general neural network compression techniques such as quantization and weight pruning is vast, we discuss compression techniques specifically targeting Transformers (Vaswani et al., 2017).

Original: Due to the heavy computational cost of the self-attention, several recent works have focused to reduce its time complexity from O ( N 2 ) to about O ( N ) given the N length of the sequence by pruning the attentions in the self-attention layers.
Revised: Due to the heavy computational cost of the self-attention, several recent works have focused to reduce its time complexity from O ( N 2 ) to about O ( N ) given the N length of the sequence by pruning the attentions in the self-attention layers.

Original: 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.
Revised: 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.

Original: Similarly,
Revised: Similarly,

Original: 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.
Revised: 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.

Original: However, these prior works do not consider the content of the input sequence, as they make use of static sparsity patterns.
Revised: However, these prior works do not consider the content of the input sequence, as they make use of static sparsity patterns.

Original: To remedy this, some recent works aim to dynamically sparsify attentions of the Transformer considering the input data.
Revised: To remedy this, some recent works aim to dynamically sparsify attentions of the Transformer considering the input data.

Original: Specifically, Reformer (Kitaev et al., 2020) divides the tokens into several groups through Locally Sensitive Hashing (LSH) and calculates attentions by sparsely connecting them.
Revised: Specifically, Reformer (Kitaev et al., 2020) divides the tokens into several groups through Locally Sensitive Hashing (LSH) and calculates attentions by sparsely connecting them.

Original: Routing Transformer (Roy et al., 2021) uses the k-means clustering algorithm to cluster input tokens.
Revised: Routing Transformer (Roy et al., 2021) uses the k-means clustering algorithm to cluster input tokens.

Original: Token Pruning for Transformers
Revised: Token Pruning for Transformers

Original: 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 et al., 2022).
Revised: 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 et al., 2022).

Original: Moreover, fully connected layers require significant computations that are non-negligible, although they have been relatively overlooked.
Revised: Moreover, fully connected layers require significant computations that are non-negligible, although they have been relatively overlooked.

Original: 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.
Revised: 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.

Original: Notably,
Revised: Notably,

Original: PoWER-BERT (Goyal et al., 2020) progressively eliminates tokens based on their importance scores, computed from the attentions at each layer of the Transformer encoder, and Learned Token Pruning (LTP) (Kim et al., 2022) aims to learn to threshold the importance scores to prune the token.
Revised: PoWER-BERT (Goyal et al., 2020) progressively eliminates tokens based on their importance scores, computed from the attentions at each layer of the Transformer encoder, and Learned Token Pruning (LTP) (Kim et al., 2022) aims to learn to threshold the importance scores to prune the token.

Original: However, both methods prune the tokens during the feed-forward pass, assuming that important input tokens are also important output tokens.
Revised: However, both methods prune the tokens during the feed-forward pass, assuming that important input tokens are also important output tokens.

Original: Therefore, they may prune out tokens that are important for the output vector at earlier layers.
Revised: Therefore, they may prune out tokens that are important for the output vector at earlier layers.

Original: Compared to them, our method STTABT can consider the impact of each token on the output vector and the prediction task by attention back-tracking, thanks to our novel lightweight attention approximation network.
Revised: Compared to them, our method STTABT can consider the impact of each token on the output vector and the prediction task by attention back-tracking, thanks to our novel lightweight attention approximation network.

Original:  S PARSE T OKEN T RANSFORMER WITH A TTENTION B ACK T RACKING
Revised:  S PARSE T OKEN T RANSFORMER WITH A TTENTION B ACK T RACKING

Original: We start by briefly describing the Transformer model that our method focuses on.
Revised: We start by briefly describing the Transformer model that our method focuses on.

Original: Then, we propose two components for our Sparse Token Transformer with Attention Back Tracking (STTABT):
Revised: Then, we propose two components for our Sparse Token Transformer with Attention Back Tracking (STTABT):

Original: Attention 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.
Revised: Attention 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.

Original: STTABT is a post-hoc token pruning method and assumes that the target Transformer encoder has previously been optimized on the training dataset for the target task.
Revised: STTABT is a post-hoc token pruning method and assumes that the target Transformer encoder has previously been optimized on the training dataset for the target task.

Original: Our goal is to lower the inference cost of the given pretrained Transformer by pruning the tokens at each layer.
Revised: Our goal is to lower the inference cost of the given pretrained Transformer by pruning the tokens at each layer.

Original: 3. T RANSFORMER
Revised: 3. T RANSFORMER

Original: 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.
Revised: 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.

Original: Let us denote the attention score and probability at l -th Transformer layer as follows:
Revised: Let us denote the attention score and probability at l -th Transformer layer as follows:

Original: [Equation]
Revised: [Equation]

Original: where Q l ∈ R L ×
Revised: where Q l ∈ R L ×

Original: d k is the query matrix, K l ∈ R L ×
Revised: d k is the query matrix, K l ∈ R L ×

Original: d k is the key matrix, d k is the dimension of query and key matrix, and l ∈ [1 , L ] is the layer index.
Revised: d k is the key matrix, d k is the dimension of query and key matrix, and l ∈ [1 , L ] is the layer index.

Original: 3. A TTENTION B ACK -T RACKING ( ABT )
Revised: 3. A TTENTION B ACK -T RACKING ( ABT )

Original: We now describe the details of our method,
Revised: We now describe the details of our method,

Original: Attention Back Tracking ( ABT ).
Revised: Attention Back Tracking ( ABT ).

Original: The fundamental idea of
Revised: The fundamental idea of

Original: 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.
Revised: 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.

Original: Specifically, as shown in Figure 1 (right), we start with the tokens with higher attention probabilities from the last layer, and backtrack their accumulated attention probabilities to the first layer of the Transformer encoder consisting of L layers with regards to the token for the sequence classification task (specifically, [CLS] token).
Revised: Specifically, as shown in Figure 1 (right), we start with the tokens with higher attention probabilities from the last layer, and backtrack their accumulated attention probabilities to the first layer of the Transformer encoder consisting of L layers with regards to the token for the sequence classification task (specifically, [CLS] token).

Original: However, to apply such a technique into practice, we must feed-forward the input sequence through all of the Transformer layers in order to compute attention scores beforehand.
Revised: However, to apply such a technique into practice, we must feed-forward the input sequence through all of the Transformer layers in order to compute attention scores beforehand.

Original: Thus a straightforward implementation of the technique could be quite inefficient because it will require us to perform forward pass of the model twice for each input.
Revised: Thus a straightforward implementation of the technique could be quite inefficient because it will require us to perform forward pass of the model twice for each input.

Original: Attention Approximation Network To alleviate the cost, we use an attention approximation network (ApproxNet), a distilled Transformer encoder that approximates the attention probabilities of the original network with a significantly less number of parameters.
Revised: Attention Approximation Network To alleviate the cost, we use an attention approximation network (ApproxNet), a distilled Transformer encoder that approximates the attention probabilities of the original network with a significantly less number of parameters.

Original: We perform knowledge distillation based on the distillation loss from TinyBERT (Jiao et al., 2020) to build the ApproxNet.
Revised: We perform knowledge distillation based on the distillation loss from TinyBERT (Jiao et al., 2020) to build the ApproxNet.

Original: Specifically, we leverage embedding based distillation loss L embd , hidden state based distillation loss L hidn , attention-based distillation loss L attn , and prediction-layer distillation loss L pred for the distillation.
Revised: Specifically, we leverage embedding based distillation loss L embd , hidden state based distillation loss L hidn , attention-based distillation loss L attn , and prediction-layer distillation loss L pred for the distillation.

Original: Especially, for the attention-based distillation, we use the KL divergence loss from
Revised: Especially, for the attention-based distillation, we use the KL divergence loss from

Original: MiniLM (Wang et al., 2020) instead of the mean squared error loss as follows:
Revised: MiniLM (Wang et al., 2020) instead of the mean squared error loss as follows:

Original: [Equation]
Revised: [Equation]

Original:  where A lprob ∈ R H × T × T denotes l -th layer attention probability of the original network (teacher), 
Revised: j = T (cid:88) i = ( [Equation] (2) where A lprob ∈ R H × T × T denotes l -th layer attention probability of the original network (teacher), ˆ

Original: ˆ A lprob ∈
Revised: A lprob ∈

Original: R H × T × T denotes l -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. 
Revised: R H × T × T denotes l -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. [Figure]

Original: 1 def UpdateTokenMaskOfLayer( ˆ A lprob , i l +1 , s l +1 , p , k ) //
Revised: def UpdateTokenMaskOfLayer( ˆ A lprob , i l +1 , s l +1 , p , k ) //

Original: Reduce attention heads and gather required output tokens’ attention probabilities with previous importance score.
Revised: Reduce attention heads and gather required output tokens’ attention probabilities with previous importance score.

Original:  A ← gather(mean( ˆ A lprob , dim=0), i l +1 , dim=0)
Revised: [Figure] A ← gather(mean( ˆ A lprob , dim=0), i l +1 , dim=0)

Revised: [Figure]

Original: a ← mean( diag ( s l +1 )
Revised: a ← mean( diag ( s l +1 )

Original: A ∗ p + A ∗ (1 − p ) , dim=0)
Revised: A ∗ p + A ∗ (1 − p ) , dim=0)

Original: i l ← unique(cat( i l +1 , topk( a , k )))

Original: s l ← gather( a , i l ) s l ← s l / sum( s l )

Original: [Figure]
Revised: [Figure]

Original: s
Revised: ← unique(cat( i l +1 , topk( a , k )))

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: Token Pruning with Attention Back-tracking on Approximated Attention Then, using the approximated attentions from the ApproxNet, we can now prune the tokens at each layer of the main network.
Revised: Token Pruning with Attention Back-tracking on Approximated Attention Then, using the approximated attentions from the ApproxNet, we can now prune the tokens at each layer of the main network.

Original: In Algorithm 1, we define the pruning function UpdateTokenMaskOfLayer , which updates the mask at each Transformer layer l for token pruning.
Revised: In Algorithm 1, we define the pruning function UpdateTokenMaskOfLayer , which updates the mask at each Transformer layer l for token pruning.

Original: The input arguments of the pruning function are the approximated attention probability ˆ A lprob of l -th layer from the ApproxNet, the token mask i
Revised: The input arguments of the pruning function are the approximated attention probability ˆ A lprob of l -th layer from the ApproxNet, the token mask i

Original: l +1 ∈
Revised: l +1 ∈

Original: Z o l +1 , and the importance scores of the tokens s l +1 ∈ R
Revised: Z o l +1 , and the importance scores of the tokens s l +1 ∈ R

Original: o l +1 from the upper layer l + 1 , where l = 1 is the first layer of the Transformer model and o l +1 is the number of selected tokens in ( l + 1) -th layer.
Revised: o l +1 from the upper layer l + 1 , where l = 1 is the first layer of the Transformer model and o l +1 is the number of selected tokens in ( l + 1) -th layer.

Original: i L is initialized with indices of required tokens, which is [CLS] token in the classification task.
Revised: i L is initialized with indices of required tokens, which is [CLS] token in the classification task.

Original: s
Revised: s

Original: L ∈ R
Revised: L ∈ R

Original: o T is initialized as the one vector 1 /o T where o t is the number of output tokens, which is generally a single token.
Revised: o T is initialized as the one vector 1 /o T where o t is the number of output tokens, which is generally a single token.

Original: p ∈ (0 , 1) ⊂
Revised: p ∈ (0 , 1) ⊂

Original: R determines how much importance scores should be propagated from the upper layer.
Revised: R determines how much importance scores should be propagated from the upper layer.

Original: If p = 0 , the importance scores from the upper layer will be completely ignored when pruning the tokens in l -th layer.
Revised: If p = 0 , the importance scores from the upper layer will be completely ignored when pruning the tokens in l -th layer.

Original: k ∈ 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 to retain at each layer, by adjusting k .
Revised: k ∈ 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 to retain at each layer, by adjusting k .

Original: The function returns the index of remaining input tokens at l -th layer and their importance scores given those arguments.
Revised: The function returns the index of remaining input tokens at l -th layer and their importance scores given those arguments.

Original: 3. C ONCRETE M ASKING
Revised: 3. C ONCRETE M ASKING

Original: 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, a mechanism for dynamically altering the pruning ratio is also necessary.
Revised: 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, a mechanism for dynamically altering the pruning ratio is also necessary.

Original: Therefore, previous work, such as LTP (Kim et al., 2022) proposed to learn a layer-wise threshold for token pruning.
Revised: Therefore, previous work, such as LTP (Kim et al., 2022) proposed to learn a layer-wise threshold for token pruning.

Original: In detail, during fine-tuning,
Revised: In detail, during fine-tuning,

Original: LTP keeps tokens that have larger scores than the threshold, which is jointly optimized with the network parameters.
Revised: LTP keeps tokens that have larger scores than the threshold, which is jointly optimized with the network parameters.

Original: 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.
Revised: 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.

Original: To address the problem, we suggest a novel dynamic thresholding method named as Concrete Masking , inspired by the learnable dropout function idea in Concrete dropout (Gal et al., 2017).
Revised: To address the problem, we suggest a novel dynamic thresholding method named as Concrete Masking , inspired by the learnable dropout function idea in Concrete dropout (Gal et al., 2017).

Original: Let us consider a single row of one of multi-head attention score matrix
Revised: Let us consider a single row of one of multi-head attention score matrix

Original: a lt = A scorelh,t ∈ R T in this section for brevity, where h is the head index
Revised: a lt = A scorelh,t ∈ R T in this section for brevity, where h is the head index

Original: and i is the row index.
Revised: and i is the row index.

Original: We further assume that the attention scores follow a normal distribution.
Revised: We further assume that the attention scores follow a normal distribution.

Original: We transform a t into the uniform distribution where a t,i ∈ [0 , 1] so that the importance score becomes uniformly distributed.
Revised: We transform a t into the uniform distribution where a t,i ∈ [0 , 1] so that the importance score becomes uniformly distributed.

Original: Formally,
Revised: Formally,

Original: [Equation]
Revised: [Equation]

Original: where Φ is the cumulative distribution function of the standard normal distribution, µ a = (cid:80)
Revised: where Φ is the cumulative distribution function of the standard normal distribution, µ a = (cid:80)

Original: Tt =
Revised: Tt =

Original: a lt /T is the mean of a l , and
Revised: a lt /T is the mean of a l , and

Original: σ a = (cid:113)(cid:80) Tt =1 ( a lt − µ a ) 2 /T is the standard deviation of the a l .
Revised: σ a = (cid:113)(cid:80) Tt =1 ( a lt − µ a ) 2 /T is the standard deviation of the a l .

Original: Then, we apply above transformation to all the rows of the predicted attention score matrix from ApproxNet ˆ
Revised: Then, we apply above transformation to all the rows of the predicted attention score matrix from ApproxNet ˆ

Original: A lscore to acquire ˜
Revised: A lscore to acquire ˜

Original: A l where all of rows are uniformly distributed.
Revised: A l where all of rows are uniformly distributed.

Original: Similar to the line 3 of Algorithm 1, we then propagate the transformed importance score ˜ a as follows:
Revised: Similar to the line 3 of Algorithm 1, we then propagate the transformed importance score ˜ a as follows:

Original: [Equation]
Revised: [Equation]

Original: where the token mask m l +1 ∈ (0 , 1) T ⊂ R T is the Concrete mask vector.
Revised: where the token mask m l +1 ∈ (0 , 1) T ⊂ R T is the Concrete mask vector.

Original: Then, the updated
Revised: Then, the updated

Original: Concrete token mask m t is defined as follows (Gal et al., 2017):
Revised: Concrete token mask m t is defined as follows (Gal et al., 2017):

Original: [Equation]
Revised: [Equation]

Original: where sigmoid is the sigmoid function, τ is the temperature hyperparameter, and p θ
Revised: where sigmoid is the sigmoid function, τ is the temperature hyperparameter, and p θ

Original: l = sigmoid ( θ l ) where θ l is a learnable parameter and initialized with θ init .
Revised: l = sigmoid ( θ l ) where θ l is a learnable parameter and initialized with θ init .

Original: During training, the model learn the proper value of θ l with the soft mask function applied to each layer.
Revised: During training, the model learn the proper value of θ l with the soft mask function applied to each layer.

Original: After training for reproducing the base accuracy with sufficient epochs, we change the soft mask into a hard mask, then fine-tune the parameters to fit the hard mask as in LTP (Kim et al., 2022).
Revised: After training for reproducing the base accuracy with sufficient epochs, we change the soft mask into a hard mask, then fine-tune the parameters to fit the hard mask as in LTP (Kim et al., 2022).

Original: For hard masking, we threshold the soft Concrete token mask m l using a set threshold, usually 0 .
Revised: For hard masking, we threshold the soft Concrete token mask m l using a set threshold, usually 0 .

Original: To train our Concrete masking mechanism, we need regularization terms in addition to the task loss.
Revised: To train our Concrete masking mechanism, we need regularization terms in addition to the task loss.

Original: First of all, we add the regularization term L p , which prevents θ l from deviating too far from the initialized value and another regularization term L mask so that the average token retention ratio stay around p θ init as follows:
Revised: First of all, we add the regularization term L p , which prevents θ l from deviating too far from the initialized value and another regularization term L mask so that the average token retention ratio stay around p θ init as follows:

Original: [Equation]
Revised: [Equation]

Original: where λ p and λ mask are hyperparameters.
Revised: where λ p and λ mask are hyperparameters.

Original:  E XPERIMENTS 4. E XPERIMENTAL S ETUP
Revised:  E XPERIMENTS 4. E XPERIMENTAL S ETUP

Original: To show the general applicability of our method, STTABT, we validate it on both Natural Language
Revised: To show the general applicability of our method, STTABT, we validate it on both Natural Language

Original: Processing (NLP) and vision tasks.
Revised: Processing (NLP) and vision tasks.

Original: 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 a metric that measures how many tokens are retrained after pruning.
Revised: 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 a metric that measures how many tokens are retrained after pruning.

Original: We use the average token retention ratio values across all the layers of the target Transformers.
Revised: We use the average token retention ratio values across all the layers of the target Transformers.

Original: 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 FLOPs in Section 5.
Revised: 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 FLOPs in Section 5.

Original: 4.1. P RUNING T RANSFORMER - BASED T RAINED L ANGUAGE M ODELS
Revised: 4.1. P RUNING T RANSFORMER - BASED T RAINED L ANGUAGE M ODELS

Original: Dataset and metric.
Revised: Dataset and metric.

Original: We use nine datasets from GLUE (Wang et al., 2019) benchmark for the text classification and use the BERT base (Devlin et al., 2019) as the base model.
Revised: We use nine datasets from GLUE (Wang et al., 2019) benchmark for the text classification and use the BERT base (Devlin et al., 2019) as the base model.

Original: For evaluation metric, we use the metrics from the original paper (Wang et al., 2019).
Revised: For evaluation metric, we use the metrics from the original paper (Wang et al., 2019).

Original: Baselines.
Revised: Baselines.

Original: We compare
Revised: We compare

Original: STTABT against two notable token pruning baselines – the manual top-k method and Learned Token
Revised: STTABT against two notable token pruning baselines – the manual top-k method and Learned Token

Original: Pruning (LTP) (Kim et al., 2022).
Revised: Pruning (LTP) (Kim et al., 2022).

Original: The manual top-k method is the feed-forward token pruning which keeps the k l tokens with highest importance scores for each l -th layer.
Revised: The manual top-k method is the feed-forward token pruning which keeps the k l tokens with highest importance scores for each l -th layer.

Original: LTP is a baseline with a learnable threshold, which requires to fine-tune the whole parameters to learn the appropriate threshold value for each input.
Revised: LTP is a baseline with a learnable threshold, which requires to fine-tune the whole parameters to learn the appropriate threshold value for each input.

Original: We match the training settings of STTABT and LTP, then follow the hyperparameter search strategies described in (Kim et al., 2022).
Revised: We match the training settings of STTABT and LTP, then follow the hyperparameter search strategies described in (Kim et al., 2022).

Original: 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.
Revised: 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.

Original: 4.1. P RUNING V ISION T RANSFORMERS
Revised: 4.1. P RUNING V ISION T RANSFORMERS

Original: Dataset and metric.
Revised: Dataset and metric.

Original: 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.
Revised: 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.

Original: Specifically, we use DeiT-small as the base architecture for all experiments on the image classification.
Revised: Specifically, we use DeiT-small as the base architecture for all experiments on the image classification.

Original: Baselines.
Revised: Baselines.

Original: We use DynamicViT
Revised: We use DynamicViT

Original: (Rao et al., 2021) as the primary baseline.
Revised: (Rao et al., 2021) as the primary baseline.

Original: We follow the experimental setups, including data preprocessing and training pipeline of DynamicViT for all experiments on the image classification task.
Revised: We follow the experimental setups, including data preprocessing and training pipeline of DynamicViT for all experiments on the image classification task.

Original: 4. E XPERIMENTAL R ESULTS ON NLP T ASKS
Revised: 4. E XPERIMENTAL R ESULTS ON NLP T ASKS

Original: Attention back-tracking.
Revised: Attention back-tracking.

Original: We measure the performance of our attention back-tracking (ABT) method with the ApproxNet (§3.2).
Revised: We measure the performance of our attention back-tracking (ABT) method with the ApproxNet (§3.2).

Original: We experiment with different k in Algorithm 1 to vary the average token retention ratios.
Revised: We experiment with different k in Algorithm 1 to vary the average token retention ratios.

Original: The experimental results in Figure 3 show that the our method significantly outperforms the manual top-k with the same average token retention ratio.
Revised: The experimental results in Figure 3 show that the our method significantly outperforms the manual top-k with the same average token retention ratio.

Original: Furthermore, we find that using the ApproxNet (blue lines) obtains similar performance to using the attention from the full network (green lines).
Revised: Furthermore, we find that using the ApproxNet (blue lines) obtains similar performance to using the attention from the full network (green lines).

Original: 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.
Revised: 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.

Original: Concrete masking.
Revised: Concrete masking.

Original: In Figure 3, we further present the performance evaluations with or without the Concrete masking, to examine its effect.
Revised: In Figure 3, we further present the performance evaluations with or without the Concrete masking, to examine its effect.

Original: One of the advantages of our model with Concrete masking is that the user can flexibly adjust the token retention ratio by changing θ init value.
Revised: One of the advantages of our model with Concrete masking is that the user can flexibly adjust the token retention ratio by changing θ init value.

Original: 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.
Revised: 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.

Original: For the lower token retention ratio of LTP, we must increase the scaling factor, then the regularization term becomes significantly bigger than task loss at some point.
Revised: For the lower token retention ratio of LTP, we must increase the scaling factor, then the regularization term becomes significantly bigger than task loss at some point.

Original: 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.
Revised: 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.

Original: For instance, in QNLI, MRPC, CoLA, and SST-2 of Figure 3,
Revised: For instance, in QNLI, MRPC, CoLA, and SST-2 of Figure 3,

Original: LTP baseline cannot cover the low token retention ratio, whereas our method can.
Revised: LTP baseline cannot cover the low token retention ratio, whereas our method can.

Original: At relatively high token retention ratios around 30-50%, both STTABT with concrete masking and
Revised: At relatively high token retention ratios around 30-50%, both STTABT with concrete masking and

Original: LTP work well, although our method still outperforms LTP in most cases.
Revised: LTP work well, although our method still outperforms LTP in most cases.

Original: However, at lower than 30% token retention ratios, our method with Concrete masking show significantly superior performance
Revised: However, at lower than 30% token retention ratios, our method with Concrete masking show significantly superior performance

Original: [Figure]
Revised: [Figure]

Original: over LTP, better preserving the accuracy of the full model (Please see the results on MNLI, QNLI,
Revised: over LTP, better preserving the accuracy of the full model (Please see the results on MNLI, QNLI,

Original: MRPC, CoLA, SST-2, QQP, and RTE).
Revised: MRPC, CoLA, SST-2, QQP, and RTE).

Original: To examine which tokens are pruned with our method and
Revised: To examine which tokens are pruned with our method and

Original: LTP, we further visualize the pruning mask at each layer of BERT base in Figure 6.
Revised: LTP, we further visualize the pruning mask at each layer of BERT base in Figure 6.

Original: We observe that
Revised: We observe that

Original: LTP keeps most of the important tokens, but unimportant tokens as well such as dots in Figure 6.
Revised: LTP keeps most of the important tokens, but unimportant tokens as well such as dots in Figure 6.

Original: On the other hand, our method with Concrete masking prunes useless tokens such as dots at much earlier layers than LTP, yielding higher sparsity in total.
Revised: On the other hand, our method with Concrete masking prunes useless tokens such as dots at much earlier layers than LTP, yielding higher sparsity in total.

Original: The Concrete masking method have significantly higher accuracy in most of cases in GLUE benchmark than ABT without Concrete masking and baselines.
Revised: The Concrete masking method have significantly higher accuracy in most of cases in GLUE benchmark than ABT without Concrete masking and baselines.

Original: The Concrete masking method works especially well at high pruning rates, while have competitive performance to baselines at low pruning rates too.
Revised: The Concrete masking method works especially well at high pruning rates, while have competitive performance to baselines at low pruning rates too.

Original: [Figure]
Revised: [Figure]

Original: (SST-2) (Wang et al., 2019)of trained and
Revised: SST-2) (Wang et al., 2019)

Original: Concrete masking modelsare used
Revised: Concrete masking models

Original: ≈ 89 %).
Revised: ≈ 89 %).

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: 4. E XPERIMENTAL R ESULTS ON C OMPUTER V ISION T ASKS
Revised: 4. E XPERIMENTAL R ESULTS ON C OMPUTER V ISION T ASKS

Original: The experimental results show that our Attention back-tracking (ABT) performs competitively to
Revised: The experimental results show that our Attention back-tracking (ABT) performs competitively to

Original: DynamicViT even without Concrete masking, on computer vision tasks, and our full model with
Revised: DynamicViT even without Concrete masking, on computer vision tasks, and our full model with

Original: Concrete masking significantly outperforms Dynamic ViT 5, especially at lower token retention ratios.
Revised: Concrete masking significantly outperforms Dynamic ViT 5, especially at lower token retention ratios.

Original: For example, our method with Concrete masking achieves 74 .
Revised: For example, our method with Concrete masking achieves 74 .

Original: 1% ( − 5 .
Revised: 1% ( − 5 .

Original: 8% p ) accuracy with 18 .
Revised: 8% p ) accuracy with 18 .

Original: 7% token retention ratio, but DynamicViT requires 37 .
Revised: 7% token retention ratio, but DynamicViT requires 37 .

Original: 5% token retention ratio with 74 .
Revised: 5% token retention ratio with 74 .

Original: 9% ( − 5 .
Revised: 9% ( − 5 .

Original: 0% p ) accuracy.
Revised: 0% p ) accuracy.

Original: Moreover, since our ABT requires training ApproxNet only once regardless of the target token retention ratio, it has a huge advantage in training cost than DynamicViT.
Revised: Moreover, since our ABT requires training ApproxNet only once regardless of the target token retention ratio, it has a huge advantage in training cost than DynamicViT.

Original: The DynamicViT used twice more tokens than our Concrete masking, but the accuracy improvement is only 0 .
Revised: The DynamicViT used twice more tokens than our Concrete masking, but the accuracy improvement is only 0 .

Original: Even considering attention approximation overhead, our model performs better than DynamicViT since the overhead is only about 1 .
Revised: Even considering attention approximation overhead, our model performs better than DynamicViT since the overhead is only about 1 .

Original: 95 − 6 .
Revised: 95 − 6 .

Original: 86% of the base model, as shown in Table 2.
Revised: 86% of the base model, as shown in Table 2.

Original: We visualize the pruned tokens at different layers of the vision Transformers in Figure 6.
Revised: We visualize the pruned tokens at different layers of the vision Transformers in Figure 6.

Original: The visualization shows that our model with Concrete masking concentrates on the important tokens in an earlier layer than DynamicViT.
Revised: The visualization shows that our model with Concrete masking concentrates on the important tokens in an earlier layer than DynamicViT.

Original: Moreover, our model preserves more important tokens while DynamicViT attends to a larger number of tokens, including unimportant ones.
Revised: Moreover, our model preserves more important tokens while DynamicViT attends to a larger number of tokens, including unimportant ones.

Original: For instance, DynamicViT 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.
Revised: For instance, DynamicViT 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.

Original:  C OMPUTATIONAL E FFICIENCY OF STTABT 5. FLOP S C OMPARISION
Revised:  C OMPUTATIONAL E FFICIENCY OF STTABT 5. FLOP S C OMPARISION

Original: [Table]
Revised: [Table]

Original: For the GLUE benchmark, we compare the attention backtracking (ABT) method and the manual top-k method with respect to FLOPs.
Revised: For the GLUE benchmark, we compare the attention backtracking (ABT) method and the manual top-k method with respect to FLOPs.

Original: In Figure 7a, we compare the performance with respect to FLOPs over the test set.
Revised: In Figure 7a, we compare the performance with respect to FLOPs over the test set.

Original: We find that the ABT performs better than the manual top-k baseline on most datasets.
Revised: We find that the ABT performs better than the manual top-k baseline on most datasets.

Original: 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.
Revised: 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.

Original: We examine the relative overheads of each additional component in Table 2.
Revised: We examine the relative overheads of each additional component in Table 2.

Original: We see that ApproxNet
Revised: We see that ApproxNet

Revised: [Figure]

Revised: Average GLUE

Original: [Figure]
Revised: [Figure]

Original: is the major bottleneck, as the ApproxNet with factor 4 requires about 6 .
Revised: is the major bottleneck, as the ApproxNet with factor 4 requires about 6 .

Original: 86% of the BERT base
Revised: 86% of the BERT base

Original: FLOPs, which is significantly larger than the manual top-k ( 0 . 01% ).
Revised: FLOPs, which is significantly larger than the manual top-k ( 0 . 01% ).

Original: However, the overhead is more convincing with higher reduction factor of 8 ( 1 . 95% ) than 4.
Revised: However, the overhead is more convincing with higher reduction factor of 8 ( 1 . 95% ) than 4.

Original: Therefore in Figure 7a, the ABT results with factor 8 shows better result than 4.
Revised: Therefore in Figure 7a, the ABT results with factor 8 shows better result than 4.

Original: We further compare the performance of STTABT with different reduction factors of the ApproxNet.
Revised: We further compare the performance of STTABT with different reduction factors of the ApproxNet.

Original: Table 1 shows that the ApproxNet performs accurately approximations of the actual attention scores with a lower reduction factor, thus results in better preformance.
Revised: Table 1 shows that the ApproxNet performs accurately approximations of the actual attention scores with a lower reduction factor, thus results in better preformance.

Original: However, when considering the accuracy-efficiency tradeoff, the performance of the model with the reduction factor of 4 are often worse than the one with the reduction factor of 8.
Revised: However, when considering the accuracy-efficiency tradeoff, the performance of the model with the reduction factor of 4 are often worse than the one with the reduction factor of 8.

Original: In Figure 7b, we show the accuracy over FLOPs comparison of our model against baselines on
Revised: In Figure 7b, we show the accuracy over FLOPs comparison of our model against baselines on

Original: ImageNet-1k (Deng et al., 2009), using DeiT small as the base model to prune.
Revised: ImageNet-1k (Deng et al., 2009), using DeiT small as the base model to prune.

Original: The baselines include
Revised: The baselines include

Original: DynamicViT (Rao et al., 2021) is used as the primary baseline, DeiT small (Touvron et al., 2021),
Revised: DynamicViT (Rao et al., 2021) is used as the primary baseline, DeiT small (Touvron et al., 2021),

Original: SPViT (Kong et al., 2021), IA-RED 2 (Pan et al., 2021a), S 2 ViTE (Chen et al., 2021), and HVT (Pan et al., 2021b).
Revised: SPViT (Kong et al., 2021), IA-RED 2 (Pan et al., 2021a), S 2 ViTE (Chen et al., 2021), and HVT (Pan et al., 2021b).

Original: We observe that the fine-tuned DeiT small model pruned by our Concrete masking significantly outperforms the fine-tuned DeiT tiny model with the same FLOPs trained from scratch.
Revised: We observe that the fine-tuned DeiT small model pruned by our Concrete masking significantly outperforms the fine-tuned DeiT tiny model with the same FLOPs trained from scratch.

Original: DynamicViT with DeiT small , on the other hand, results in a model with much lower performance than DeiT tiny .
Revised: DynamicViT with DeiT small , on the other hand, results in a model with much lower performance than DeiT tiny .

Original: This is good news for practitioners, as this suggests that we may not need to train
Revised: This is good news for practitioners, as this suggests that we may not need to train

Original: Transformers at different scales from scratch, to deploy the model to systems or applications with different computational budgets.
Revised: Transformers at different scales from scratch, to deploy the model to systems or applications with different computational budgets.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this paper, we propose a sparse token Transformer architecture based on a novel token pruning method, Sparse Token Transformer with Attention Back Tracking (STTABT), with two novel ideas – Attention Back-Tracing (ABT) and Concrete masking.
Revised: In this paper, we propose a sparse token Transformer architecture based on a novel token pruning method, Sparse Token Transformer with Attention Back Tracking (STTABT), with two novel ideas – Attention Back-Tracing (ABT) and Concrete masking.

Original: 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).
Revised: 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).

Original: Then, we propose a novel thresholding method named as Concrete masking, which adapts the idea from Concrete dropout to learn the appropriate token-wise pruning threshold given the desired amount of tokens to be retained.
Revised: Then, we propose a novel thresholding method named as Concrete masking, which adapts the idea from Concrete dropout to learn the appropriate token-wise pruning threshold given the desired amount of tokens to be retained.

Original: We evaluate the performance of our method on classification tasks from NLP and CV domains, by pruning tokens from the trained Transformer encoder-based language models and Vision Transformers.
Revised: We evaluate the performance of our method on classification tasks from NLP and CV domains, by pruning tokens from the trained Transformer encoder-based language models and Vision Transformers.

Original: Our method shows a better accuracy-efficiency trade-off 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.
Revised: Our method shows a better accuracy-efficiency trade-off 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.

Original: 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.
Revised: 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.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: We will introduce code construction and data collection in this section for reproducibility.
Revised: We will introduce code construction and data collection in this section for reproducibility.

Original: Model Implementation.
Revised: Model Implementation.

Original: First we construct our model with PyTorch (Paszke et al., 2019) and
Revised: First we construct our model with PyTorch (Paszke et al., 2019) and

Original: Huggingface (Wolf et al., 2020).
Revised: Huggingface (Wolf et al., 2020).

Original: We modify bert-base-uncased from Huggingface Models (Wolf et al., 2020).
Revised: We modify bert-base-uncased from Huggingface Models (Wolf et al., 2020).

Original: ApproxNet.
Revised: ApproxNet.

Original: We create ApproxNet (§3.2, trainer/glue base.by ) with the pretrained models.
Revised: We create ApproxNet (§3.2, trainer/glue base.by ) with the pretrained models.

Original: The approximation result is shown in Section A. ABT.
Revised: The approximation result is shown in Section A. ABT.

Original: We implement Algorithm 1 on top of bert-base-uncased (Devlin et al., 2019) ( models/sparse token.py ).
Revised: We implement Algorithm 1 on top of bert-base-uncased (Devlin et al., 2019) ( models/sparse token.py ).

Original: LTP,
Revised: LTP,

Original: Concrete masking.
Revised: Concrete masking.

Original: Then, we modify forward path to support LTP (Kim et al., 2022) and fix
Revised: Then, we modify forward path to support LTP (Kim et al., 2022) and fix

Original: ApproxSparseBertModel to implement Concrete masking (§3.3).
Revised: ApproxSparseBertModel to implement Concrete masking (§3.3).

Original: NLP training.
Revised: NLP training.

Original: The training for NLP tasks, we implemented our own training loop ( trainer/concrete trainer.py ) and hyper-parameter tuning settings inherits
Revised: The training for NLP tasks, we implemented our own training loop ( trainer/concrete trainer.py ) and hyper-parameter tuning settings inherits

Original: LTP hyper-parameter settings (Kim et al., 2022) ( main/ltp glue plot.py , main/concrete glue plot.py ).
Revised: LTP hyper-parameter settings (Kim et al., 2022) ( main/ltp glue plot.py , main/concrete glue plot.py ).

Original: We use the GLUE dataset from Huggingface Dataset (Wolf et al., 2020).
Revised: We use the GLUE dataset from Huggingface Dataset (Wolf et al., 2020).

Original: Vision training.
Revised: Vision training.

Original: We start from DynamicViT implementation (Rao et al., 2021) including data loader, training loop, training loss, and evaluation codes ( main/vit concrete end2end.py ).
Revised: We start from DynamicViT implementation (Rao et al., 2021) including data loader, training loop, training loss, and evaluation codes ( main/vit concrete end2end.py ).

Original: Therefore, we did not apply hyper parameter grid search on ViT pipeline.
Revised: Therefore, we did not apply hyper parameter grid search on ViT pipeline.

Original: And we remove hard token mask training stage, and we use larger λ mask = 100 than NLP task, because of lack of padding.
Revised: And we remove hard token mask training stage, and we use larger λ mask = 100 than NLP task, because of lack of padding.

Original: Also, we did not use prediction distillation loss for training ApproxNet for ViT.
Revised: Also, we did not use prediction distillation loss for training ApproxNet for ViT.

Original: The training dataset is ImageNet-1k (ILSVRC2012)
Revised: The training dataset is ImageNet-1k (ILSVRC2012)

Original: (Deng et al., 2009).
Revised: (Deng et al., 2009).

Original: We replace the Transformer encoder in ViT model (Dosovitskiy et al., 2021b) into our model code.
Revised: We replace the Transformer encoder in ViT model (Dosovitskiy et al., 2021b) into our model code.

Original: Our PyTorch model is modified to apply ViT changes from Transformer encoder (Dosovitskiy et al., 2021b)
Revised: Our PyTorch model is modified to apply ViT changes from Transformer encoder (Dosovitskiy et al., 2021b). A CKNOWLEDGMENTS

Revised: This work was partly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT)

Revised: (No. RS-2022-00187238,

Revised: Development of Large Korean Language Model Technology for Efficient Pre-training), (No.2022-000124, Development of Artificial Intelligence Technology for Self-Improving Competency-Aware

Original:  A PPENDIX A E XPERIMENTAL D ETAILS A. D ETAILS ON A TTENTION A PPROXIMATION N ETWORK
Revised: Learning Capabilities), and (No.2019-0-00075, 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 (NRF2018R1A5A1059921). A PPENDIX A E XPERIMENTAL D ETAILS A. D ETAILS ON A TTENTION A PPROXIMATION N ETWORK

Original: [Table]
Revised: [Table]

Original: In this section, we provide additional details on the attention approximation networks (ApproxNet).
Revised: In this section, we provide additional details on the attention approximation networks (ApproxNet).

Original: We train the ApproxNet following the instructions in Section 3.2.
Revised: We train the ApproxNet following the instructions in Section 3.2.

Original: As we mentioned, the architecture of the ApproxNet is identical with the original network such as BERT (Devlin et al., 2019) and Vision
Revised: As we mentioned, the architecture of the ApproxNet is identical with the original network such as BERT (Devlin et al., 2019) and Vision

Original: Transformer (ViT) (Touvron et al., 2021), except for the dimension of hidden and intermediate hidden unit size.
Revised: Transformer (ViT) (Touvron et al., 2021), except for the dimension of hidden and intermediate hidden unit size.

Original: We downsize the architecture by reducing the hidden unit dimension with different reducing factors.
Revised: We downsize the architecture by reducing the hidden unit dimension with different reducing factors.

Original: We use 4 and 8 as the reducing factors in all experiments of the paper.
Revised: We use 4 and 8 as the reducing factors in all experiments of the paper.

Original: We then train the
Revised: We then train the

Original: ApproxNet with the knowledge distillation (Wang et al., 2020; Jiao et al., 2020).
Revised: ApproxNet with the knowledge distillation (Wang et al., 2020; Jiao et al., 2020).

Original: Specifically, we use the embedding based distillation loss L embd , hidden state based distillation loss L hidn , attentionbased distillation loss L attn , and prediction-layer distillation loss L pred to train the ApproxNet for
Revised: Specifically, we use the embedding based distillation loss L embd , hidden state based distillation loss L hidn , attentionbased distillation loss L attn , and prediction-layer distillation loss L pred to train the ApproxNet for

Original: BERT.
Revised: BERT.

Original: For Vision Transformer, we do not use the prediction-layer distillation loss.
Revised: For Vision Transformer, we do not use the prediction-layer distillation loss.

Original: In Table A.1, we present the approximation error between the attention probability in the original network and the ApproxNet.
Revised: In Table A.1, we present the approximation error between the attention probability in the original network and the ApproxNet.

Original: We observe that the distillation helps the ApproxNet to approximate almost identical attention probability of the original network.
Revised: We observe that the distillation helps the ApproxNet to approximate almost identical attention probability of the original network.

Original: We also observe that the smaller reduce factor leads to lower attention approximation error.
Revised: We also observe that the smaller reduce factor leads to lower attention approximation error.

Original: However, we also find that the ApproxNet with a larger reducing factor also accurately approximate the attention probability of the original network.
Revised: However, we also find that the ApproxNet with a larger reducing factor also accurately approximate the attention probability of the original network.

Original: To boost the training speed for ApproxNet on downstream tasks, we first distill the pre-trained
Revised: To boost the training speed for ApproxNet on downstream tasks, we first distill the pre-trained

Original: Transformer encoders prior to the distillation of the fine-tuned network on each task.
Revised: Transformer encoders prior to the distillation of the fine-tuned network on each task.

Original: This procedure is not necessary if the downstream task is large enough or there is no downstream task such as
Revised: This procedure is not necessary if the downstream task is large enough or there is no downstream task such as

Original: ImageNet-1k.
Revised: ImageNet-1k.

Original: For instance, we directly train ApproxNet by setting the fine-tuned ViT as the teacher network with ImageNet-1k dataset.
Revised: For instance, we directly train ApproxNet by setting the fine-tuned ViT as the teacher network with ImageNet-1k dataset.

Original: However, for experiments on GLUE benchmark, we first train the ApproxNet by setting the original pre-trained BERT as the teacher network with the Wikitext103 dataset (Merity et al., 2017) for 200 epochs.
Revised: However, for experiments on GLUE benchmark, we first train the ApproxNet by setting the original pre-trained BERT as the teacher network with the Wikitext103 dataset (Merity et al., 2017) for 200 epochs.

Original: 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.
Revised: 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.

Original: B A DDITIONAL E XPERIMENTAL R ESULTS
Revised: B A DDITIONAL E XPERIMENTAL R ESULTS

Original: [Figure]
Revised: [Figure]

Original: a c y

Original: [Figure]

Original: s
Revised: y

Original: [Figure]

Original: l a

Original: [Figure]
Revised: [Figure]

Original: benchmark.

Original: STTABT@f x means the attention back-tracking by
Revised: STTABT@f x means the attention back-tracking by

Original: ApproxNet with factor x .
Revised: ApproxNet with factor x .

Original: B. A DDITIONAL R ESULTS ON GLUE EXPERIMENTS WITH FLOP S
Revised: B. A DDITIONAL R ESULTS ON GLUE EXPERIMENTS WITH FLOP S

Original: In Figure B.1, we present evaluation results on all GLUE benchmark with regards to FLOPs.
Revised: In Figure B.1, we present evaluation results on all GLUE benchmark with regards to FLOPs.

Original: We average every data point across 9 datasets in Figure B.1 to plot Figure 7a in the main paper.
Revised: We average every data point across 9 datasets in Figure B.1 to plot Figure 7a in the main paper.

Original: B. A DDITIONAL V ISUALIZATION E XAMPLES
Revised: B. A DDITIONAL V ISUALIZATION E XAMPLES

Original: In Figure B.2 and B.3, we visualize the additional visualization of the token pruning on each layer with STTABT (Concrete) and baselines (Kim et al., 2022;
Revised: In Figure B.2 and B.3, we visualize the additional visualization of the token pruning on each layer with STTABT (Concrete) and baselines (Kim et al., 2022;

Original: Rao et al., 2021).
Revised: Rao et al., 2021).

Original: We further confirm that our method prunes more tokens throughout all of transformer layers than baselines in both sentence and image classification tasks.
Revised: We further confirm that our method prunes more tokens throughout all of transformer layers than baselines in both sentence and image classification tasks.

Original: [Table]
Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure] [Table] 
Revised: [Figure] [Table] [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: C A BLATION S TUDY(
Revised: C A BLATION S TUDY

Original: S ECTION A DDED

Original: C.
Revised: C.

Original: P ROPAGATION F ACTORp
Revised: P ROPAGATION F ACTOR

Original: [Figure]
Revised: [Figure]

Original: We perform an ablation study on the importance score propagation factor p , which is introduced in
Revised: We perform an ablation study on the importance score propagation factor p , which is introduced in

Original: Algorithm 1.
Revised: Algorithm 1.

Original: We use CoLA dataset in GLUE benchmark (Wang et al., 2019) for this ablation study.
Revised: We use CoLA dataset in GLUE benchmark (Wang et al., 2019) for this ablation study.

Original: In Figure C.4, we observe that Attention Back-Tracking method with p value between the range of 0 . 0 - 0 .
Revised: In Figure C.4, we observe that Attention Back-Tracking method with p value between the range of 0 . 0 - 0 .

Original: 5 shows the similar task performance.
Revised: 5 shows the similar task performance.

Original: However, we also observe that using p value as 0 .performs the best.
Revised: However, we also observe that using p value as 0 .performs the best.

Original: C. C ONCRETE λ mask AND λ p
Revised: C. C ONCRETE λ mask AND λ p

Original: [Figure]
Revised: [Figure]

Original: We perform an ablation study on two balancing terms λ mask , λ p used in Concrete masking training, which is introduced in Section 3.3.
Revised: We perform an ablation study on two balancing terms λ mask , λ p used in Concrete masking training, which is introduced in Section 3.3.

Original: We use CoLA dataset in GLUE benchmark (Wang et al., 2019) for this ablation study.
Revised: We use CoLA dataset in GLUE benchmark (Wang et al., 2019) for this ablation study.

Original: In Figure C.5, we could not observe any clear pattern from varying λ p and λ mask .
Revised: In Figure C.5, we could not observe any clear pattern from varying λ p and λ mask .

Original: Therefore, it is better to adjust both values individually depending on the task.
Revised: Therefore, it is better to adjust both values individually depending on the task.

Original: For instance, for tasks in GLUE benchmark, we observe θ l does not change a lot during the training.
Revised: For instance, for tasks in GLUE benchmark, we observe θ l does not change a lot during the training.

Original: Therefore, we use a low λ p around 1 e − 3 in GLUE task training.
Revised: Therefore, we use a low λ p around 1 e − 3 in GLUE task training.

Original: On the other hand, we observe the model with concrete masking tend to fail in keeping the target token retention ration in the image classification task with ViT.
Revised: On the other hand, we observe the model with concrete masking tend to fail in keeping the target token retention ration in the image classification task with ViT.

Original: In this case, we use a high λ mask around 100 to keep the token retention ratio close to the desired target value.
Revised: In this case, we use a high λ mask around 100 to keep the token retention ratio close to the desired target value.

Original: D E XTRA B ASELINES ON V I T E XPERIMENT(S ECTION A DDED ) D. T RADE - OFF BETWEEN T OP -1 A CCURACY AND T OKEN R ETENTION R ATIO
Revised: D E XTRA B ASELINES ON V I T E XPERIMENT D. T RADE - OFF BETWEEN T OP -1 A CCURACY AND T OKEN R ETENTION R ATIO

Original: Technically, it is possible to apply our proposed method on any transformer-based architecture.
Revised: Technically, it is possible to apply our proposed method on any transformer-based architecture.

Original: To validate this, we perform additional experiments on the image classification task with another ViT architecture (Jiang et al., 2021).
Revised: To validate this, we perform additional experiments on the image classification task with another ViT architecture (Jiang et al., 2021).

Original: In Figure D.6, we additionally plot results of STTABT (Concrete masking) on LVViT (Jiang et al., 2021).
Revised: In Figure D.6, we additionally plot results of STTABT (Concrete masking) on LVViT (Jiang et al., 2021).

Original: Please note that Figure D.6 is the extended version of
Revised: Please note that Figure D.6 is the extended version of

Original: Figure 5 where we add experimental results on STTABT (Concrete) with LVViT and the original
Revised: Figure 5 where we add experimental results on STTABT (Concrete) with LVViT and the original

Original: LVViT (Jiang et al., 2021).
Revised: LVViT (Jiang et al., 2021).

Original: Concrete masking on LVViT results in accuracy drop of 3 .
Revised: Concrete masking on LVViT results in accuracy drop of 3 .

Original: 68% p when we drop a token retention ratio from 76% to 18% , while Concrete masking on DeiT suffers accuracy drop of 5 .
Revised: 68% p when we drop a token retention ratio from 76% to 18% , while Concrete masking on DeiT suffers accuracy drop of 5 .

Original: 53% p in the same range of token retention ratio.
Revised: 53% p in the same range of token retention ratio.

Original: This result shows our Concrete masking method preserves the performance of the original model by preserving important tokens better.
Revised: This result shows our Concrete masking method preserves the performance of the original model by preserving important tokens better.

Original: D. T RADE - OFF BETWEEN FLOP S AND T OKEN R ETENTION R ATIO
Revised: D. T RADE - OFF BETWEEN FLOP S AND T OKEN R ETENTION R ATIO

Original: We additionally plot the results of STTABT (Concrete masking) with LVViT small
Revised: We additionally plot the results of STTABT (Concrete masking) with LVViT small

Original: (Jiang et al., 2021) and DynamicViT (Rao et al., 2021) with LVViT small in Figure D.7.
Revised: (Jiang et al., 2021) and DynamicViT (Rao et al., 2021) with LVViT small in Figure D.7.

Original: Figure D.7 is an extended version of Figure 7b.
Revised: Figure D.7 is an extended version of Figure 7b.

Original: When we use LVViT as the backbone model, Concrete masking outperforms
Revised: When we use LVViT as the backbone model, Concrete masking outperforms

Original: DynamicViT.
Revised: DynamicViT.

Original: The additional baselines include MViTv2 (Li et al., 2022),
Revised: The additional baselines include MViTv2 (Li et al., 2022),

Original: SwinT (Liu et al., 2021), and LVViT (Jiang et al., 2021).
Revised: SwinT (Liu et al., 2021), and LVViT (Jiang et al., 2021).

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: DyViT-LV-S means DynamicViT with LVViT small .
Revised: DyViT-LV-S means DynamicViT with LVViT small .

Original: Note that this figure is the  version 
Revised: Note that this figure is the extended version of