On the Surprising Effectiveness of Attention Transfer for Vision Transformers

Conventional wisdom suggests that pre-training Vision Transformers (ViT) improves downstream performance by learning useful representations. Is this actually true? We investigate this question and find that the features and representations learned during pre-training are not essential. Surprisingly, using only the attention patterns from pre-training (i.e., guiding how information flows between tokens) is sufficient for models to learn high quality features from scratch and achieve comparable downstream performance. We show this by introducing a simple method called attention transfer, where only the attention patterns from a pre-trained teacher ViT are transferred to a student, either by copying or distilling the attention maps. Since attention transfer lets the student learn its own features, ensembling it with a fine-tuned teacher also further improves accuracy on ImageNet. We systematically study various aspects of our findings on the sufficiency of attention maps, including distribution shift settings where they underperform fine-tuning. We hope our exploration provides a better understanding of what pre-training accomplishes and leads to a useful alternative to the standard practice of fine-tuning

Paper

Similar papers

Peer review

Reviewer 9FPV7/10 · confidence 4/52024-07-04

Summary

In this work, the authors demonstrate that a large part of the benefits of pre-training in ViT models actually comes not from the pre-trained features but the pre-trained knowledge of attention maps. Specifically, the authors propose an alternative to fine-tuning call Attention Transfer, and they use this method to transfer attention maps from pre-trained ViTs to teach from-scratch ViTs. This method achieves results comparable to results from finetuning. The authors present two variations of Attention Transfer, the first of which directly copies the attention maps from the teacher network, while the second teaches the student network to make its own attention maps with a distillation objective derived from the teacher attention maps. The authors demonstrate impressive and surprising results on ImageNet-1k, and present further analysis of this approach in a variety of configurations. While the proposed Attention Transfer method does have some clear limitations, specifically for domain gaps, overall, the work has very interesting implications for ViT pretraining and finetuning.

Strengths

The main result of the work, which is to demonstrate the key importance of pre-trained attention maps over pre-trained features in ViT finetuning/transfer learning, is very interesting, and has important implications for the use of ViTs. While I think it is unlikely that Attention Transfer (in its current form) will replace standard finetuning (see notes in the following section), I think the implications of the results and analysis are still very important. The potential to use Attention Transfer with ensembles is also interesting, and the authors show that it has the potential to boost the performance of self-supervised ViTs further. This work also helps to explain some of the properties of MAE. In particular, prior works have found that MAE lends itself well to fine-tuning, but not as well to direct linear probes. This work shows a possible explanation: that the strength of MAE is not in its pretrained features but instead in its pretrained attention maps. The work is clearly presented and has wide coverage of many model and training configurations in the main work and appendix. They also include additional analysis on the impact of partial transfer, and assess if the student is or is not re-learning the same features as the teacher. Overall, their analysis is quite comprehensive.

Weaknesses

While the proposed Attention Transfer method has very interesting implications for ViTs, I’m not sure if it will make its way into practical use, for either the Attention Copy or Attention Distillation variants. The main issue is that the proposed method does not always match the performance of the regular fine-tuning approach, particularly in cases with a domain gap between the pretraining data and the downstream task. In addition, Attention Transfer is more expensive to train than standard Fine Tuning, as acknowledged in Section A.2. For some of the analysis results, it would be very helpful to see how Attention Distillation performs as compared to Attention Copy. In particular, in any analysis where there is a domain gap (Table 2 for example), it would be interesting to see if Attention Distillation performs better, as it is suggested that Attention Distillation allows more flexibility in the student and thus may perform better in such cases. On a less important note, I find that the visualizations in Figures 7, 10, and 11 are somewhat difficult to see due to the combination of the heat map colors and the background images. I would suggest revising these figures to try and make them clearer.

Questions

In the Attention Copy setting, are the unnecessary Q and K layers removed from the student network? There is an important conclusion in the work that seems somewhat under-acknowledged. In lines 210-217 and in Table 1, it is found that transferring the teacher’s Q is more effective than transferring the attention map. Why was more attention and testing not performed on Q transfer?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors discuss the limitations of Attention Transfer and present results for it with many configurations. Overall, they are quite transparent about acknowledging the situations where Attention Transfer underperforms full Fine Tuning. I do not see any risks for negative societal impact.

Reviewer p7aT6/10 · confidence 4/52024-07-12

Summary

The paper introduces attention transfer as an alternative to fine-tuning in Vision Transformers (ViT), separating intra-token and inter-token operations to enhance feature extraction and combination. Attention transfer contains attention copy and attention distillation. Attention distillation matches the performance of fine-tuning on ImageNet-1K while learning different representations, facilitating the performance of feature ensemble. The method is verified across various model scales and datasets.

Strengths

1. The paper first points out the sufficiency of attention maps in pretrainings and provides extensive analyses on attention transfer. 2. The proposed attention distillation is verified across various model scales and datasets.

Weaknesses

1. Compared with full fine-tuning, the proposed attention transfer method introduces additional computation costs on the forward process of the teacher model. Therefore, to make the article more comprehensive, authors should consider including a baseline of distillation on features. 2. Transfer tasks (e.g., Table 2 and 3) should contain the results of training from scratch. 3. There is an important result in L169, the result of attention copy from a fine-tuned model (85.6), not appearing in any Tables or Figures. 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.

Questions

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? If so, is attention really important and sufficient for transfer tasks?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

Reviewer 7AbG6/10 · confidence 4/52024-07-12

Summary

The authors propose a novel perspective on the utility of pretraining vision transformers by demonstrating that the actual features and representations learned during pre-training are not crucial. Instead, they find that simply re-using the self attention from pre-training specifically, the way information flows between tokens—is sufficient for models to learn high-quality features from scratch and achieve performance comparable to pre-trained models. To support their claim, the authors introduce a method called attention copy and attention distill. This method involves transferring the attention from a pre-trained teacher ViT to a student ViT, either by copying or distilling the attention maps. This approach allows the student model to learn its own features while still benefiting from the pre-trained attention patterns. The authors also highlight several drawbacks of previous works that heavily rely on pre-trained features and representations. They point out that the conventional fine-tuning approach may not be as effective under distribution shift settings, where the pre-trained features might not generalize well. In contrast, their attention transfer method provides a more robust alternative that maintains high performance even when the distribution shifts. Through systematic experiments, the authors examine various aspects of their findings, particularly focusing on the sufficiency of attention maps. They provide evidence that attention patterns alone can guide the learning process effectively, thus questioning the necessity of the entire pretraining paradigm.

Strengths

The following are some strengths of this work - The problem is well motivated and also also been discussed in previous works. Not specifically the distillation approach, but reusing attention maps has been interesting not only in the vision community, but widely studied in NLP. - I found that the paper is beautifully written, it basically shows the efforts that the authors took in trying to explain each aspect of their work. They keep the language simple and easy to understand, with concise explanations together with proper visualizations and plots where necessary. - The analysis of different components such as transferring attention from different layers, different heads, CKA analysis etc. is very well thought and it was a joy to read through the findings. So I thank the authors for this and encourage them to do these kind of analysis in all their future works as well. - The experiments on different tasks such as image classification, model robustness etc. show the effectiveness of the approach.

Weaknesses

The following are some queries: - In Figure 5, it is surprising to see that attn-copy has least correlation with respect to the fine-tuned model as compared to attn-distill. 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? - In general, with the CKA computation. I think it would be more interesting to understand the correlation across the features at different layers of the model. I would refer the authors to the work in [51], where the authors show correlation plot of features across every layer before and after their method is applied. This would help understand how attn-copy and attn-distill affect representations learned by the model. The authors can take a look at Figure 2 and Figure 4 in [2*] for reference. - Continuing from above, it would also be interesting to see this correlation across different heads after the model is pretrained with attn-copy and attn-distill. In [1*], the author shows that there exists high correlation across attention heads in ViTs, so it would be interesting to see if attn-copy or attn-distill mitigates this. [1*] https://github.com/sayakpaul/probing-vits?tab=readme-ov-file#visualizing-mean-attention-distances [2*] Zhou et al., Refiner: Refining Self-attention for Vision Transformers, arxiv 2021 - 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, BYOL, iBOT or distillation approaches like SimSiam, etc. - I think I might be missing something, but the visualization of attention from the [CLS] token, 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. Can the authors please comment on this. - 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.

Questions

There have been works such as [51], that illustrate the performance of copying attention across intermediate layers of the network. However, 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 [3*] Xiao et al., Sharing attention weights for fast transformer, IJCAI 2019 [4*] Wang et al., Evolving attention with residual convolutions, ICML 2021 [5*] Ying et al., Lazyformer: Self attention with lazy update, arXiv 2021

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes the authors discuss the limitations

Authorsrebuttal2024-08-14

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). | | tune | copy | distill | |-|-|-|-| DINO | 83.2 | 82.2 | 82.8 |

Reviewer qaog3/10 · confidence 4/52024-07-13

Summary

This paper investigates how transferring attention patterns from a pre-trained ViT to a student affect the student's downstream performance. By applying the attention copy strategy, the paper shows that when the pre-trained dataset and downstream dataset are the same, the trained student may achieve performance superior than the student training from scratch. Moreover, it largely recovers the performance of a pretrained-and-finetuned model. Further, the authors propose an attention transfer (distillation) scheme. With attention distillation scheme, the student may achieve performance comparable or even better than the pretrained-and-finetuned model. The authors provide extensive experiments to verify the effectiveness of attention transfer and show when attention transfer works.

Strengths

- The paper is well-written and easy to follow. - The findings is meaningful and interesting to an extent. - Extensive experiments are conducted.

Weaknesses

- 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. - 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". The empirical results in the paper cannot sufficiently support this claim. For example, we only see comparable results on one setting where the model is pretrained on ImageNet (unsupervised/self-supervised pretraining) and finetuned on ImageNet. For other settings, like out-of-distribution tasks, detection tasks, etc., the attention transfer does not work as well as fine-tuning. - The experiments which study the effect of transferring a subset of Q, K, V seems not to support the main claim of this paper. 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. - 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".

Questions

Please see the weakness part.

Rating

3

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors adequately addressed the limitations.

Reviewer 7AbG2024-08-10

Response to author's rebuttal

I thank the authors for the responses. I think they have answered most of my queries and am satisfied with it. Im looking forward to the results of the DINO experiment. I keep my initial rating. Thanks

Reviewer 9FPV2024-08-12

I thank the authors for their responses and discussion. I agree that the authors should be careful about phrasing to avoid a possible "overclaim" for Attention Transfer vs Transfer Learning. I also support the revised visualizations. Overall I think this work has very interesting implications about transfer learning in ViTs and I maintain my original rating in support of accepting this work.

Reviewer qaog2024-08-12

Re:rebuttal

Thanks for the authors' detailed response. However, I still feel concerned about the novelty and the conclusion of this work. The effectiveness of transferring Q implies that the attention may not be the key to the surprising results of performing attention transfer to an extent. The results are interesting and the story is generally good. But it is still not convincing to me that attention is the underlying key to the performance though it is intuitive. The authors show transferring attention works but transferring Q also works. Is it possible attention is not the key factor? Further evidence should be provided beyond performance numbers, to show the relationship between attention transfer and the performance. Based on the above reasons, I still think the current version is not acceptable.

Authorsrebuttal2024-08-14

> The authors show transferring attention works but transferring Q also works. Is it possible attention is not the key factor? - We emphasize that the *only way that copying Q affects the student model* is by *making it easier to learn a useful attention map* $softmax(QK^\top)$. Thus, the strong performance of copying Q should support our surprising findings on the importance of the pre-trained attention maps. - During the rebuttal period, we experimented with distilling only the Q activations from the teacher to the student. We found that this achieves 85.0 when distilling from a pre-trained MAE ViT-L, which is worse than copying Q. We suspect that this is because the student must first learn to match Q and then learn the K that creates a good attention map. Only then does the Q-distillation attention map provide reasonable guidance to learn good features. - Q distillation does not match the 85.7 that Attention Distillation achieves. This supports our hypothesis that copying Q does well because it allows the model to slightly modify the pre-trained attention maps to better match the downstream task. Indeed, Attention Copy does well when the teacher maps are well-suited for the downstream task: copying from a *fine-tuned* MAE ViT-L achieves the same 85.6 accuracy (L169). - In our original response to Reviewer p7aT, we reported the result of distilling the features from a pre-trained MAE model. This achieved an accuracy of 81.3, significantly lower than what Attention Transfer achieves, which further corroborates our story that features are not as important as previously thought. - We will add these new results and discussion on copying Q to the paper, around L217. - Finally, Figure 7, 10, and 11 visualize the attention maps learned by attention distillation and qualitatively show that it matches the teacher’s attention maps well for the layers that are distilled. This links the pre-trained attention maps to the high downstream task performance. Overall, we believe that our findings on the sufficiency of attention maps are already quite surprising and useful for the research community. We have run extensive experiments and ablations and hope that our paper provides valuable insights that potentially motivate the next generation of pre-training and fine-tuning algorithms.

Reviewer p7aT2024-08-13

Their rebuttal has addressed most of my concerns. I find this paper interesting and believe it will benefit the community. As a result, I have increased my score to weak accept. However, I am still not convinced that the attention map is the primary factor for the high performance, rather than the Q features. More qualitative or quantitative evidence would support the claims presented in the paper.

Authorsrebuttal2024-08-14

We're glad that we have addressed most of your concerns! Below, we clarify the result on copying Q: > I am still not convinced that the attention map is the primary factor for the high performance, rather than the Q features. - We emphasize that the *only way that copying Q affects the student model* is by *making it easier to learn a useful attention map* $softmax(QK^\top)$. Thus, the strong performance of copying Q should support our surprising findings on the importance of the pre-trained attention maps. - During the rebuttal period, we experimented with distilling only the Q activations from the teacher to the student. We found that this achieves 85.0 when distilling from a pre-trained MAE ViT-L, which is worse than copying Q. We suspect that this is because the student must first learn to match Q and then learn the K that creates a good attention map. Only then does the Q-distillation attention map provide reasonable guidance to learn good features. - Q distillation does not match the 85.7 that Attention Distillation achieves. This supports our hypothesis that copying Q does well because it allows the model to slightly modify the pre-trained attention maps to better match the downstream task. Indeed, Attention Copy does well when the teacher maps are well-suited for the downstream task: copying from a *fine-tuned* MAE ViT-L achieves the same 85.6 accuracy (L169). - In our original response to Reviewer p7aT, we reported the result of distilling the features from a pre-trained MAE model. This achieved an accuracy of 81.3, significantly lower than what Attention Transfer achieves, which further corroborates our story that features are not as important as previously thought. - We will add these new results and discussion on copying Q to the paper, around L217. - Finally, Figure 7, 10, and 11 visualize the attention maps learned by attention distillation and qualitatively show that it matches the teacher’s attention maps well for the layers that are distilled. This links the pre-trained attention maps to the high downstream task performance. Overall, we believe that our findings on the sufficiency of attention maps are already quite surprising and useful for the research community. We have run extensive experiments and ablations and hope that our paper provides valuable insights that potentially motivate the next generation of pre-training and fine-tuning algorithms.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC