Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes a lightweight prediction module to estimate the importance score of each token in the vision transformer. The module is added to different layers to prune redundant tokens hierarchically. To optimize the prediction module in an end-to-end manner, an attention masking strategy is proposed to differentiably prune a token by blocking its interactions with other tokens. Experiments show that the proposed method can significantly reduce the computational cost with a small drop in accuracy.
main_review: I think that it is an interesting idea to utilize the natural properties of the vision transformer for network pruning. Maybe it is the first paper in this topic. The experimental results are also pretty good. The writing is clear and professional. However, this paper is only evaluated for image classification. I’m wondering if the proposed pruning affects much for downstream vision tasks (e.g., object detection and semantic segmentation)?
limitations_and_societal_impact: None
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The intuition of this paper is that attention is sparse in vision transformers. To this end, this paper proposes a pruning method to estimate the importance score of each token before self-attention layer. An attention masking strategy is proposed to differentiably prune a token by blocking its interactions with other tokens. Experiments on DeiT and LV-ViT demonstrate the effectiveness of the proposed method.
main_review: No ablation studies on the loss functions

Although it's a pruning strategy on the DeiT and LV-ViT, the proposed method is act as a downsampling operation (L120 use C′ = C/2 in our implementation).
However, the baseline is only set to the model doesn't use the pruning method, so baseline comparison with the convolution based downsampling architecture is needed. Comparing with those method in Table 2 (eg PVT) is not enough. The authors need to design fair experiment setting to show the pruning downsampling is better than the convolution downsampling.
limitations_and_societal_impact: see above
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: Based on the observation that the final prediction in ViT is only decided by the most informative tokens, the author proposes DynamicViT which adopts dynamic token sparsification framework to prune redundant tokens progressively and hierarchically dependent on the input. To optimize the pruning in an end-to-end and differentiable way, the author proposes the attention masking strategy. The hierarchically pruning largely reduces the FLOPs and increases the throughout with the price of a very negligible performance drop for ViT, which demonstrates a very good complexity-accuracy trade-off achieved by the proposed DynamicViT.
main_review: Strengths:
1) The paper is generally well-written and easy to follow. The code is also provided for reproducibility.
2) The DynamicViT model is inspired by a very interesting observation and the author conducts extensive experiments to confirm the correctness of the assumption.
3) The proposed attention masking strategy enables the DynamicViT module to be trained in a differentiable end-to-end manner.
4) DynamicViT achieves a very promising complexity-accuracy trade-off compared to CNNs and ViT approaches.

Weaknesses:
1) The hierarchically pruning of DynamicViT works well for the classification tasks. But it seems hard for DynamicViT to extend to object detection & semantic segmentation tasks, which limits the generalization and application scenarios of the proposed approach.
limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper introduces a dynamic network to achieve efficient inference for vision transformers. Specifically, the proposed framework adopts a multi-stage architecture. In each stage, a learnable mask with the Gumbel-softmax training strategy is generated to remove some tokens for efficiency. The distillation loss and fine-tuning process are used for higher performance.
main_review: Strength:
+ Albeit the dynamic mechanism has been widely explored on Transformers in NLP tasks and CNN networks in Image classification tasks, introducing dynamic token sparsiﬁcation for Vision Transformers is new and interesting.
+ The visualization and empirical analysis are clear and make sense.
+ The paper is generally well written.

Weakness:
- The experimental comparison is not that fair. The proposed framework uses additional distillation losses on features and predictions but the baseline methods do not. According to the observations in DeiT[21], the distillation procedure is essential for classification accuracy. Although some ablation studies are provided in the supplementary material, the performance of the baseline framework with the distillation procedure has not been revealed.
- Compared with the baseline methods, the DynamicViT adopts an additional fine-tuning procedure with 30 epochs. The fine-tuning procedure with distillation could significantly improve the performance of Vision Transformers, but the paper has not reported it.
- Can the DynamicViT be trained in an end-to-end manner? How about the performance if it is trained from scratch?

After Rebuttal:
The response addressed my concerns. I lean to vote for acceptance.
limitations_and_societal_impact: Only partially, more discussion is needed.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper works on accelerating the inference of ViT models by dynamically dropping tokens. This is realized by learning a light-weighted importance mask prediction between ViT blocks. The mask prediction module contains a few MLP layers and runs in $O(N)$ time instead of $O(N^2)$ in the self-attention. The authors also identified and resolved a few technical challenges in practice: 1. differentiating into the hard sampling; 2. batching after pruning. Experiments show the proposed method improves throughput by 43% ~ 54% with 0.2 ~ 0.5 accuracy drop.
main_review: ### strengths

- The overall results of improving throughput by 43% ~ 54% with 0.2 ~ 0.5 accuracy drop is favorable, and is likely to be used in practice in future ViT-based models.

 - The technical contribution of using MLPs on each individual token together with a global pooling of tokens to predict importance score makes sense. This is more efficient than the straightforward baseline of using the attention mask as the pruning weight (it will be great if such a baseline comparison is provided in the paper).

 - I like the model scaling figure (Fig. 4). It shows the proposed method works under large backbones.

 - Code is provided in the supplement.

### Weaknesses

 - While the overall throughput improvements are convincing, however I didn't understand where are the speed improvements from. In L143, the authors say in practice they don't actually remove the tokens due to batching. If this is the case, should the actually FLOPs keep unchanged? Is it the same for both training and testing? It seems that testing is also batched (Table. 1 caption). Please make sure to clarify this in the rebuttal.

 - The paper does not compare to other efficient ViT baselines. I know there are no neighbors directly close to this work, but some straightforward counterparts I have in mind are the O(n log n) transformers, or just directly pruning using the attention weight (e.g., the cross-attention weight of the class token). I will consider increasing my rating if such experiments are provided in the rebuttal and the results are interesting.

 - I don't quite like the distillation loss and KL loss. First, why do we need both distillation loss and KL loss? They seem pretty similar. It would be better if the author can provide ablations of both losses. Note this is NOT a required experiment in the rebuttal.

 - The distillation and Eq 9-11 seem to make the training much slower. Please report the training time (e.g., with/ without distillation) for reference. Much slower training time will NOT decrease my rating.

 - I don't get the problem of `the zeroed tokens will still influence other tokens through the Softmax operation` in L146. Would weight the zeroed token by negative infinite a simple solution, instead of using EQ 9-11?

 - The term `structural downsampling` in Fig. 1 is a bit confusing. From the context, it seems this just refers to pooling. `Structural downsampling` causes confusion to the term `structural pruning` in the network sliming literature, which refers to filtering the weights.

  - Should $D$ in Eq. 2 be $\hat{\bf D}$?

### Rating

  Overall this is a good paper with favorable experimental results, novel technical contribution, and good application value. Most of my negative comments are about writing, and I hope the authors can clarify them in the rebuttal.

limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: A token sparsification method is proposed in this paper for speeding up ViT. Attention masking is added into the ViT framework with learnable dropping out via back-propagations.  Unessential tokens are gradually pruned via attention masking for inference acceleration. The modified ViT, named dynamicViT, maintains the major accuracy while reducing one third FLOPs.
main_review: The main idea of this work is to develop an attention mask strategy for discarding unimportant tokens. The importance is validated by the attention mask which is learnable via gradient back-propagations. This design functions as a module that is integrated into ViTs. It splits one ViT into two parts. The first part maintains the original operations of ViT while the second part proceeds with reduced tokens.  During training, a knowledge distillation based method is introduced to train dynamicViT based on the original ViT.

The whole framework is to reduce the tokens while maintaining accuracy. This seems overall fine. There are a few concerns raised regarding the methodology, presentation, and experiments.

1. In Fig. 2, there are $N_1$ blocks before sparsification and there are $N_2$ blocks afterward. It is not clear how $N_1$ and $N_2$ are set according to different transformers.  Moreover, it seems a tradeoff to set $N_1$ and $N_2$ for acceleration and accuracy.

2. The proposed framework is posed as a knowledge distillation problem, which is specified in Sec. 3.4. It raises an interesting thinking that this approach aims at reducing tokens rather than reducing network parameters (e.g., linear projection parameters). It sounds like DeiT and DynamicViT function similarly (i.e., pruning network params v.s. pruning tokens). A discussion upon DeiT would be helpful to improve the motivation. 

3. In the experiments, dynamicViT only focuses on LV-VIT architecture. As the proposed method seems general, is it suitable for different transformers? If so, how? (see q1).

4. The attention mask is learned by gradient back-propagation. How the computed loss will affect this mask to discard unimportant tokens. More explanation on this aspect would make the technique more convincing. For instance, how the mask looks like? What if the majority value of the mask is 0? Since the mask is influenced by the gradients, the values of the mask are sometimes not satisfactory. Have the authors considered any regularization upon that?

5. Eq. 9-11 seems to overcomplicate the zero token issue. Reshaping the remaining tokens to make them compact is a way to ensure the following computations.

----------------------------------------------post rebuttal---------------------------------------------

While the proposed token sparsification seems not suitable for SWIN-based transformers where the tokens are densely organized for further inference. The authors have addressed most of the concerns. This reviewer turns positive on the current submission.
limitations_and_societal_impact: See above. The negative societal impact is not discussed. The checklist is shown in the supplementary file, which does not follow the guideline that it must appear in the submitted PDF, immediately after references.
review_title: null

Meta-review: All the reviewers agree that the observations made by this submission are interesting and the proposed dynamic token sparsification method is novel and effective. Some concerns on the experiment details are raised by the reviewers but the authors address them well in the response. A clear acceptance. 
Decision: Accept (Poster)