Skip-Attention: Improving Vision Transformers by Paying Less Attention

We appreciate the Reviewer f6zn's valuable feedback. We address the concerns as follows:

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.

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.

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.

4. Comparison with DynamicViT and EViT

We compare SkipAT with DynamicViT in Table 7 in the Appendix. When using ViT-T/16 as the backbone, SkipAT outperforms DynamicViT by 2% in terms of top-1 accuracy and also achieves higher throughout (6900 vs 6100 im/sec).

Comparing EViT [1*] with SkipAT using ViT-S/16 as backbone, SkipAT has comparable top-1 accuracy (80.2% of SkipAT vs. 79.8% of EViT) and comparable throughput (3800 vs/ 3900 im/sec). We shall add these results in Table 7.

[1*] Liang et al., Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations, ICLR 2022

We appreciate the Reviewer iYVP's valuable feedback. We address the concerns as follows:

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

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.

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.

We appreciate the Reviewer VJcB's valuable feedback. We address the concerns as follows:

1. Additional ablations

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.

2. Using larger models

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.

3. Complexity for dense prediction tasks

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

4. Error bars of throughput

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:

5. Minor comments and suggestions

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