Source-linked AI summary
Layer by layer, module by module: Choose both for optimal OOD probing of ViT
Ambroise Odonnat, Vasilii Feofanov, Laetitia Chapel, Romain Tavenard, Ievgen Redko
TL;DR
The paper asks why intermediate ViT layers can outperform final layers and examines this question with linear probing across layers, modules, and diverse classification benchmarks. It finds that distribution shift drives deeper-layer degradation and that optimal probing locations depend on shift strength, with feedforward activations favored under strong shift and normalized attention outputs under weak shift.
Problem
The paper addresses limited and conflicting evidence about when intermediate vision-transformer representations outperform final-layer representations under distribution shift.
Method
The authors linearly probe layers and transformer modules of a pretrained 86M-parameter ViT across 11 diverse image-classification benchmarks.
Results
Distribution shift drives final-layer degradation; feedforward activations perform best under strong shift, while LayerNorm outputs before feedforward networks perform better when shift is negligible.
Takeaways & Limitations
Standard transformer-block-output probing is suboptimal, so layer and module selection should depend on whether downstream data is in-distribution or out-of-distribution.
Takeaways & Limitations
The proposed explanation that higher-dimensional feedforward representations promote feature disentanglement remains a hypothesis.
Abstract
from arXiv · showhide
Recent studies have observed that intermediate layers of foundation models often yield more discriminative representations than the final layer. While initially attributed to autoregressive pretraining, this phenomenon has also been identified in models trained via supervised and discriminative self-supervised objectives. In this paper, we conduct a comprehensive study to analyze the behavior of intermediate layers in pretrained vision transformers. Through extensive linear probing experiments across a diverse set of image classification benchmarks, we find that distribution shift between pretraining and downstream data is the primary cause of performance degradation in deeper layers. Furthermore, we perform a fine-grained analysis at the module level. Our findings reveal that standard probing of transformer block outputs is suboptimal; instead, probing the activation within the feedforward network yields the best performance under significant distribution shift, whereas the normalized output of the multi-head self-attention module is optimal when the shift is weak.
1 INTRODUCTION
The paper investigates when intermediate ViT representations outperform final-layer features under distribution shift. It studies this question through linear probing across layers and transformer modules, identifying shift-dependent probing choices.
- Motivation: Distribution shift between pretraining and downstream data is identified as the driving factor behind intermediate layers outperforming final layers.The study attributes the phenomenon to downstream distribution differences rather than only the pretraining objective.
- Approach: The paper studies pretrained ViTs on out-of-distribution image-classification tasks using linear probing across layers.It also examines when and why intermediate representations outperform final-layer representations.
- Approach: The analysis probes normalization, multi-head attention, residual, and feedforward modules because their robustness to distribution shift may differ.The paper conducts a fine-grained module-level study in addition to layer-wise probing.
- Takeaways: In in-distribution settings, final layers yield better performance than intermediate layers.This is presented as the first actionable takeaway.
- Takeaways: In out-of-distribution settings, probing intermediate feedforward inputs and activations performs better.This is presented as the second actionable takeaway.
2 EXPERIMENTAL SETUP
The experiments use an 86M-parameter ViT pretrained on ImageNet-21k and evaluate pooled hidden representations with logistic-regression linear probes across 11 classification benchmarks.
- Transformer modules: Each transformer block alternates multi-head attention and feedforward networks, preceded by LayerNorms and followed by residual connections.The study tracks outputs from eight operations within each layer.
- Model and evaluation: The model is an 86M-parameter vision transformer pretrained on ImageNet-21k.The setup follows the paper’s stated pretrained ViT configuration.
- Model and evaluation: Linear probing pools CLS-token embeddings and applies logistic regression with the L-BFGS solver.This procedure is used for each selected hidden representation.
- Benchmarks: The benchmark suite contains 11 image-classification datasets spanning Cifar10, Cifar100, five Cifar10-C variants, two DomainNet domains, Flowers102, and Pets.The datasets provide varied downstream classification settings.
3 DISTRIBUTION SHIFT DEGRADES THE PERFORMANCE OF FINAL LAYERS
The layer-wise experiments show that distribution shift degrades deeper representations, while final layers remain best in-distribution. This pattern supports choosing layers according to the relationship between pretraining and downstream data.
- Layer-wise results: Intermediate layers become more robust than final layers as the distribution shift between pretraining and downstream data increases.The paper observes worsening deeper representations in OOD settings as shift increases.
- Layer-wise results: In out-of-distribution settings, deeper representations become worse as the shift increases across downstream datasets.The reported plots are ordered from lower to higher shift, with similar patterns observed across all datasets.
- Layer-wise results: In in-distribution settings, the best visual embedding occurs at the end of the network.The final layers outperform intermediate layers in the ID scenario.
- Interpretation: Intermediate representations may be more robust because layers specialize closer to the classification head, making layer selection important when finetuning is prohibitive.The paper frames identifying ID versus OOD conditions as necessary for anticipating accuracy profiles and selecting a layer.
4 NOT ALL TRANSFORMER MODULES ARE WORTH PROBING ON OOD DATA
Module choice materially affects linear probing under distribution shift: standard transformer-block outputs are usually suboptimal, while feedforward activations perform best when the shift is strong.
- Probing transformer-block outputs (RC2) is suboptimal on all datasets except Flowers102, where most components perform equally well.
- FC2 is the worst module, achieving the lowest accuracy on 10 of 12 datasets.
- Act has the highest overall win rate and outperforms other components by a large margin under strong distribution shift.
- As distribution shift increases, Act is strongest in intermediate layers while its performance plunges in final layers; LN2 and RC2 remain less accurate but more stable.
- Residual-stream information may explain why LN2 and RC2 have less concave accuracy profiles than Act and FC2.
5 DISCUSSION
Across diverse classification benchmarks, the paper finds that distribution shift explains final-layer degradation and that the best module to probe depends on shift strength.
- Final-layer performance degrades under distribution shift, whereas intermediate representations are more robust.
- Standard probing of transformer-block outputs is not optimal across the studied settings.
- Feedforward activations in intermediate layers are richest under significant shift, while LayerNorm outputs preceding the feedforward network are better when shift is almost negligible.
A IMPLEMENTATION DETAILS
The implementation uses a patch-based ViT and evaluates representations with linear probing alongside finetuning experiments across varied classification benchmarks.
- ViT inputs are split into square patches, linearly embedded, prepended with a CLS token, and processed through transformer layers.
- The experiments use varied benchmarks including Cifar10, Cifar100, Cifar10-C variants, DomainNet domains, Flowers102, and Pets.
- Finetuning uses SGD with momentum 0.9, no weight decay, cosine learning-rate decay, batch size 512, and gradient clipping at norm 1.
- Full finetuning reports best top-1 test accuracy over each dataset’s learning-rate grid, with means and standard deviations from three runs.
B ADDITIONAL EXPERIMENTS
Additional figures extend the layer-wise analyses across benchmarks and compare pretrained with finetuned models as distribution shift increases.
- Figures 5 and 6 provide additional results for the experiments in Sections 3 and 4 across all benchmarks.
- Figure 5 compares pretrained and finetuned models across layers, ordering datasets by increasing distribution shift from left to right.
- Figure 5 states that stronger distribution shift is associated with worse final-layer performance.
- Figure 6 presents layer-wise, module-wise linear-probing performance as the shift between pretraining and downstream data increases.