StochCA: A Novel Approach for Exploiting Pretrained Models with Cross-Attention

TL;DR

StochCA addresses the challenge of efficiently exploiting large pretrained models for target tasks by integrating cross-attention with a frozen pretrained ViT. It introduces a stochastic per-layer switch between self-attention and cross-attention, enabling the target model to selectively reference pretrained knowledge while fine-tuning only the queries and MLPs, formalized as $h_l=(1-\beta)SA(...) + \beta CA(...)$ with $\beta\sim\mathrm{Bernoulli}(p)$. Empirical results across transfer learning and domain generalization show StochCA often outperforms strong baselines and remains complementary when combined with methods like MIRO, with analysis revealing a favorable balance between leveraging pretrained representations and avoiding negative transfer. The method scales across ImageNet and CLIP backbones and demonstrates robustness and data efficiency, offering practical impact for deploying pretrained models on diverse visual tasks.

Abstract

Utilizing large-scale pretrained models is a well-known strategy to enhance performance on various target tasks. It is typically achieved through fine-tuning pretrained models on target tasks. However, naïve fine-tuning may not fully leverage knowledge embedded in pretrained models. In this study, we introduce a novel fine-tuning method, called stochastic cross-attention (StochCA), specific to Transformer architectures. This method modifies the Transformer's self-attention mechanism to selectively utilize knowledge from pretrained models during fine-tuning. Specifically, in each block, instead of self-attention, cross-attention is performed stochastically according to the predefined probability, where keys and values are extracted from the corresponding block of a pretrained model. By doing so, queries and channel-mixing multi-layer perceptron layers of a target model are fine-tuned to target tasks to learn how to effectively exploit rich representations of pretrained models. To verify the effectiveness of StochCA, extensive experiments are conducted on benchmarks in the areas of transfer learning and domain generalization, where the exploitation of pretrained models is critical. Our experimental results show the superiority of StochCA over state-of-the-art approaches in both areas. Furthermore, we demonstrate that StochCA is complementary to existing approaches, i.e., it can be combined with them to further improve performance. Our code is available at https://github.com/daintlab/stochastic_cross_attention

Figures (4)

• Figure 1: Illustration of the attention mechanisms: (Left) The process of self-attention, where a model focuses on different parts of a single input sequence to compute its representations. (Right) The process of cross-attention, where a model attends to a different sequence (e.g., from a pretrained model) while processing an input sequence, facilitating an interaction between the two.
• Figure 2: Two approaches to combining self-attention and cross-attention: (Left) A simple framework that integrates both self-attention and cross-attention in a straightforward manner. (Right) The StochCA Framework that incorporates stochasticity in the combination of self-attention and cross-attention.
• Figure 3: Samples of domain generalization benchmarks.
• Figure : Training procedure of Stochastic Cross-Attention