On the Surprising Effectiveness of Attention Transfer for Vision Transformers

We’re glad that you found our results interesting and our analysis comprehensive. We address your questions below:

I’m not sure if it will make its way into practical use, for either the Attention Copy or Attention Distillation variants.

We agree that the attention transfer methods are not currently ready for practical use. One of our main objectives was to use attention transfer to understand the role of features vs attention in pre-trained models. But beyond scientific utility, attention transfer may have some advantages over fine-tuning in the future:

• An attention map (of size LxL, L is the sequence length) does not depend on the model’s dimension, which means the map from a model can be directly transferred to a different-sized model.
• Attention transfer gets rid of layer-wise learning rate decay. This is a crucial hyper-parameter almost used everywhere for tuning pre-trained vision models (beyond ViT). The fundamental prior here is that early layers should change less compared to later layers. But such a prior can be a restriction for next-generation models and getting rid of it opens up new opportunities.
• Sharing weights can incur security risks (e.g., white-box attacks). In such settings, we need an effective way to transfer knowledge from the pre-trained model, and attention transfer offers such a possibility.

We will add more detail on avenues for future research in the paper, and we hope other researchers find more practical ways to use attention transfer!

For some of the analysis results, it would be very helpful to see how Attention Distillation performs as compared to Attention Copy.

In our main rebuttal above, we have updated tables so that we have both Attention Copy and Attention Distillation. In general, Attention Distillation performs better than Attention Copy, which follows our existing findings.

Revising attention map visualizations

We have modified the figures to be easier to see. We have included samples in the 1 page PDF. Let us know if you have any further suggestions!

In the Attention Copy setting, are the unnecessary Q and K layers removed from the student network?

Yes, we remove them from the student network.

Why was more attention and testing not performed on Q transfer?

Thank you for pointing this out! See [a] in the main rebuttal above.

Thank you for the feedback on our work! Below, we respond to your questions and comments:

In Table 1, transfer Q is better than transfer Q,K (which is equivalent to the attention map), does it mean distillation on features would achieve better performance?

See the main rebuttal above for a detailed discussion of the Q-copying result. For feature distillation, we followed your suggestion and tried distilling the residual stream features from a pre-trained MAE ViT-L. In our preliminary results, we obtain a downstream accuracy of 81.3 on ImageNet-1k. This is significantly lower than the 85.7 that can be achieved through fine-tuning or attention distillation. This makes sense: the features learned during self-supervised pre-training are not directly well-suited for classification, so trying to match them can hurt performance. CKA analysis of the features (Fig. 5) supports this hypothesis – the fine-tuned MAE does well by significantly changing the features in the latter half of the network. Overall, transferring attention appears to do much better than distilling the features.

Transfer tasks (e.g., Table 2 and 3) should contain the results of training from scratch.

In our main rebuttal above, we have updated Table 2-3 with the results of training from scratch.

There should be an extra Table or Figure containing the result with other similar settings (e.g., attention copy from only pre-trained model) for better comparison.

Thanks for the suggestion! We will add a new subsection and table that shows the effect of the teacher (pre-trained vs fine-tuned) on attention copy and distillation. We hope this will be more organized for readers.

We’re glad that you found our work well-written, with extensive experiments and interesting findings. Below, we respond to your questions and comments:

The findings is somewhat similar to previous works which apply attention transfer in ConvNets. The only different things is in ViT, attention can be explicitly represented, which eases the operation of attention transfer.

Previous works on ConvNets have mainly looked at transferring different properties of the features, like spatial feature magnitudes. They have also mainly been conducted in the knowledge distillation paradigm, where a task-specific (not pre-trained) downstream teacher is distilled into a smaller student. Our work differs from prior work in a few ways. First, as you point out, ViTs allow us to explicitly decouple the attention patterns from the features that they combine. Second, we extensively investigate the pre-training/fine-tuning paradigm, and we surprisingly find that the attention maps, by themselves, are often sufficient to achieve the same downstream performance. This result is completely new and calls into question the “feature learning” story that typically motivates pre-training in vision.

The authors seem to over-claim something, e.g., "offer a potential alternative to the decade-long practice of fine-tuning pre-trained vision models"

You’re right – we intended to say that attention transfer achieves surprisingly good results, and we didn’t mean to imply that attention transfer was already sufficiently developed to be a practical alternative to fine-tuning. We tried to make this clear in L327-330 in the conclusion. We will clarify our wording in the introduction and write “with further research, attention transfer could be a potential alternative to fine-tuning pre-trained vision models.”

The results show that transferring Q is the best of all. So why do we choose to transferring Q and K (or attention)? It implies that transferring the attention pattern may not be the key to the superior performance on downstream task, although transferring attention patterns sounds reasonable and interpretable.

This is a great question! See [a] in the main rebuttal above.

For Table 2-6, it is strange to see only partial results in one table. It is better to show the results for "training from scratch", "fine-tuned", "copy" and "distill".

For some of our tables, we had run either “copy” or “distill” due to compute limitations. During the rebuttal period so far, we have been running more experiments to ensure that Tables 2-6 are comprehensive – see [b] in the main rebuttal above.

We’re really glad that you liked the motivation, analysis, and writing of our paper! We respond to your comments below:

Fig. 5: In attn-distill from layer 20-24, the correlation increases much more than the pretrained model, which is not the case with attn-copy. Do the authors have any intuition on why this is the case?

Our main hypothesis is that this is because we transferred all 24 layers for Attention Copy but only distilled 18 layers for Attention Distillation. This means that Attention Distillation is more flexible in how it combines the features for the last 6 layers, so it can find a strategy similar to the fine-tuned MAE. In contrast, Attention Copy is constrained to be more similar to the pre-trained MAE.

Do the authors have the same observation for attn-copy and attn-distill when using ViT-L pretrained in a supervised setting. Also, im curious to know if the same observation holds for methods pretrained using self-distillation approaches such as DINO

We are running this right now! We hoped to have the results by now but had an issue with our cluster. We will post an update below with the results as soon as possible.

in Figure 7 seems that attn-distill has worse attention than attn-copy at the deeper and intermediate layers. Also the localization of the object seems to be bad as well.

Attention Distillation only really deviates from the pre-trained attention maps at the later layers of the network. This makes sense, since we don’t distill the last 6 layers’ attention maps. Furthermore, the “noisy” attention pattern can sometimes be quite useful in Vision Transformers, since they often tend to store information in low-entropy regions of the background (see “ViTs Need Registers” [6*]). The fine-tuned MAE shows similar “noisy” patterns in its later layers, which we can see after fixing a small issue with our attention map visualizations: since we follow the standard practice of using global average pooling, which averages the representations at all spatial locations, the CLS token representation is not used after the 24th layer. This means that its attention map has no signal to improve. We fix this by now showing the pattern after the 23rd layer instead. The examples in the 1 page PDF show that the fine-tuned MAE also exhibits the same patterns in the later layers as Attention Distillation.

I would also like to see the comparison with different state-of-the-art methods that use ViT-L as their backbone is tasks like classification, object detection and robustness.

To the best of our knowledge, fine-tuned MAE is the SOTA ViT-L model without using extra data on ImageNet classification and the OOD robustness benchmarks. For object detection, the SOTA ViT-B-based model using ImageNet-1k is ViTDet [32] with 56.0 $AP^{box}$ and 48.0 $AP^{mask}$. Note that the detection results cannot be directly compared with ours due to further architectural modifications within ViTDet on top of the ViT backbone.

I would urge the authors to also include the following works for completeness, which have shown that copying attention works well in the domain of NLP

Thank you for suggesting these NLP papers – we will definitely add them in our related works section!

[6*] Darcet, T., Oquab, M., Mairal, J., & Bojanowski, P. (2023). Vision transformers need registers. arXiv preprint arXiv:2309.16588.

Here are the results of fine-tuning, Attention Copy, and Attention Distillation from a DINO ViT-B/16 pre-trained model. Note that our DINO fine-tuning recipe achieves an accuracy of 83.2, higher than the 82.8 that has been achieved in previous papers. Overall, we find similar results as we originally reported in Table 5 with MoCo (another representation learning method based on self-similarity), as well as FLIP (a vision-language contrastive method).