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 propose Head2Toe, a method that exploits intermediate representations of DNNs to improve performance, and OOD generalization of transfer learning. The key idea of the Head2Toe is to augment traditional linear probing with intermediate representations. It uses group lasso as regularization to learn weights of different features, and then select features based on the learned weights and validation performance.Experiments on VTAB benchmark show that Head2Toe outperforms linear probes and is competitive compared with fine-tuning. An important finding is that Head2Toe improves OOD generalization.