Beyond the final layer: Attentive multilayer fusion for vision transformers

TL;DR

This work shows that valuable information for downstream tasks in vision transformers is distributed across layers, not confined to the final representation. It introduces an attention-based multilayer fusion that attends over CLS and AP tokens from all intermediate layers, enabling task-aware integration while keeping the backbone frozen. Across 20 datasets and nine pretrained ViTs, this approach yields consistent gains over standard linear probing, with the largest improvements arising for tasks distant from pretraining domains. Attention heatmaps reveal that intermediate layers contribute differently depending on the task, capturing a balance between hierarchical depth and spatial cues. Overall, the method provides a scalable, efficient pathway for probing-based adaptation of vision transformers by leveraging the model’s hierarchical representations.

Abstract

With the rise of large-scale foundation models, efficiently adapting them to downstream tasks remains a central challenge. Linear probing, which freezes the backbone and trains a lightweight head, is computationally efficient but often restricted to last-layer representations. We show that task-relevant information is distributed across the network hierarchy rather than solely encoded in any of the last layers. To leverage this distribution of information, we apply an attentive probing mechanism that dynamically fuses representations from all layers of a Vision Transformer. This mechanism learns to identify the most relevant layers for a target task and combines low-level structural cues with high-level semantic abstractions. Across 20 diverse datasets and multiple pretrained foundation models, our method achieves consistent, substantial gains over standard linear probes. Attention heatmaps further reveal that tasks different from the pre-training domain benefit most from intermediate representations. Overall, our findings underscore the value of intermediate layer information and demonstrate a principled, task aware approach for unlocking their potential in probing-based adaptation.

Figures (18)

• Figure 1: Schematic of our multi-layer Attentive Probe. The method applies cross-attention to CLS and AP tokens from multiple transformer layers, automatically discovering which representations contain the most task-relevant features.
• Figure 2: Absolute accuracy gain (percentage points) of linear (blue) and attentive probes (orange) when fusing an increasing number of intermediate layer representations ($\mathcal{L}_{\text{last}}$, $\mathcal{L}_{\text{mid+last}}$, $\mathcal{L}_{\text{quarterly}}$, and $\mathcal{L}_{\text{all}}$), as well as AAT (grey) aggregated across datasets for the three base models. Including more intermediate layers improves for all models, with our attentive probe over all layers achieving the highest median gain and consistently outperforms the simple linear probe (zero line).
• Figure 3: Balanced accuracy distributions of baseline (left panel) and absolute accuracy gains in percentage points (right panel) for different representation fusion methods across model architectures, aggregated over all 20 datasets. The substantial benefits from attentive probing of intermediate layers [All layers (CLS+AP, attentive)] persist even for large models, indicating that large models fail to encode all task-relevant information in their final layer's CLS token.
• Figure 4: Attention weights across layers and datasets for base models, averaged over heads and samples, are distributed across multiple layers, demonstrating their relevance for downstream tasks.
• Figure 5: Accuracies per model and dataset