Paper: Skip-Attention: Improving Vision Transformers by Paying Less Attention (ICLR-2024)

Reviewer comment (Novelty):
It seems that the motivation of this paper has been mentioned in Zhou et al. (Refiner: Refining self-attention for vision transformers). They observe that reusing the attention maps in the next transformer block does not brings performance drop. The authors should more clearly explain the differences between this paper and the work mentioned above.

Author response:
**1. Comparison with Refiner**

Thank you for the pointer. We discuss Refiner (Zhou et al.) in Section 2 under ``Efficient attention" (6th line from the end of the paragraph). 

We acknowledge that both SkipAT and Refiner leverage the high correlation across layers of a transformer, but for two different purpose: increasing the classification accuracy (by Refiner) and increasing the computational efficiency (by SkipAT). Refiner introduces additional convolutions on top of ViT to de-correlate the attention maps to increase the expressive power of the model. This obviously comes at an additional computational cost both in terms of parameters and latency compared to the vanilla ViT. In contrast, SkipAt removes the correlated MSA blocks and approximate them using a cheap convolutional alternative leading to a reduced computation compared to the vanilla ViT.

On ImageNet-1K, Refiner achieves a top-1 accuracy of 80.3\% using 2M more parameters than SkipAT, which achieves a similar accuracy of 80.2\%. We shall clarify and further emphasize the novelty of SkipAT with respect to Refiner in the paper.

Reviewer comment (Experiments):
The baselines used in this paper are not recently proposed. The results are already not state-of-the-art compared to recent works, like CMT (CVPR'2022). I would like to see how would the performance go when the proposed approach is applied to recent state-of-the-art ViT models as they mostly did not change the self-attention part.

Author response:
**2. Comparison with CMT**

We thank the reviewer for providing the reference. In Table 2(a), we compare SkipAT with methods that focus on improving the efficiency of transformers without changing the backbone. We also provide a comprehensive list of SoTA methods in Table 7 in the Appendix. We shall add CMT to Table 7. 

In Table 2(a), we apply SkipAT to SoTA transformer backbones such as PvT and LIT, and show that SkipAT is agnostic to transformer architecture and improves both top-1 accuracy and throughput of these backbones.

Reviewer comment (Experiments):
Many ViT models are based on window self-attention, which is originally proposed in Swin Transformer (ICCV'2021). The authors have shown that the proposed method works well for the original self-attention. So, how would the performance go when the proposed method is applied to ViTs with window self-attention.

Author response:
**3. Use of SkipAT with window self-attention**

In Section 4.1 -- ``Image Denoising" and Table 4, we describe the application of SkipAT to Uformer, which is a U-shaped hierarchical network with Swin transformer blocks as the encoder and decoder. We skip the whole window self-attention (WSA) block in each decoder block by reusing attention of the corresponding encoder block via the SkipAT parametric function. A detailed description is provided in Section 6.2 in Appendix.

We observe from Table 4, that  SkipAT outperforms the baseline Uformer variants with the 25\% higher
throughput on average.

Reviewer comment (Novelty):
I care more about the novelty of this paper as the originality of this paper has been mentioned in a previous work. If the authors make this clear, I would like to raise the ranking score.

Author response:
**1. Comparison with Refiner**

Thank you for the pointer. We discuss Refiner (Zhou et al.) in Section 2 under ``Efficient attention" (6th line from the end of the paragraph). 

We acknowledge that both SkipAT and Refiner leverage the high correlation across layers of a transformer, but for two different purpose: increasing the classification accuracy (by Refiner) and increasing the computational efficiency (by SkipAT). Refiner introduces additional convolutions on top of ViT to de-correlate the attention maps to increase the expressive power of the model. This obviously comes at an additional computational cost both in terms of parameters and latency compared to the vanilla ViT. In contrast, SkipAt removes the correlated MSA blocks and approximate them using a cheap convolutional alternative leading to a reduced computation compared to the vanilla ViT.

On ImageNet-1K, Refiner achieves a top-1 accuracy of 80.3\% using 2M more parameters than SkipAT, which achieves a similar accuracy of 80.2\%. We shall clarify and further emphasize the novelty of SkipAT with respect to Refiner in the paper.

Reviewer comment (Experiments):
Based on the analysis in Section 3.2, it makes sense for the authors to apply their method from layer 2 to 8. However, it is not convincing for different pretrained ViTs to skip layer 2 to 8 as well if considering different training objectives or pretrained datasets. Thus, it would be better for the authors to study if other pretrained ViTs (MAE [A], DINOv2 [B], SAM [C]), have the same phenomenon.

Author response:
**1. Different pretrained ViTs**

Thanks for raising this interesting question: To what extent the correlation analysis from a supervised image classification model, as in Section 3.2, generalize to other tasks and datasets?
We empirically observe that the optimal configurations of parametric function is determined mostly by the transformer architecture rather than the task.

 When considering different training objectives, we applied SkipAT to DINO in Section 7 of the Appendix under ``Self-Supervised Learning with DINO". We observe that SkipAT achieves almost the same performance as fully trained DINO with around 26\% less training time (73.3\% in 96 GPUhours vs. 73.6\% in 131 GPU-hours). 
Moreover, we also use SkipAT on semantic segmentation on ADE20K (Table 3) and Image Denoising on SIDD (Table 4), showing SkipAT is generalizable to different training objectives.

For other architectures, a principled way could be to compute the correlation matrix by computing the CKA between different layers and using a threshold to identify the layers on which SkipAT can be applied, as we have done for PvT and LIT on Table 2(a).

Reviewer comment (Experiments):
Introducing convolution into ViTs has shown to be effective in related works [D], which is intuitive to me to achieve performance gain for SKIPAT. In this paper, SKIPAT adopts convFFN as a parametric function to replace MSAs, which still needs to be trained from scratch in order to achieve efficiency gain. It would be promising if this parametric function can be used as a drop-in replacement for existing large ViTs.

Author response:
**2. Using SkipAT in pretrained models**

Indeed finding a parametric function, which can be plugged into any pretrained model is practically coveted but is very challenging given the complex interplay between the layers of a deep neural network.
Surprisingly, our experiments demonstrate that using an identity function on a pretrained ViT-T/16 in a plug-and-play fashion could still perform reasonably well: only 9.3\% drop in performance, when copying self-attention matrix $A$ in layers 3-8 (Table 9 of the Appendix). 

Regularizing transformer blocks to be more correlated during training could be an step toward more robust plug-and-play adaptors.

Reviewer comment (Experiments):
Can the authors specify more on the experimental setting of applying SKIPAT into hierarchical ViTs? I can understand that SKIPAT works for layer 2 to 8 in plain ViTs. But it is not intuitive to me how to select the layers to skip in PVT, LIT, etc.

Author response:
**3. Experimental setting in hierarchical ViTs**

*Using SkipAT in PvT*

As reported in Table 2, we applied SkipAT on a hierarchical architecture PvT-S, which consist of 4 stages and each stage consists of 3, 3, 6 and 3 transformer blocks respectively. We observe high correlation of $Z^{MSA}$ in the third stage and apply SkipAT in the intermediate blocks i.e. between 2 through 5 blocks. In this case, we still have applied SkipAT to adapt the features of the same dimensionality. However, by increasing the stride of the depth-wise separable convolutions in the parametric function, SkipAT can adapt features across different resolutions. 

*Using SkipAT in LIT*

As reported in Table 2, we applied SkipAT in LIT-T and LIT-S, which consist of 4 stages. The first two stages consists of convolution layers and the last two stages consists of transformer blocks. We apply SkipAT in the intermediate transformer blocks in the third stage of LIT i.e. we skip layers 2 to 5 in the third stage.

We will clarify this in the text.

Reviewer comment (Experiments):
Some experiments can be added which decouple the improvements obtained with convolutions vs the SkipAt formulation.

Author response:
Thank you for this interesting suggestion. Based on the reviewer's suggestion, we implemented the SkipAT parametric function across all layers of VIT-T/16 and train for 100 epochs on ImageNet-1K. The default configuration, which employs SkipAT at layers 3 to 8, yields a top-1 accuracy of 67.7\% (Table 5 in the main paper). When SkipAT is applied to layers 1 through 12, the top-1 accuracy decreases to 54.1\%, indicating a significant drop of 13.6\%. This suggests that leveraging the parametric function is only beneficial for approximating MSA blocks with high correlation. 

Moreover, we observe that using an Identity function, where we simply copying $Z^{MSA}$ features across 3 to 8 layers, achieves decently high top-1 accuracy of 61.1\% (Table 5), witnessing high correlation in MSA blocks. Thus, we hypothesize that the performance gains is due to the SkipAT formulation. Additionally, using more powerful parametric function such as convolution layer marginally increases performance (top-1 accuracy: 65.4\%) but increases throughput by 39\% (8500 vs. 5200 im/sec) (Table 5).

We understand the reviewer asks us to replace SkipAT parametric function $\phi(Z^{MSA}_{l-1})$ with an Identity function (i.e. using features from prev. blocks). If so, we mentioned the effect of using an Identity function above. 

We are happy to follow-up with the reviewer for any clarifications or additional ablations.

Reviewer comment (Experiments):
The main motivation of SkipAt is that the output representations of MSA in vision transformers are redundant. So, the paper claims that just refining the outputs of the previous MSA layers is sufficient. However, SkipAt consists of depthwise convolutions which are also quite powerful modules themselves, so it is unclear if the throughput gains come just by the convolutions rather than the SkipAt formulation. I suggest that the authors run the couple of ablations below, If these ablations reach lower accuracy, it would be convincing that the SkipAt formulation is responsible for the accuracy gains.

Author response:
Thank you for this interesting suggestion. Based on the reviewer's suggestion, we implemented the SkipAT parametric function across all layers of VIT-T/16 and train for 100 epochs on ImageNet-1K. The default configuration, which employs SkipAT at layers 3 to 8, yields a top-1 accuracy of 67.7\% (Table 5 in the main paper). When SkipAT is applied to layers 1 through 12, the top-1 accuracy decreases to 54.1\%, indicating a significant drop of 13.6\%. This suggests that leveraging the parametric function is only beneficial for approximating MSA blocks with high correlation. 

Moreover, we observe that using an Identity function, where we simply copying $Z^{MSA}$ features across 3 to 8 layers, achieves decently high top-1 accuracy of 61.1\% (Table 5), witnessing high correlation in MSA blocks. Thus, we hypothesize that the performance gains is due to the SkipAT formulation. Additionally, using more powerful parametric function such as convolution layer marginally increases performance (top-1 accuracy: 65.4\%) but increases throughput by 39\% (8500 vs. 5200 im/sec) (Table 5).

We understand the reviewer asks us to replace SkipAT parametric function $\phi(Z^{MSA}_{l-1})$ with an Identity function (i.e. using features from prev. blocks). If so, we mentioned the effect of using an Identity function above. 

We are happy to follow-up with the reviewer for any clarifications or additional ablations.

Reviewer comment (Experiments):
Replace all layers with SkipAt instead of just layers from 3 through 8. According to the authors hypothesis, since layers 9 though 12 have lesser correlation, using SkipAt at these layers should hurt accuracy.

Author response:
Thank you for this interesting suggestion. Based on the reviewer's suggestion, we implemented the SkipAT parametric function across all layers of VIT-T/16 and train for 100 epochs on ImageNet-1K. The default configuration, which employs SkipAT at layers 3 to 8, yields a top-1 accuracy of 67.7\% (Table 5 in the main paper). When SkipAT is applied to layers 1 through 12, the top-1 accuracy decreases to 54.1\%, indicating a significant drop of 13.6\%. This suggests that leveraging the parametric function is only beneficial for approximating MSA blocks with high correlation. 

Moreover, we observe that using an Identity function, where we simply copying $Z^{MSA}$ features across 3 to 8 layers, achieves decently high top-1 accuracy of 61.1\% (Table 5), witnessing high correlation in MSA blocks. Thus, we hypothesize that the performance gains is due to the SkipAT formulation. Additionally, using more powerful parametric function such as convolution layer marginally increases performance (top-1 accuracy: 65.4\%) but increases throughput by 39\% (8500 vs. 5200 im/sec) (Table 5).

We understand the reviewer asks us to replace SkipAT parametric function $\phi(Z^{MSA}_{l-1})$ with an Identity function (i.e. using features from prev. blocks). If so, we mentioned the effect of using an Identity function above. 

We are happy to follow-up with the reviewer for any clarifications or additional ablations.

Reviewer comment (Experiments):
The authors test their module on Ti, B and S which all have 12 layers so they recommend to use SkipAt layers from 3 through 8. How do they recommend tuning these for larger depths?

Author response:
The reviewer raises an interesting point. 

For large models, a principled approach would be to compute the correlation matrix using the CKA analysis of $Z^{MSA}$ for every layer. Using a threshold on this correlation matrix would indicate the layers with high correlation, on which the SkipAT parametric function can be applied. 
We use this approach to apply SKipAT on pyramid (PvT) and hybrid (LIT) transformer architectures (Table 2(a)) where arrangement of transformer blocks are different than ViT. For example, PvT-S consists of 4 stages with 3, 3, 6 and 3 transformer blocks at each stage. We observe high correlation of $Z^{MSA}$ in the third stage and apply SkipAT in the intermediate blocks i.e. between 2 through 5 blocks.

Another approach could be using more advanced techniques such as Neural Architectural Search (NAS). While NAS might propose valid solutions, it is computationally more expensive.

Reviewer comment (Theory):
In page 6, authors say that $n >> d$ and so $O(n^d)$ term dominates. I would suggest the authors add that this is specific to dense prediction tasks, since for image classification even for a S model (d=384, and n=196), so the claim that n >> d is not general.

Author response:
Thank you for the precise point: Indeed the claim holds for dense prediction tasks. We shall make this change.

Reviewer comment (Evaluation):
Are the throughput increase in Fig a) significant? What do the error bars look like?

Author response:
We report the mean and standard deviation of the throughput of SkipAT for 10 runs on ImageNet-1K using ViT-T/16, ViT-S/16 and ViT-B/16. For a single run, we measure throughput (image/sec $\times 10^3$) with a batch size of 1024 on a single NVIDIA A100 GPU, averaged over the validation set of ImageNet-1K. The results are in the table below:

| Arch               | Throughput (im/sec $\times 10^3$) |
| ------------------ | ---------------------------------- |
| ViT-T/16           | 5.9 $\pm$ 0.3                     |
| SkipAT + ViT-T/16  | 6.8 $\pm$ 0.2                     |
| ViT-S/16           | 3.1 $\pm$ 0.3                     |
| SkipAT + ViT-S/16  | 3.7 $\pm$ 0.3                     |
| ViT-B/16           | 1.2 $\pm$ 0.2                     |
| SkipAT + ViT-B/16  | 1.6 $\pm$ 0.1                     |

Reviewer comment (Presentation):
The authors can consider making the numbers in Figure 3 and Figure 5 bigger.

Author response:
Thank you very much for these insightful comments and suggestions. We shall incorporate all of them in the camera-ready version.

Reviewer comment (Presentation):
It is clear by comparing Fig 3 b) and Fig 5 b), that Fig 5 b) has lower correlation. The authors can also have a line plot where the x-axis is the layer id and y-axis is the average correlation of layers before it. If we plot the baseline ViT and the ViT with SkipAt on the same graph, it will make the comparison even clearer.

Author response:
Thank you very much for these insightful comments and suggestions. We shall incorporate all of them in the camera-ready version.

Reviewer comment (Presentation):
In Fig 1), are the circles to indicate #params necessary since they are roughly the same? It gives the impression that the improvements are not significant even if they are, since the circles overlap.

Author response:
Thank you very much for these insightful comments and suggestions. We shall incorporate all of them in the camera-ready version.