H EAD 2T OE : U TILIZING I NTERMEDIATE R EPRESENTATIONS FOR B ETTER OOD G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Transfer-learning methods aim to improve performance in a data-scarce target domain using a model pretrained on a data-rich source domain.
A cost-efﬁcient strategy, linear probing , involves freezing the source model and training a new classiﬁcation head for the target domain.
This strategy is outperformed by a more costly but state-of-the-art method— ﬁne-tuning all parameters of the source model to the target domain—possibly because ﬁne-tuning allows the model to leverage useful information from intermediate layers which is otherwise discarded by the later pretrained layers.
We explore the hypothesis that these intermediate layers might be directly exploited by linear probing.
We propose a method, Head-to-Toe probing (H EAD 2T OE ), that selects features from all layers of the source model to train a classiﬁcation head for the target-domain.
In evaluations on the Visual Task
Adaptation Benchmark (VTAB), Head2Toe matches performance obtained with ﬁne-tuning on average, but critically, for out-of-distribution transfer, Head2Toe outperforms ﬁne-tuning.
I NTRODUCTION
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.
; Zhuang et al.
(2020) for recent application-speciﬁc 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 ( L INEAR ) , often yields signiﬁcant improvements in performance on the target task over training the network from scratch (Kornblith et al., 2019).
An alternative to L INEAR is ﬁne-tuning ( F INE T UNING ) , 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 signiﬁcantly higher than L INEAR .
Furthermore, since the entire network is ﬁne-tuned, the entire set of new weights needs to be stored for every target task, making F INE T UNING impractical when working on edge devices or with a large number of target tasks.
However, F INE T UNING is often preferred over L INEAR since it consistently leads to better performance on a variety of target tasks even when data is scarce (Zhai et al., 2019).
F INE T UNING ’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 F INE T UNING learn from few examples successfully?
We conjecture that F INE T UNING better leverages existing internal representations rather than discovering entirely new representations; F INE T UNING exposes existing features buried deep in the net for use by the classiﬁer.
Under this hypothesis, features needed for transfer are already present in the pretrained network and might be identiﬁed directly without ﬁne-tuning the backbone itself .
[Figure]
Head2Toe
[Figure]
[Figure]
Domain Affinity
A H
T
A L i n e a r
C a l t e c h
C
I F
A RC l e v r - D
i s t
C l e v r - C o un t R e t i n o p a t h y
D M L a b d S
p r - O r i e n t d S p r - L
o c
D T D E u r o S
A T
K I TT I - D
i s t
F l o w e r s
P e t s C a m e l y o n
R e s i s
cs N O R B - A z
i m s N O R B - E l e v
S un
S V H N natural structured specialized
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 L INEAR transfer approach to the augmented representation (Fig.
1-left).
Our approach leads to signiﬁcant improvements over L INEAR as shown in Fig.
1-right.
Our key contributions are as follows:
We perform analyses to better understand two factors that affect the beneﬁt of incorporating intermediate representations: feature selection and the degree to which a target domain is out-of-distribution (OOD).
We introduce head-to-toe probing (H EAD 2T OE ), a simple recipe for selecting relevant features from intermediate representations for few-shot generalization.
On the VTAB collection of data sets, we show that H EAD 2T OE outperforms L INEAR and matches the performance of more computationally costly F INE T UNING .
Critically, H EAD 2T OE outperforms F INE T UNING 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 H EAD 2T OE improves on the state of the art for transfer learning
P RELIMINARIES
Source domain and backbone models.
In our experiments, we use source models pretrained on ImageNet-2012 (Russakovsky et al., 2015), a large scale image classiﬁcation benchmark with 1000 classes and over million natural images.
We benchmark H EAD 2T OE 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 classiﬁcation tasks, each having between 2 toclasses 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 artiﬁcial 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 classiﬁcation in the source domain should also be relevant for linear classiﬁcation in the target domain, and L INEAR should yield performance beneﬁts. 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 d S p r - L
o c
C l e v rD i s t s N O
R B - E l e v
S V H N .
s
t
r u c t u r e d s N O
R B - A z i m d S p rO r i e n t
D M L a b
R e t i n o p a t h y
C l e v rC o un t
E u r o S
A T
C a m e l y o n .
s p e c
i a li z e d .
a v g
K I TT ID i s t
R e s i s c
S un .
n a t u r a l
C
I F
A RF l o w e r s 102 D T D
C a l t e c h 101 P e t s
[Figure]
Scratch the target domain.
Therefore, we might quantify the source-target distribution shift in terms of how beneﬁcial L INEAR is relative to training a model from scratch (denoted S CRATCH ):
[Equation]
In Fig.
2, we sort the 19 VTAB-1k target tasks by their domain afﬁnity 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 L INEAR , F INE T UNING , and S CRATCH baselines.
Consistent with the literature, F INE T UNING performs as well or better than the other two baselines.
For in-distribution targets (right side of graph), L INEAR matches F INE T UNING ; for OOD targets (left side of graph), L INEAR is worse than F INE T UNING .
With distribution mismatch, the source network may ﬁlter information available in lower layers because it is not needed for the source task but is essential for the target task.
Observing F INE T UNING performs better than S CRATCH even in OOD tasks, we hypothesize that intermediate features are key for the F INE T UNING , since if learning novel features was possible using limited data, training the network from scratch would work on-par with F INE T UNING .
Motivated by this observation, H EAD 2T OE probes the intermediate features of a network directly and aims to eliminate the need for ﬁne-tuning.
H EAD 2T OE U TILIZATION OF P RETRAINED B ACKBONES 3. Y OUR R EPRESENTATIONS ARE R ICHER T HAN Y OU T HINK
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 ResNetbackbone (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 L INEAR 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 L INEAR using this O RACLE is shown in Fig. 3-left. We observe a negative correlation (-0.745, Spearman)
[Figure]
Domain Affinity
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure]
[Table]
natural structured specialized
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
T e s t
A cc u r a c y
Pets
Linear
Method
Avg
Spcl
Strc
Natr
Linear 56.
65.79 Control 58.
67.03 Oracle 60.
between the domain afﬁnity 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 signiﬁcantly 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 L INEAR 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 O RACLE 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 (C ONTROL ).
Note that this experiment bears similarity to ensembling (Zhou et al., 2002), which is known to bring signiﬁcant 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. H EAD 2T OE
Motivated by our observations in the previous section, we hypothesize that we can attain—or possibly surpass—the performance of F INE T UNING without changing the backbone itself by using
L INEAR 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 overﬁtting.
In this section, we explore various regularization and feature selection methods, which we then incorporate into our method,
H EAD 2T OE .
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:
[Equation]
where the subscript denotes a layer index, h
= x is the input, f is the activation function, W (cid:96) is the weight matrix of layer (cid:96) , and z L is the logit vector used for classiﬁcation.
When transferring a pretrained network to a target task using L INEAR , we discard the last layer of the pretrained network and train a new set of linear weights, W (cid:48)
L , such that predictions (logits) for the new task are obtained by z (cid:48)
L = h
L − 1 W (cid:48) L .
Head2Toe.
Consider a simple scheme that augments the backbone embedding with activations from all layers of the network, such that:
[Equation]
where a (cid:96) ( . ) denotes a ﬁxed function to reduce the dimensionality of the activation vector at a given layer (cid:96) .
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, h all can exceed a million elements, and W all is underconstrained by the training data, leading to overﬁtting.
Further, without care W all may become so large as to be impractical for deploying this model.
We can address these issues by selecting a subset of features before training the targetdomain classiﬁer.
Equivalently, we can zero out the non-selected rows of W all .
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 (cid:96)norm of the rows of the matrix to be sparse, and is deﬁned as:
[Equation]
To determine which features are most useful for the task, the linear head is trained with group-lasso regularization on W all .
In contrast to the approach taken by Argyriou et al.
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 (cid:96) 2 norms of the rows of W all , s , are used to decide which features are most relevant.
We refer to s i 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 ﬁnal logit mapping.
Feature selection alone provides strong regularization, therefore during the ﬁnal training we don’t use any additional regularization.
We make two comments on this.
First, because the initial round of training W all with the group-lasso regularizer is used only to rank features by importance, the method is quite robust to the regularization coefﬁcient; it simply needs to be large enough to distinguish the contribution of the individual features.
Second, interpreting s i as the importance of feature i depends on all features having the same dynamic range.
This constraint will often be satisﬁed 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 H EAD 2T OE .
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 W all to obtain s ) due to the reduced number of features in the second step.
Overall, H EAD 2T OE together with its validation procedure requires 18% more operations compared to training W all alone (details shared in Appendix B).
Verifying H EAD 2T OE .
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 [Figure] by their relevance score, s i , 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 ﬁgure shows, test accuracy decreases monotonically with the offset, indicating that the relevance score predicts the importance of including a feature in the linear classiﬁer.
H EAD 2T OE 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 overﬁt 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 (cid:96) 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 ﬁgures for all 19 datasets in Appendix E. Fig.
4-right shows the distribution of features selected over different layers for both strategies.
Note that H EAD 2T OE ’s use of a ﬁxed 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 F INE T UNING .
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.
R ELATED W ORK
Transfer Learning is studied extensively in the literature and used widely in practice.
Most similar to our work is ELMo (Peters et al., 2018), which averages intermediate representations of a language model (2 LSTM embeddings) using a learned linear combination (a softmax).
ELMo requires embeddings to be same size and is most similar to the suboptimal layer-selection baseline shown in
Fig.
4-center.
Peters et al.
compared F INE T UNING with L INEAR and reported a small gap in performance for ELMo and BERT models.
Kornblith et al.
(2019) showed that ImageNetperformance correlates highly with L INEAR and F INE T UNING .
Using a large pool of pretrained models, Renggli et al.
(2020) examined various metrics to decide which backbone to transfer.
Given the ever increasing size and never saturating performance of pretrained models, the importance of reducing the cost of F INE T UNING models is stated in the literature regularly.
Methods like featurewise transformations (Bilen & Vedaldi, 2017), residual adapters (Houlsby et al., 2019; Rebufﬁ et al.,
2017), Diff-pruning (Guo et al., 2021) and selective ﬁne-tuning (Guo et al., 2019; Fu et al., 2021) are proposed in order to reduce the cost of storing ﬁne-tuned models.
Zhang et al.
(2021) showed that selecting a sub-network of a pretrained model can increase OOD generalization.
However, none of these methods match the simplicity of training (and storing) a linear classiﬁer.
Teerapittayanon et al.
(2016), Kaya et al.
(2019), and Zhou et al.
(2020) studied intermediate representations to reduce ”overthinking” and thus provide better early-exit criteria.
Similarly Baldock et al.
(2021) showed a correlation between early classiﬁcation of a sample and how easy its classiﬁcation is.
Alternatives to transfer learning exist to learn from few labeled examples.
Meta-learning (Thrun &
Pratt, 1998; Schmidhuber, 1987; Schmidhuber et al., 1997) is used (among other things) to tackle the few-shot learning problem (i.e learning from limited data) from a learning-to-learn perspective; see Hospedales et al.
(2020) for a comprehensive survey.
In-between the representation learning and learning-to-learn paradigms lie approaches that use multi-domain training as an inductive bias, either without (Dvornik et al., 2020;
Triantaﬁllou et al., 2021; Li et al., 2021b;a) or with (Liu et al., 2021) meta-learning.
However, in the large-scale setting F INE T UNING remains a top-performing approach to few-shot classiﬁcation (Dumoulin et al., 2021).
Feature selection aims to ﬁnd the most relevant features from the input and is studied extensively in the machine-learning literature.
Approaches can be grouped according to whether labeled data is used—supervised (Nie et al., 2010) or unsupervised (He et al., 2005; Balın et al., 2019;
Atashgahi et al., 2020)—or what high-level approach is taken—ﬁlter methods (Blum & Langley, 1997), wrapper methods (Kohavi & John, 1997), or embedded methods (Yuan & Lin, 2006).
Most relevant to our work are embedded supervised methods as they have good scaling properties while achieving the best performance which is vital in our setting with over a million features.
Embedded supervised feature selection methods use a cost function to iteratively reﬁne the subset of features selected and popular approaches include forward selection (Viola & Jones, 2001; Borboudakis &
Tsamardinos, 2019), backward selection (pruning) (Mozer & Smolensky, 1989; Guyon et al., 2004) and regularization/feature-ranking based methods (Yuan & Lin, 2006; Blum & Langley, 1997; Zhao et al., 2010).
Most relevant to our work is Argyriou et al.
; Nie et al.
(2010), both of which uses (cid:96) 2 , 1 regularization to select features, however their approach requires matrix inversions which is not practical in our setting.
We point interested readers to the survey of Gui et al.
(2017) and book of Boln-Canedo et al.
for a detailed discussion on the topic.
E VALUATING H EAD 2T OE
We evaluate H EAD 2T OE on the VTAB-1k benchmark using two popular vision architectures,
ResNet-50 (Wu et al., 2018) and ViT-B/16 (Dosovitskiy et al., 2021), both pretrained on ImageNet2012.
ResNet-50 consists of 50 convolutional layers.
To utilize the intermediate convolutional features, we aggregate spatial information using average pooling on non-overlapping patches, as explained in Section 3.
We adjust the pooling size to target a ﬁxed dimensionality of the representation.
For example, for a feature map of shape 20 × 20 × 128 and a target size of 512, we average-pool disjoint patches of size 5 × 5 , resulting in 4 features per channel and 1024 features in total.
This helps us to balance different layers in terms of number features they contribute to the concatenated embedding.
We normalize features coming from each layer to a unit norm.
This scaling preserves the relative magnitude of features within a layer while accounting for inter-layer differences and works better than normalizing each feature separately or not normalizing at all.
ViT-B/16 consists of 12 multi-headed self-attention layers.
When we aggregate the output of the self-attention layers, we perform 1-D average pooling over the patch/token dimension choosing the pooling size to match a target number of features as before.
Given that token dimension is permutation invariant, 1-D average pooling might not be the best choice here and better aggregation function can provide further gains.
We share a detailed list of intermediate representations utilized for each architecture in Appendix H.
H EAD 2T OE selects a subset of features and trains a linear classiﬁer without regularization on top of the selected features.
We compare H EAD 2T OE with regularization baselines that utilize all features without doing feature selection.
These baselines are denoted as +All- (cid:96) 1 , +All- (cid:96) 2 and +All- (cid:96) 2 ,according to the norm they use.
We perform ﬁve-fold cross validation for each task and method in order to pick the best hyperparameters.
All methods search over the same learning rates and training steps (two values of
[Table]
yon
[Table]
nop
[Table] l
[Table] v r [Table]
oun t
[Table]
l
[Table] v r- [Table] i s
[Table]
M L [Table] b
K I TT
I- [Table]
i [Table] t d
[Table]
p r
[Table] c [Table]
r [Table] r i [Table]
NO
[Table] [Table]
i m
[Table]
each). Methods that leverage intermediate features (i
[Table]
e., regularization baselines and H EAD 2T OE ) additionally search over regularization coefﬁcients and the size of the aggregated representation at each layer.
The F INE T UNING baseline searches over 4 hyper-parameters; thus the comparison of
H EAD 2T OE , which searches over 24 values, to ﬁne-tuning might seem unfair.
However, this was necessary due to ﬁne-tuning being signiﬁcantly more costly than training a linear classiﬁer, even with intermediate features.
We repeat each evaluation using 3 different seeds and report median values and share standard deviations in Appendix C. More details on hyper-parameter selection and best values can be found in Appendix A.
5. R ES N ET -
The top half of Table 1 presents results on the 19 VTAB-1k target domains when transferring from a pretrained ResNet-50 architecture.
On average, H EAD 2T OE slightly outperforms all other methods, including F INE T UNING (see rightmost column of Table).
H EAD 2T OE , along with the regularization baselines that use intermediate layers, is far superior to L INEAR , indicating the value of the intermediate layers.
And H EAD 2T OE is superior to the regularization baselines, indicating the value of explicit feature selection.
Among the three categories of tasks, H EAD 2T OE excels relative to the other methods for the specialized category, but does not outperform F INE T UNING for the natural and structured categories. Does H EAD 2T OE select different features for each task?
Which layers are used more frequently? In Appendix F, we show the distribution of features selected across different layers and the amount of intersection among features selected for different tasks and seeds. We observe a high variety among layers and features motivating the importance of performing feature selection for each task separately.
H EAD 2T OE exceeds F INE T UNING performance, but requires only 0.5% of FLOPs during training on average.
Similarly, the cost of storing the adapted model is reduced to 1% on average.
We discuss H EAD 2T OE ’s computational and storage costs in detail in Appendix B.
[Figure]
5. V I T-B/
Results for ViT-B/16 are shared in the bottom half of Table 1.
As with the ResNet-50 architecture, H EAD 2T OE achieves the best accuracy among methods that keep the backbone ﬁxed:
H EAD 2T OE improves accuracy over L INEAR by about 10% on average (in absolute performance), and H EAD 2T OE outperforms the regularization baselines that include intermediate features but that do not explicitly select features.
The story for H EAD 2T OE versus F INE T UNING is a bit more complicated.
The F INE T UNING recipe used by Dosovitskiy et al.
(2021) includes many tweaks such as selective data-augmentation, image upsampling, learning rate warm-up, and gradient clipping.
In order to make a direct comparison with L INEAR , in the next-to-last row of the Table we report the performance of F INE T UNING without all of these tweaks; H EAD 2T OE obtains about a 1% absolute accuracy improvement over F INE T UNING in this setting.
However, when some of the tweaks are incorporated into F INE T UNING —in particular gradient clipping—F INE T UNING achieves a 3% improvement over H EAD 2T OE . Since it is a linear classiﬁer, adding gradient clipping to H EAD 2T OE doesn’t improve the results.
That being said, H EAD 2T OE retains its competitive advantage in terms of resource costs and its accuracy would likely similarly beneﬁt from the development of more sophisticated tricks (notably moving beyond the naive 1-D pooling of ViT features), much like
F INE T UNING has gained over time.
5. S CALING H EAD 2T OE
In Fig.
5 we vary the number of intermediate features used for each of our pretrained backbones.
We observe that including all layers always performs better on average.
However when varying the number of target features for each layer, we observed 2 distinct sub-group of target tasks that behave differently as the number of features increases.
This observation informed our discussion to include both small and large target sizes in our validation hyper parameter search.
C ONCLUSION
In this work, we introduced H EAD 2T OE , an approach that extends linear probing (L INEAR ) by selecting the most relevant features among a pretrained network’s intermediate representations.
We show that doing so greatly improves performance over L INEAR and allows the approach to reach a performance competitive with—and in some cases superior to—F INE T UNING .
Our ﬁndings challenge the conventional belief that F INE T UNING is required to achieve good performance on OOD tasks.
While more work is needed before H EAD 2T OE can realize the full computational beneﬁts of linear probing, our work paves the way for applying new and more efﬁcient feature selection approaches and for experimenting with H EAD 2T OE probing in other domains such as regression, video classiﬁcation, object detection, reinforcement learning, and language modelling domains.
A V ALIDATION P ROCEDURE FOR D IFFERENT A LGORITHMS
We perform hyper-parameter validation for each VTAB task separately.
For all methods, we chose the learning rate and the total number of training steps using the grid lr = 0 .
01 and steps = 500 , 5000 , following the lightweight hyper-parameter sweep recommended by the VTAB benchmark (Zhai et al., 2019).
For regularization baselines (cid:96) 1 , (cid:96) 2 and (cid:96) 2 , 1 we search for regularization coefﬁcients using (0 . 00001 , 0 . 0001 , 0 . 001) .
We include an extra value in this setting in order to account for the overhead introduced by H EAD 2T OE .
For H EAD 2T OE we choose (cid:96) 2 , 1 regularization coefﬁcients from (0 . 001 , 0 . 00001) and target feature sizes from (1024 , 16384 , 40000) for ResNet-50 and (768 , 15360 , 32448) for
ViT-B/16.
After calculating feature scores H EAD 2T OE validates the following fractions: (0 . 0005 , 0 .
05 , 0 . 1) and thus requires 18% more operations compared to other regularization baselines.
Note that this is because initial training to obtain feature scores s i is performed once and therefore searching for optimal number of features has a small overhead.
Hyper parameters selected by H EAD 2T OE for each VTAB task are shared in Table 2 and Table 3.
Next we explain how this overhead is estimated. [Table]
B C OST OF H EAD 2T OE
We evaluate different values of F and pick the value with best validation performance.
Cost of H EAD 2T OE consists of three parts: (1) C I : Cost of calculating the representations using the pretrained backbone.
(2) C all : cost of training the initial head W all (in order to obtain s i ’s) (3) (cid:80) f C F = f : total cost of validating different values of F .
Cost of validating a fraction value f , assuming equal number of training steps, is equal to C f = C all ∗ f .
Therefore relative cost of searching for F is equal to the sum of fractions validated (in comparison to the initial training of ( s )).
In Table 4, we compare cost of running H EAD 2T OE adaptation with F INE T UNING .
H EAD 2T OE uses the backbone once in order to calculate the representations and then trains the W all , whereas F INE T UNING requires a forward pass on the backbone at each step. Therefore we calculate the cost [Table]
[Table]
of ﬁnetuning for t steps as C I · t .
Similarly, cost of H EAD 2T OE is calculated as C I + C all · t .
The overall relative cost of C all increases with number of classes C and number of features considered
N .
As shown in Table 1-top, H EAD 2T OE obtains better results than F INE T UNING , yet it requires 0.5% of the FLOPs needed on average.
After adaptation, all methods require roughly same number of FLOPs for inference due to all methods using the same backbone.
The cost of storing models for each task becomes important when the same pre-trained backbone is used for many different tasks.
In Table 4 we compare the cost of storing different models found by different methods.
A ﬁnetuned model requires all weights to be stored which has the same cost as storing the original network, whereas L INEAR and H EAD 2T OE requires storing only the output head: W linear .
H EAD 2T OE also requires to store the indices of the features selected using a bitmap.
Even though H EAD 2T OE considers many more features during adaptation, it selects a small subset of the features and thus requires a much smaller ﬁnal classiﬁer (on average 1% of the F INE T UNING ).
Note that hyper parameters are selected to maximize accuracy, not the size of the ﬁnal model.
We expect to see greater savings with an efﬁciency oriented hyper-parameter selection.
C S TANDARD D EVIATIONS FOR T ABLE
In Table 5, we share the standard deviations for the median test accuracies presented in Table 1.
On average we observe L INEAR obtains lower variation as expected, due to the limited adaptation and convex nature of the problem.
Head2Toe seem to have similar (or less) variation then the other regularization baselines that use all features.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
- D i
[Table]
t d
[Table]
p r-
L o c d S p r-
O r
i s NO
RB - A z i m
[Table]
h
[Table]
yon
[Table]
[Table]
nop
[Table]
oun
[Table]
C
[Table]
r
[Table]
D i
[Table]
D M
[Table]
b
[Table]
D A DDITIONAL P LOTS FOR E XPERIMENTS USING S INGLE A DDITIONAL L AYER
Test accuracies when using a single additional intermedieate layer from a pretrained ResNetbackbone are shown in [Table] [Figure]
E A DDITIONAL P LOTS FOR L AYER /F EATURE - WISE S ELECTION C OMPARISON
We compare layer-wise selection strategy discussed in Section 3 to H EAD 2T OE in Fig.
For allmost all datasets, feature-wise selection produces better results.
Retinopathy and Flowers-102 are the only two datasets where the layer-wise strategy performs better.
F A DDITIONAL P LOTS FOR R ES N ET -50 R ESULTS Improvement
of H EAD 2T OE over ﬁne-tuning test accuracy for ResNet-50 backbone is shown in
Fig. 8. Similar to earlier plots, we observe a clear trend between being OOD and improvement in accuracy: H EAD 2T OE obtains superior few-shot generalization for most of OOD tasks. We also share the distribution of features selected for each task in Fig. Since different tasks might have different number features selected, we normalize each plot such that bars add up to 1.
Overall, features from later layers seem to be preferred more often. Early layers are preferred, especially for
OOD tasks like Clevr and sNORB. We observe a diverse set of distributions for selected features, which reﬂects the importance of selecting features from multiple layers. Even when distributions match, H EAD 2T OE can select different features from the same layer for two different tasks.
Therefore, next, we compare the indices of selected features directly to measure the diversity of features selected across tasks and seeds.
Transfer across Tasks In Fig.
9-left we investigate how features selected using a task i performs when evaluated on a task j .
Each pixel represents the average accuracy over 3 seeds.
For each column we subtract the diagonal term (i.e. self-transfer, i = j ) to obtain delta accuracy for each
[Figure]
[Figure]
[Figure]
[Figure]
dataset.
For most tasks, using a separate task for feature selection hurts performance.
Results in
Flowers-102 and Pets get better when other datasets like sNORB-Elev is used.
Crucially no single dataset (rows of Fig.
9-left) seems to get best transfer, which highlights the importance of doing the feature selection during the adaptation (not beforehand).
Note that, in practice (and in the literature) evaluation for each task is done in isolation, i.e. the data from other tasks are not available
[Figure]
[Table]
[Table]
i
[Table]
c
[Table]
o n
Transfer of Selected Features
[Table]
[Figure]
[Figure]
Similarity of features selected In Fig.
9-right we investigate the intersection of features selected by H EAD 2T OE .
We select 2048 features for each task from the pool of 29800 features (same experiments as in Fig.
4-center).
For each task H EAD 2T OE selects a different subset of features.
We calculate the fraction of features shared between each subset.
For each target task we run 3 seeds resulting in 3 sets of features kept.
When comparing the similarity across seeds (diagonal terms), we average over 6 possible combinations among these 3 sets.
Similarly, when comparing two different datasets we average over 9 possible combinations.
Results show that features selected by complementary tasks like dSprites-Loc and dSpr-Orient (or sNORB tasks) have high overlap.
Both dSprites and sNORB images are generated artiﬁcially and have similar characteristics.
This is reﬂected in the heat-map (about 40% overlap).
Similarly, features selected for CIFAR-100 and Caltech101 tasks seem to have high overlap.
Features selected also vary across different runs (seeds).
This variation seem to be higher for highly in-domain datasets like Flowers102 and Pets, which seems to correlate with H EAD 2T OE ’s poor performance.
Similarly dSprites-Loc seems to have higher variance across different seeds, which correlates with the outlier performance of H EAD 2T OE on this task when compared with F INE T UNING (see Fig.
We believe that understanding and reducing the inconsistency among different runs can be a promising direction to improve head-to-toe utilization performance.
Apart from a small fraction of datasets, most datasets seem to share less than 20% of the features, which highlights the importance of doing the feature selection for each target task separately.
Effect of Training Data In Fig.
11, we compare the performance of H EAD 2T OE with other baselines using reduced training data.
Fraction ( d f )=1 indicates original tasks with 1000 training samples.
For other fractions we calculate number of shots for each class as int (1000 ∗ f d /C ) where C is the number of classes and then sample a new version of the task accordingly.
For SUN-397 task, number of shots are capped at 1 and thus fractions below 0.75 lead to 1-shot tasks and thus results are all the same.
Overall we observe the performance of H EAD 2T OE improves with amount of data available, possibly due to the reduced noise in feature selection.
G A DDITIONAL P LOTS FOR S CALING B EHAVIOUR OF H EAD 2T OE
Scaling behaviour of H EAD 2T OE when using different feature target size and number of layers
over 19 VTAB-1k tasks is shown in Fig.
[Table]
1 b 4 u 2 b 4 up r e l o
g
i t s p
r
e
l
o g i t s[-
 p r e-] l o g i t s Flowers H D ETAILS OF I NTERMEDIATE R EPRESENTATIONS I NCLUDED
ResNet-50 has 5 stages.
First stage includes a single convolutional layer (root) followed by pooling.
We include features after pooling.
Remaining stages include 3,4,6 and 3 bottleneck units (v2) each.
Bottleneck units start with normalization and activation and has 3 convolutional layers.
We includes features right after the activation function is applied, resulting in 3 intermediate feature sets per unit.
Output after 5 stages are average-pooled and passed to the output head.
We include features before and after pooling.
with the output of the ﬁnal layer (logits), total number of locations where features are extracted makes 52.
ViT-B/16 models consists of multiple encoder units.
Each attention head consists of a self attention module and followed by 2 MLPs.
For each encoder unit, we include features (1) after layer-norm (but before self-attention) (2) features after self-attention (3-4) features after MLP layers (and after gelu activation function).
Additionally we use patched image (i.e. tokenized image input), pre-logits and logits.
[Figure]
[Figure]
[Table]
[Figure]
[Figure]
[Figure]
e t
S i z e
[Figure] [Figure]
t
S z e
[Figure]
T a r g e t
S i z e sNORB
[Figure]