Layer by layer, module by module: Choose both for optimal OOD probing of ViT

AI summary

Overview

Research area: Computer vision, specifically representation analysis of pretrained vision transformers (ViTs) under distribution shift.

Technical level: Intermediate. The paper is readable without deep mathematical background, but it assumes familiarity with transformer block internals (LayerNorm, multi-head attention, feedforward networks, residual connections) and with linear probing as an evaluation protocol.

Scope: A layer-by-layer and module-by-module linear probing study of an 86M-parameter ViT pretrained on ImageNet-21k, evaluated across 11 image classification benchmarks spanning in-distribution and out-of-distribution conditions.

What This Paper Is About

Prior work argued that intermediate layers of foundation models beat the final layer mainly because of autoregressive pretraining, and that for vision transformers the final layer stays optimal. This paper shows that conclusion holds only when the downstream data is in-distribution (ID). Once the downstream data differs from the pretraining data, the deeper layers degrade and intermediate layers become the better choice. The authors then go one level finer, asking not just which layer to probe but which module inside a transformer block to probe.

Key Contributions

1. Diagnosis of the cause. The paper identifies distribution shift between pretraining and downstream data as the primary reason intermediate layers of a pretrained ViT outperform final layers — not the pretraining objective itself, as earlier work on autoregressive vision models had suggested.

2. A fine-grained module-level probing study. Instead of probing only the transformer block output, the authors track and probe the outputs of 8 operations inside each block — LN1, MHA, RC1, LN2, FC1, Act, FC2, and RC2 — across layers and across datasets.

3. A concrete finding that standard probing is suboptimal. Probing RC2 (the transformer block output, the standard choice in prior work) is not the best option on all datasets but one. FC2 is the worst module to probe, with the lowest accuracy on 10 out of 12 datasets (as stated in the paper), while Act achieves the highest overall win rate.

4. Practical guidance for choosing where to probe. Under significant distribution shift, probe the feedforward activation (Act) in intermediate layers; when the shift is weak or negligible, probe a normalized module instead — the paper's abstract points to the normalized output of the multi-head self-attention module, while the discussion section names the LayerNorm preceding the feedforward network (LN2). A conservative fallback when the shift is hard to detect is LN2 rather than the standard RC2.

Main Findings

• Deeper is not always better under shift. With the pretrained model frozen, probing performance at deeper layers degrades as the distribution shift grows. On in-distribution data, the best visual embedding is at the end of the network. The paper uses the gap between frozen and finetuned encoders as a proxy for the strength of the shift, and orders datasets from weakest to strongest shift as Flowers102, Cifar10, Contrast, and Speckle Noise.

• Intermediate layers are more robust than final layers. The paper explains this by the tendency of layers to specialize closer to the classification head.

• Act wins under strong shift. The activation after FC1 is the best-performing module overall and beats other components by a large margin when the shift is strong, even though it is less good on easier datasets such as Cifar10, Flowers102, or Pets. Examples from the best-over-layers table: Act reaches 80.20 on Contrast, 71.15 on Motion Blur, 63.35 on Speckle Noise, 49.34 on Clipart, and 34.90 on Sketch, each higher than every other module on that dataset.

• FC2 is consistently the worst choice. FC2 shows the lowest accuracy on 10 out of 12 datasets, and its profile degrades most sharply with depth. Sample values: 89.98 on Cifar10, 60.92 on Cifar100, 70.85 on Contrast, 56.70 on Gaussian Noise, 40.97 on Clipart, 28.45 on Sketch.

• Other modules are comparable, with LN2 slightly ahead. On Cifar10, RC1 reaches 92.19, LN2 92.20, FC1 92.28, and RC2 92.07; on Flowers102 most components perform equally well (LN1 96.58, RC1 96.58, LN2 96.62, FC1 96.44, RC2 96.62), the single dataset where standard RC2 probing is not clearly suboptimal.

• Residual stream explains stability. Because the residual stream carries information from previous layers, the accuracy profiles of LN2 and RC2 are less concave across depth than those of Act and FC2.

• A dimensional argument. All modules map from (R^d)^n to (R^d)^n except the feedforward network, where FC1 expands tokens to 4d and FC2 contracts them back to d. The authors hypothesize that operating in higher dimension helps feature disentanglement, and that Act filters noise induced by the projection, while FC2's compression may hurt linear separability. A complementary reading treats the feedforward as a key-value memory, where FC1 and Act capture semantic content of inputs while FC2 reflects a distribution over tokens.

• Finetuning upper bounds. Full finetuning reaches 99.02 ± 0.02 on Cifar10, 92.74 ± 0.05 on Cifar100, 97.23 ± 0.18 on Contrast, 87.14 ± 1.16 on Gaussian Noise, 94.67 ± 0.14 on Motion Blur, 95.42 ± 0.13 on Snow, 89.58 ± 0.43 on Speckle Noise, 78.50 ± 0.49 on Clipart, 71.30 ± 0.26 on Sketch, 99.15 ± 0.05 on Flowers102, and 94.57 ± 0.29 on Pets.

Experimental setup, for context: all experiments use an 86M-parameter ViT pretrained on ImageNet-21k. Linear probing pools the CLS token embedding and applies logistic regression with the L-BFGS solver. The 11 benchmarks are Cifar10, Cifar100; five Cifar10-C variants at severity 5 (Contrast, Gaussian Noise, Motion Blur, Snow, Speckle Noise); two DomainNet domains (Clipart, Sketch); Flowers102; and Pets. Finetuning uses SGD with momentum 0.9, no weight decay, cosine learning rate decay, batch size 512, gradient clipping at norm 1, resolution 224×224, a sweep over 4 learning rates, and 3 runs with different seeds; each dataset has 4000 to 20000 training steps depending on the configuration.

Methodology in Plain English

The authors take a single frozen pretrained ViT and, for every layer, extract the representation produced by each internal operation. They then fit a simple logistic regression on top of those frozen features — no updating of the backbone — which tells them how linearly separable the classes are in each representation. Doing this across the depth of the network produces an accuracy-versus-depth curve for each module.

To know how much the downstream data differs from the pretraining data, they finetune the same model on each dataset and use the gap between the frozen and finetuned accuracy as a proxy for shift severity. Datasets are then ordered from weakest to strongest shift, and the shape of each module's accuracy curve is compared across that ordering. The comparison is repeated on every benchmark, with residual connections and LayerNorms included so that the contribution of each part of the block can be separated from the others.

Why This Matters

The paper reframes a debate about whether intermediate layers are "better" — a framing that depended on evaluating mostly on data the model had already seen. By showing that the answer flips with distribution shift, it gives practitioners a concrete rule for choosing a probing location, and it suggests that knowing whether you are in an ID or OOD regime is a prerequisite for that choice.

Impact on research: the work challenges the conclusion of Skean et al. (2025) that the benefit of intermediate layers is a byproduct of autoregressive pretraining, and it argues that the standard practice of probing the transformer block output is a suboptimal default. It also opens a module-level framing that prior layer-level analyses could not express.

Real-world applications:

• Medical imaging, where models pretrained on natural images are applied to clinical scans that differ markedly from the pretraining distribution.
• Autonomous driving and remote sensing, where weather, lighting, and sensor differences create exactly the kind of corruption-style shift represented by the Cifar10-C variants used here.
• Industrial visual inspection, where defect images from a new production line or camera setup will not match pretraining data.
• Fine-grained recognition in the wild, such as camera-trap wildlife or plant identification, which resembles the DomainNet Clipart and Sketch domain gaps.

Industry relevance: when finetuning is too expensive or labels are scarce, teams rely on frozen features. The paper's practical recommendation — probe the feedforward activation if the shift is significant, otherwise probe a normalized module and use LN2 as a safe default — is directly actionable for anyone building a linear probe, retrieval index, or lightweight classifier on top of a frozen ViT.

Future Directions

1. Extend the analysis of Skean et al. (2025) with information-theoretic, geometric, and invariance measures computed at the level of individual transformer modules, rather than at the level of whole layers.

2. Build methods that detect the degree of distribution shift at deployment, since the optimal probing target depends on shift severity and the paper's own proxy requires finetuning the model.

3. Determine automatically which layer and which module to probe for a given unlabeled downstream dataset, removing the need for a sweep over all layer-module pairs.

4. Test whether the layer-level and module-level conclusions transfer to other backbones, pretraining corpora, and modalities, and to objectives beyond the ImageNet-21k supervised pretraining used here.

Target Audience

Researchers and engineers working on transfer learning, robustness to distribution shift, or representation analysis of vision transformers. It is also useful for practitioners who deploy frozen foundation models as feature extractors and need to decide where to tap the network, and for students who want a concrete, well-scoped example of how probing studies are designed and interpreted. Readers looking for new architectures, training objectives, or state-of-the-art accuracy numbers will not find them here — the contribution is diagnostic and prescriptive rather than a new model.

Authors’ abstract

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.