Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
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. 

Although 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.

In terms of computational efficiency, latency instead of FLOPs might be more important for practitioners.

An ablation study and intrinsic evaluation of each component (ABT and Concrete Masking) are necessary.

Although 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. 

clarity,_quality,_novelty_and_reproducibility: This paper is clearly written and easy to follow. For example, Algorithm 1 and all equations are helpful in understanding the exact mechanism of attention back tracking and Concrete Masking. In my view, figures are somewhat complicated. Making it clearer might be better.

The flexibility in choosing initial values is good. On the other hand, I presume it may increase the complexity of hyperparameter space. The authors should provide the detail on how they set those values and how difficult or costly to decide them.

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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.

The 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.

The 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.

The paper then presents an extensive set of experiments, both for language and vision models.

Figure 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.

Figure 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).

Figure 5 shows STTABT beats DynamicVIT for pretty much every retention ratio.
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.

In 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.

A few questions:

- What's the size for VIT? is it Base?

- What's the impact of hyper parameter p (defined at the end of page 4)? Can we see performance as a function of p? 

- 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.

[1] = Learning to Merge Tokens in Vision Transformers, Renggli et al.
[2] = TokenLearner: Adaptive Space-Time Tokenization for Videos, Ryoo et al.
clarity,_quality,_novelty_and_reproducibility: The paper is clearly written, and the idea seems novel.
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.

While the algorithm is a bit intricate, I think there's value in showing that using attention scores to select tokens can indeed work well.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
strength_and_weaknesses: Strength:
1. The idea of using approximated network during the token pruning seems new in literature.
2. Promising accuracy and computation trade-off in the low token retention ratio regime.

Weaknesses:
1. 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. 
2. 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.

clarity,_quality,_novelty_and_reproducibility: Clarity:
The experiments section is well written. However, the method section (subsection 3.2 and subsection 3.3) may need some modifications to further improve the readability.

Quality:
The paper proposes a novel backtracking pruning method and validates it in both NLP and CV tasks. 

Novelty:
The proposed method is new in the token pruning literature.

Reproducibility:
Good reproducibility. The authors include a reproducibility statement section and the code is attached for review. I checked the major part of the codes and it matches the proposed method in the main paper.
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.

Minor issues:

1. 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.

2. 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.


I 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.

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.

I have some concerns about the clarity of the method and empirical evaluation.

How is ApproxNet trained? Is it trained along with the main Transformer model or separately trained after the main model training is done?

What does recursive pruning mean in Figure 1? I understand the main idea of the caption but find the figure hard to interpret.

What 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?

I 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.

How 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.

clarity,_quality,_novelty_and_reproducibility:  The idea makes sense but the clarity of the method needs to be greatly improved. Given the current form, it would be difficult to reproduce the presented method.

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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper proposes an algorithm to prune tokens at inference time in transformer models. Two main ingredients are proposed: an attention approximation network (ApproxNet) and concrete masking. The authors addressed the main concerns of the reviewers in their rebuttal and in the updated manuscript. All reviewers feel positively about this paper, and so do I. I therefore recommend acceptance. 


justification_for_why_not_higher_score: While the paper is interesting and I recommend acceptance, I don't see enough novelty in the proposed method nor potential impact to justify a spotlight or oral presentation. 
justification_for_why_not_lower_score: This is a valuable paper that seems of interest to the research community.
Decision: Accept: poster