Intermediate Layer Classifiers for OOD generalization

TL;DR

This work questions the conventional focus on penultimate-layer features for out-of-distribution generalization and demonstrates that intermediate-layer representations, accessed via Intermediate Layer Classifiers (ILCs), often generalize more robustly under distribution shifts. By freezing a pre-trained model and training linear probes on intermediate layers, and by selecting the best layer up to L-2, the approach yields substantial OOD gains in both few-shot and zero-shot settings across diverse datasets and architectures. The findings show that intermediate layers are less sensitive to shifts than the penultimate layer and remain informative even with limited or no OOD data, challenging the prevailing emphasis on last-layer retraining. These results provide a practical, scalable baseline for improving robustness and suggest revisiting the role of intermediate features in real-world OOD scenarios.

Abstract

Deep classifiers are known to be sensitive to data distribution shifts, primarily due to their reliance on spurious correlations in training data. It has been suggested that these classifiers can still find useful features in the network's last layer that hold up under such shifts. In this work, we question the use of last-layer representations for out-of-distribution (OOD) generalisation and explore the utility of intermediate layers. To this end, we introduce \textit{Intermediate Layer Classifiers} (ILCs). We discover that intermediate layer representations frequently offer substantially better generalisation than those from the penultimate layer. In many cases, zero-shot OOD generalisation using earlier-layer representations approaches the few-shot performance of retraining on penultimate layer representations. This is confirmed across multiple datasets, architectures, and types of distribution shifts. Our analysis suggests that intermediate layers are less sensitive to distribution shifts compared to the penultimate layer. These findings highlight the importance of understanding how information is distributed across network layers and its role in OOD generalisation, while also pointing to the limits of penultimate layer representation utility. Code is available at https://github.com/oshapio/intermediate-layer-generalization

Figures (26)

• Figure 1: Usinglastvsintermediatelayers for OOD generalisation. A common way to address distribution shift is to fine-tune the last layer of a network on the target distribution (few-shot learning). We show that earlier-layer representations often generalise better than the last layer. Moreover, even when only the in-distribution (ID) data is available, earlier-layer representations are often better than the last layer (zero-shot learning).
• Figure 2: Intermediate Layer Classifiers (ILC). Given a frozen pre-trained model like a ResNet or a ViT, we train a linear probe on an intermediate layer representation at intermediate layers (here, we show this process only at layer $l$). The composition of $l^\text{th}$ layer feature extractor and the intermediate layer classifier (ILC) is the final classifier. We shorthand $\mathbf{r}_l(\mathbf{x})$ as $\mathbf{r}_l$ for brevity.
• Figure 3: Information content for OOD generalization in last layer vs intermediate layers. "Last layer" refers to the OOD accuracy of the last-layer retraining approach ($\text{ILC}_{L-1}$); "Best layer" refers to the maximal OOD accuracy among the intermediate layer classifiers (ILC) ($\text{ILC}_{l^*}$). For MultiCelebA, we report the worst-group accuracy (WGA).
• Figure 4: Accuracies of ILCs and last layer retraining under varying number of OOD samples for ResNets. Performance of best ILCs and last-layer retraining on subpopulation shifts (first row) and the remaining shifts (second row) using CNN models. We used ResNet50 for Waterbirds and CelebA, and ResNet18 for the remaining datasets.
• Figure 5: WGA on OOD data for ID-trained CNNs. For the explanations for Base, Last layer, and Best layer, refer to the text. ResNet50: Waterbirds, CelebA. ResNet18: MultiCelebA.