Head2Toe: Utilizing Intermediate Representations for Better OOD Generalization

1 INTRODUCTION . Tranfer learning is a widely used method for obtaining strong performance in a variety of tasks where training data is scarce—see Zhu et al . ( 2020 ) ; Alyafeai et al . ( 2020 ) ; Zhuang et al . ( 2020 ) for recent application-specific surveys . A well-known recipe for transfer learning involves the supervised or unsupervised pretraining of a model on a source task with a large training dataset ( also referred to as upstream training ) . After pretraining , the model ’ s output head is discarded , and the rest of the network is used to obtain a feature embedding , i.e. , the output of what was formerly the penultimate layer of the network . When transferring to a target task , a new output head is trained on top of the feature extractor ( downstream training ) . This approach makes intuitive sense : if a linear combination of embedding features performs well on the source task , we expect a different linear combination of features to generalize to the target domain , provided the source and target tasks are similar . This approach of training a new output head , referred to as linear probing ( LINEAR ) , often yields significant improvements in performance on the target task over training the network from scratch ( Kornblith et al. , 2019 ) . An alternative to LINEAR is fine-tuning ( FINETUNING ) , which uses targetdomain data to adapt all weights in the feature extractor together with the new output head . This procedure requires doing forward and backward passes through the entire network at each training step and therefore its per-step cost is significantly higher than LINEAR . Furthermore , since the entire network is fine-tuned , the entire set of new weights needs to be stored for every target task , making FINETUNING impractical when working on edge devices or with a large number of target tasks . However , FINETUNING is often preferred over LINEAR since it consistently leads to better performance on a variety of target tasks even when data is scarce ( Zhai et al. , 2019 ) . FINETUNING ’ s superior generalization in the low-data regime is counterintuitive given that the number of model parameters to be adapted is often large relative to the amount of available training data . How does FINETUNING learn from few examples successfully ? We conjecture that FINETUNING better leverages existing internal representations rather than discovering entirely new representations ; FINETUNING exposes existing features buried deep in the net for use by the classifier . Under this hypothesis , features needed for transfer are already present in the pretrained network and might be identified directly without fine-tuning the backbone itself . In this work , we propose and explore methods for selecting useful features embedded in internal layers of a pretrained net , concatenating them to the embedding produced by the pretrained net , and then applying the LINEAR transfer approach to the augmented representation ( Fig . 1-left ) . Our approach leads to significant improvements over LINEAR as shown in Fig . 1-right . Our key contributions are as follows : 1 . We perform analyses to better understand two factors that affect the benefit of incorporating intermediate representations : feature selection and the degree to which a target domain is out-of-distribution ( OOD ) . 2 . We introduce head-to-toe probing ( HEAD2TOE ) , a simple recipe for selecting relevant features from intermediate representations for few-shot generalization . On the VTAB collection of data sets , we show that HEAD2TOE outperforms LINEAR and matches the performance of more computationally costly FINETUNING . 3 . Critically , HEAD2TOE outperforms FINETUNING on OOD target domains . If a practitioner can make an educated guess about whether a target domain is OOD with respect to a source , using HEAD2TOE improves on the state of the art for transfer learning 2 PRELIMINARIES . Source domain and backbone models . In our experiments , we use source models pretrained on ImageNet-2012 ( Russakovsky et al. , 2015 ) , a large scale image classification benchmark with 1000 classes and over million natural images . We benchmark HEAD2TOE using convolutional ( ResNet-50 , Wu et al. , 2018 ) and attention-based ( ViT-B/16 , Dosovitskiy et al. , 2021 ) architectures pretrained on ImageNet-2012 . Target domains . In this work , we focus on target tasks with few examples ( i.e. , few-shot ) and use Visual Task Adaptation Benchmark-1k ( Zhai et al. , 2019 ) to evaluate different methods . Visual Task Adaptation Benchmark-1k consists of 19 different classification tasks , each having between 2 to 397 classes and a total of 1000 training instances . The domains are grouped into three primary categories : ( 1 ) natural images ( natural ) , ( 2 ) specialized images using non-standard cameras ( specialized ) , and ( 3 ) rendered artificial images ( structured ) . Characterizing out-of-distribution domains . Consider the relationship between source and target domains . If the domain overlap is high , then features extracted for linear classification in the source domain should also be relevant for linear classification in the target domain , and LINEAR should yield performance benefits . If the domain overlap is low , then constraining the target domain to use the source domain embedding may be harmful relative to training a model from scratch on the target domain . Therefore , we might quantify the source-target distribution shift in terms of how beneficial LINEAR is relative to training a model from scratch ( denoted SCRATCH ) : DomainAffinity = AccLINEAR − AccSCRATCH . In Fig . 2 , we sort the 19 VTAB-1k target tasks by their domain affinity to ImageNet-2012—from low to high—for a pretrained ResNet-50 backbone . The Figure shows examples of the three target domains with the most and least distribution shift ( left and right sides , respectively ) . These examples seem consistent with intuitive notions of distribution shift . Baselines . Fig . 2 also presents transfer test accuracy of LINEAR , FINETUNING , and SCRATCH baselines . Consistent with the literature , FINETUNING performs as well or better than the other two baselines . For in-distribution targets ( right side of graph ) , LINEAR matches FINETUNING ; for OOD targets ( left side of graph ) , LINEAR is worse than FINETUNING . With distribution mismatch , the source network may filter information available in lower layers because it is not needed for the source task but is essential for the target task . Observing FINETUNING performs better than SCRATCH even in OOD tasks , we hypothesize that intermediate features are key for the FINETUNING , since if learning novel features was possible using limited data , training the network from scratch would work on-par with FINETUNING . Motivated by this observation , HEAD2TOE probes the intermediate features of a network directly and aims to eliminate the need for fine-tuning . 3 HEAD2TOE UTILIZATION OF PRETRAINED BACKBONES . 3.1 YOUR REPRESENTATIONS ARE RICHER THAN YOU THINK . In this section , we conduct a simple experiment to demonstrate the potential of using representations from intermediate layers . We concatenate the feature embedding of a pretrained ResNet-50 backbone ( features from the penultimate layer ) with features from one additional layer and train a linear head on top of the concatenated features . When extracting features from convolutional layers , we reduce the dimensionality of the convolutional stack of feature maps using strided average pooling , with a window size chosen so that the resulting number of features is similar to the number of embedding features ( 2048 for ResNet-50 ) . In order to estimate an upper bound on how much the performance can be improved over LINEAR by including a single intermediate layer , we use an oracle to select the layer which yields the largest boost in test performance ( for each target task separately ) . Percentage improvement over LINEAR using this ORACLE is shown in Fig . 3-left . We observe a negative correlation ( -0.745 , Spearman ) between the domain affinity of a target task and the accuracy gain . As predicted earlier , adding intermediate representations does not improve in-domain generalization because feature embedding already contains the most useful features . In contrast , generalization on out-of-domain tasks are improved significantly when intermediate features are used . In Fig . 3-top-right , we show test accuracy as a function of the ResNet-50 layer whose internal representation is used to augment the feature embeddings to boost LINEAR transfer . Figures for remaining tasks can be found in Appendix D. Different tasks on VTAB benchmark prefer inclusion of different layers to obtain the best generalization accuracy , emphasizing the importance of selecting the right set of features for optimal generalization . Overall , the ORACLE that selects the layer with best test performance for each task yields an average of 3.5 % improvement on VTAB-1k benchmark . One possible explanation for the improvement in performance with the augmented representation is simply that it has more degrees of freedom ( 4096 features instead of 2048 ) . To demonstrate that the improvement is due to inclusion of intermediate layers and not simply due to increased dimensionality , we formed a 4096-element feature representation by using a second feature embedding obtained from a second ResNet-50 backbone pretrained on ImageNet-2012 ( CONTROL ) . Note that this experiment bears similarity to ensembling ( Zhou et al. , 2002 ) , which is known to bring significant accuracy gains on its own ( Mustafa et al. , 2020 ) . Utilizing a second backbone doubles the amount of resources required , yet it falls 1 % shy of the Oracle performance demonstrating the extent to which using intermediate representations can be helpful for generalization . 3.2 HEAD2TOE . Motivated by our observations in the previous section , we hypothesize that we can attain—or possibly surpass—the performance of FINETUNING without changing the backbone itself by using LINEAR augmented with well-chosen intermediate activations . However , this approach introduces the problem of selecting features . The number of possible features to include is large relative to the amount of training data available , leading to the potential of overfitting . In this section , we explore various regularization and feature selection methods , which we then incorporate into our method , HEAD2TOE . Notation . Our method applies to any network with any type of layers , but here , for simplicity and without loss of generality , we consider a network with L fully connected layers , each layer receiving input from the layer below : z ` = h ` −1W ` ; h ` = f ( z ` ) ( 1 ) where the subscript denotes a layer index , h0 = x is the input , f is the activation function , W ` is the weight matrix of layer ` , and zL is the logit vector used for classification . When transferring a pretrained network to a target task using LINEAR , we discard the last layer of the pretrained network and train a new set of linear weights , W ′L , such that predictions ( logits ) for the new task are obtained by z′L = hL−1W ′ L. Head2Toe . Consider a simple scheme that augments the backbone embedding with activations from all layers of the network , such that : z′L = hallWall ; hall = concat ( a1 ( h1 ) , a2 ( h2 ) , ... , aL ( hL ) ) ( 2 ) where a ` ( . ) denotes a fixed function to reduce the dimensionality of the activation vector at a given layer ` . Such functions become valuable when considering network architectures that generate large number of intermediate features . For example , for convolutional networks , where the representations at each layer are structured as multi-channel spatial arrays of detectors , we perform two-dimensional average pooling as explained in Section 3.1 and Section 5 . Even with dimensionality reduction , hall can exceed a million elements , and Wall is underconstrained by the training data , leading to overfitting . Further , without care Wall may become so large as to be impractical for deploying this model.1 . We can address these issues by selecting a subset of features before training the targetdomain classifier . Equivalently , we can zero out the non-selected rows of Wall . We will show that appropriate selection of a subset of features leads to better generalization than using all features . Feature selection based on group lasso . Group lasso ( Yuan & Lin , 2006 ) is a popular method for selecting relevant features in multi-task adaptation settings ( Argyriou et al. , 2007 ; Nie et al. , 2010 ) . When used as a regularizer on a weight matrix W , the group-lasso regularizer encourages the ` 2 norm of the rows of the matrix to be sparse , and is defined as : |W |2,1 = |s|1 = ∑ i |si| ; si = √∑ j w2ij ( 3 ) To determine which features are most useful for the task , the linear head is trained with group-lasso regularization on Wall . In contrast to the approach taken by Argyriou et al . ( 2007 ) and Nie et al . ( 2010 ) , which use costly matrix inversions , we incorporate the regularization as a secondary loss and train with stochastic gradient descent . Following training , the ` 2 norms of the rows of Wall , s , are used to decide which features are most relevant . We refer to si as the relevance score for feature i . We pick a fraction F of all features with the greatest relevance and train a new linear head to obtain the final logit mapping . Feature selection alone provides strong regularization , therefore during the final training we don ’ t use any additional regularization . We make two comments on this . First , because the initial round of training Wall with the group-lasso regularizer is used only to rank features by importance , the method is quite robust to the regularization coefficient ; it simply needs to be large enough to distinguish the contribution of the individual features . Second , interpreting si as the importance of feature i depends on all features having the same dynamic range . This constraint will often be satisfied naturally due to the normalization layers found in the architecture itself ( we read out representations after normalization ) . To cover the remaining cases ( where no normalization exists between two layers ) we normalize the representations of each layer to a unit norm . Selection of F and the cost of HEAD2TOE . The fraction F determines the total number of features retained . One would expect the optimal value to depend on the target task . We select F based on validation performance . This validation procedure is inexpensive compared to the cost of the initial phase of the algorithm ( i.e. , training of Wall to obtain s ) due to the reduced number of features in the second step . Overall , HEAD2TOE together with its validation procedure requires 18 % more operations compared to training Wall alone ( details shared in Appendix B ) . Verifying HEAD2TOE . In Fig . 4-left , we demonstrate the effectiveness of group lasso on identifying relevant intermediate features of a ResNet-50 for transfer from ImageNet . We rank all features 1For example , using a pooling size of 2 , ResNet-50 generates 1.7 million features and storing Wall requires 6.6e8 parameters ( 2.6GB for float32 ) for SUN-397 . 0 5000 10000 15000 20000 25000 Offset 0.54 0.56 0.58 0.60 Te st A cc ur ac y Linear Oracle 102 103 104 # Features Selected 0.46 0.48 0.50 0.52 0.54 0.56 0.58 0.60 Te st A cc ur ac y Linear Oracle feature-wise layer-wise layer-wise ( group ) 0 100 200 300 400 Feature ( Acc : 63.8 % ) .ro ot b1 u1 b1 u2 b1 u3 b2 u1 b2 u2 b2 u3 b2 u4 b3 u1 b3 u2 b3 u3 b3 u4 b3 u5 b3 u6 b4 u1 b4 u2 b4 u3 pr el og its 0 pr el og its lo gi ts 0 500 1000 1500 2000 Layer ( Acc : 44.1 % ) Figure 4 : ( left ) Over all VTAB tasks , average accuracy of HEAD2TOE when selecting 2048 consecutive features sorted by their relevance score , starting with an index specified by offset . ( center ) Average accuracy over all VTAB tasks as a function of the number of features included . In this experiment , we only use features from the first layer of each of the 18 ResNet blocks and adjust pooling to have around 2048 features for each layer , totalling 29800 features . Results show that selecting layers performs worse than selecting features when adapting to a target domain . ( right ) Distribution of 2048 intermediate features retained from a ResNet-50 when using feature-wise and layer-wise scores on the SVHN transfer task . by their relevance score , si , and we select groups of 2048 consecutive features beginning at a particular offset in this ranking . Offset 0 corresponds to selecting the features with largest relevance . We calculate average test accuracies across all VTAB tasks . As the figure shows , test accuracy decreases monotonically with the offset , indicating that the relevance score predicts the importance of including a feature in the linear classifier . HEAD2TOE selects individual features independent of the layer in which they are embedded . We compare this feature-wise strategy to selecting layers as whole ( i.e. , selecting all features in a layer or none ) . One might expect layer-wise selection to be superior because it involves fewer decisions and therefore less opportunity to overfit to the validation set that is being used for selection . Further , layer-wise selection may be superior because the relevance of one feature in a layer may indicate the relevance of others . To obtain a layer-wise relevance score , we compute the mean relevance score of all features in a layer and then rank layers by relevance . We also run an alternative layer selection algorithm , in which we group weights originating from the same layer together and select layers using the ` 2 norm of the groups ( layer-wise ( group ) ) . Fig . 4-center compares feature-wise and layer-wise selection methods , matched for number of features retained . Feature-wise selection performs better than layer-wise selection on average and the best performance is obtained when around 1000 features are kept . We share figures for all 19 datasets in Appendix E. Fig . 4-right shows the distribution of features selected over different layers for both strategies . Note that HEAD2TOE ’ s use of a fixed backbone means that as we search for features to include , the extracted features themselves are not changing . Consequently , we can calculate them once and reuse as necessary , instead of recalculating at every step , as required for FINETUNING . Furthermore , since the backbone is frozen , the only additional weight storage required for each target task is the output head . Such properties are key for transfer-learning methods to be practical when using models with billions of parameters .

Summary: This paper explores the utility of intermediate layers for linear probing in transfer learning. Specifically, authors proposed Head-to-Toe probing (HEAD2TOE), that selects features from all layers of the source model to train a classification head for the target-domain. Experiments on the VTAB benchmark shows that Head2Toe matches performance obtained with fine-tuning on average, but critically, for out-of-distribution transfer, Head2Toe outperforms fine-tuning.