Fine-Tuning Distorts Pretrained Features and Underperforms Out-of-Distribution

1 INTRODUCTION . Pretraining a model on a large dataset before transferring to a downstream task ’ s training data substantially improves accuracy over training from scratch—for example , pretraining a ResNet-50 on unlabeled ImageNet boosts accuracy on CIFAR-10 from 94 % to 98 % ( Chen et al. , 2020a ; b ) . Highstakes applications such as poverty mapping in under-resourced countries ( Jean et al. , 2016 ) , selfdriving cars ( Yu et al. , 2020 ) , and medical diagnosis ( AlBadawy et al. , 2018 ) , require models that also generalize to circumstances not seen in the training distribution . In addition to testing on data drawn from the downstream task ’ s training distribution ( in-distribution ; ID ) it is increasingly important to test on data distributions unseen during training ( out-of-distribution ; OOD ) . After initializing with a pretrained model , two popular transfer methods are fine-tuning ( running gradient descent on all the model parameters ) , and linear probing ( tuning the head but freezing lower layers ) . In the ID setting it is well known that fine-tuning leads to better accuracy than linear probing ( Kornblith et al. , 2019 ; Zhai et al. , 2020 ; He et al. , 2020 ) , and even when testing OOD prior work usually fine-tunes all parameters of their model ( Hendrycks et al. , 2019a ; Miller et al. , 2021 ; Andreassen et al. , 2021 ) . Intuitively , fine-tuning all layers of a network can improve pretrained features by tailoring them to the specific task , while linear probing freezes these features . In this work , we investigate the OOD performance of fine-tuning and linear probing and find that surprisingly , fine-tuning often does worse than linear probing in the presence of large distribution shift . We experiment on six distribution shift benchmarks ( Breeds Living17 , Breeds Entity30 , DomainNet , CIFAR → STL , CIFAR10.1 , FMoW geo-shift ) , initializing with good pretrained features from MoCo-v2 ( Chen et al. , 2020b ) and CLIP ( Radford et al. , 2021 ) . While both methods offer gains over training from scratch , fine-tuning improves the average ID accuracy from 84 % to 86 % but brings down the OOD accuracy from 72 % to 67 % ( Figure 1 ) . When and why does fine-tuning underperform linear probing ? We theoretically consider fine-tuning a two-layer linear network in an overparameterized regression setting where the feature extractor layer has been pretrained to map high-dimensional inputs to useful , lower-dimensional , features . We prove that fine-tuning is worse than linear probing on worst-case OOD inputs when using high quality pretrained features . Even with an infinitesimally small learning rate , fine-tuning distorts pretrained features—the features of ID data are updated while those of OOD data remain unchanged . Since the head and feature extractor are simultaneously optimized during fine-tuning to a configuration that works well on ID training data , the head only accomodates the distorted features of ID points and performs poorly on the unchanged features of OOD points . Interestingly , we show that this feature distortion issue can not be simply fixed by early stopping—throughout the process of fine-tuning , we never pass through parameters that do well OOD . On the other hand , we show that linear-probing extrapolates better OOD because it preserves pretrained features , but does not do as well as fine-tuning ID because linear probing can not adapt the features to the downstream task . Technical challenges . Existing theoretical work on transfer learning focuses on linear probing ( Wu et al. , 2020 ; Tripuraneni et al. , 2020 ; Du et al. , 2020 ) . In contrast , the analysis of fine-tuning is scarce and challenging because it requires understanding the training dynamics , instead of only the loss function and its global minimizers . In fact , fine-tuning and training from scratch optimize the same training loss and only differ in their initializations ( pretrained vs random ) . A mathematical analysis that distinguishes them needs to capture properties of the different global minima that these algorithms converge to , a phenomenon that is sometimes theoretically referred to as the implicit regularization effect of initialization ( Neyshabur et al. , 2014 ) . Accordingly , our analysis reasons about the parameters that gradient methods pass through starting from the pretrained initialization , which is challenging because there is no known closed form for this trajectory . Two-layer linear networks are widely studied in the implicit regularization community ( Saxe et al. , 2014 ; Arora et al. , 2018 ) , however they analyze random and often small initializations which don ’ t capture pretraining . Algorithmic implications . Our theory says that fine-tuning fails because when trying to fit ID training data with a randomly initialized head , the feature extractor changes significantly for ID examples , making features for ID and OOD examples largely inconsistent . This can be fixed by initializing with a good head that does not need to be updated much during fine-tuning , reducing how much the feature extractor changes . This suggests the simple two-step strategy of first linearprobing to find a good head and then full fine-tuning ( LP-FT ) . Empirically , LP-FT outperforms fine-tuning and linear-probing , both ID and OOD . Even on CIFAR-10.1 ( small distribution shift ) , where fine-tuning is better for both ID and OOD , we find LP-FT outperforms fine-tuning on both metrics . LP-FT and vanilla fine-tuning use similar amounts of compute because the first step of linear probing is relatively very cheap . We note that while LP-FT has sometimes been used as a finetuning heuristic ( Kanavati & Tsuneki , 2021 ) , it has not been used for robustness / OOD accuracy , and we show that it addresses the ID-OOD tradeoff theoretically and empirically . Practical insights . Finally , we check whether fine-tuning fails and LP-FT works , for the reasons predicted by our feature distortion theory . As predicted by the theory , we find that : ( 1 ) fine-tuning indeed never matches the OOD accuracy of linear probing throughout the course of training , ( 2 ) finetuning changes the features for ID examples more than for OOD examples leading to distortions ( 3 ) fine-tuning can do better than linear probing OOD if the pretrained features are not very high quality ( MoCo-v1 instead of MoCo-v2 ) or the ID and OOD datasets are very close ( e.g. , CIFAR-10 and CIFAR-10.1 ) and ( 4 ) LP-FT indeed changes both ID and OOD features orders of magnitude less than fine-tuning does . 2 SETUP . Task and evaluation . Given training samples { ( x1 , y1 ) , . . . ( xn , yn ) } sampled from some distribution Pid , our goal is to learn a predictor f : Rd → Y to map inputs x ∈ Rd to outputs y ∈ Y . We evaluate predictors on their standard “ in-distribution ” ( ID ) performance Lid on new test samples drawn from Pid that the training data is also sampled from . We also evaluate classifiers on their “ out-of-distribution ” ( OOD ) performance Lood on test samples drawn from a new distribution Pood that is different from Pid . Formally , for some loss function ℓ , we evaluate classifiers on : Lid ( f ) = E ( x , y ) ∼Pid [ ℓ ( f ( x ) , y ) ] and Lood ( f ) = E ( x , y ) ∼Pood [ ℓ ( f ( x ) , y ) ] . ( 2.1 ) Models . In this work , we focus on predictors that leverage pretrained feature extractors . For convenience of such analyses , we parameterize the final predictor in terms of a linear “ head ” v ∈ V on top of some features gB ( x ) ∈ Rk for some “ base ” parameters B ∈ B . Formally , f is parameterized by base parameter B and head parameters v such that fv , B ( x ) = v⊤gB ( x ) . In our experiments ( Section 4 ) , gB is a deep network and in our theory ( Section 3 ) , gB is a linear projection . We assume access to some initial pretrained feature extractor B0 that is obtained by training on potentially large amounts of data from a distribution that could be different from Pid and Pood . We focus on two popular methods to learn a predictor fv , B given training data from Pid : ( i ) linear probing where B = B0 and the linear head is obtained by minimizing some loss ( e.g. , logistic loss for classification , squared loss for regression ) on the training data , and ( ii ) fine-tuning where both v and B are updated by performing gradient descent on some loss on the training data with B initialized at B0 . 3 THEORY : FINE-TUNING DISTORTS PRETRAINED FEATURES . We theoretically study the performance of different training methods that use a good pretrained feature extractor . In a linear setting , we characterize when and why fine-tuning , in which all model parameters are updated , can increase OOD loss due to feature distortion . We define the problem setting in Section 3.1 , and present the main result and proof sketch in Section 3.2 . Our analysis handles two key challenges which distinguish it from prior work on transfer learning in linear models ( Wu et al. , 2020 ; Tripuraneni et al. , 2020 ; Du et al. , 2020 ; Xie et al. , 2021a ) . Prior work focuses on linear probing , while we study fine-tuning where the resulting optimization problem is non-convex and we need to analyze the effect of initialization from a particular pretrained parameter setting . We also study overparameterized models where the number of training examples is less than the input dimension , reflecting the practical setting where neural networks achieve zero training loss . Here , training loss alone does not determine test performance—this fact makes the setting very relevant because both training from scratch and fine-tuning have same training loss but very different test performance , but also makes the analysis challenging . 3.1 LINEAR OVERPARAMETERIZED SETTING . Models . Recall from Section 2 that we parameterize predictors in terms of base and head parameters . In this section , we study models where the feature extractor is linear , i.e . fv , B ( x ) = v⊤Bx where B ∈ B = Rk×d , and v ∈ V = Rk . For simplicity , we assume the models are well-specified i.e . y = v⊤⋆ B⋆x where v⋆ ∈ Rk and B⋆ ∈ Rk×d . 1 Note that B⋆ and v⋆ are only unique up to rotations , i.e. , for any rotation matrix U , ( Uv⋆ ) T ( UB⋆ ) x = v T ⋆ B⋆x . 1We note that our main contribution—analysis of fine-tuning ( Theorem 3.1 ) does not require this wellspecified assumption . We compare fine-tuning with linear probing by adapting earlier work on linear probing which requires the well-specified assumption . Suppose we have a pretrained feature extractor B0 close to B⋆ , so minU∥B0 − UB⋆∥2≤ ϵ over rotation matrices U ∈ Rk×k—this follows from prior work in pretraining ( Tripuraneni et al. , 2020 ) . As in prior work suppose B⋆ , B0 have been orthogonalized to have orthonormal rows . Data distribution and evaluation . For our analysis , we focus on regression , where Y = R and ℓ is the squared loss . Let X ∈ Rn×d , X ̸= 0 be a matrix encoding n training points from Pid where each of the n rows is a training input . Let Y ∈ Rn be the corresponding outputs . We consider an overparameterized setting where 1 ≤ k < n < d and k < d− n. Intuitively , the input dimension d is high ( e.g. , 10K ) , feature dimension k is lower ( e.g. , 100 ) and n is in the middle ( e.g. , 5K ) . In this work , we are primarily interested in the out-of-distribution ( OOD ) performance . Since the OOD data can be arbitrary we follow prior work ( Rosenfeld et al. , 2021 ; Kamath et al. , 2021 ; Chen et al. , 2021b ) and consider the worst case loss over distributions ( equivalently , individual points ) of bounded norm : Lood ( v , B ) = max ∥x∥2≤1 ( v⊤⋆ B⋆x− v⊤Bx ) 2 = ∥B⊤⋆ v⋆ −B⊤v∥22 ( 3.1 ) Training methods . Given training data and a pretrained base parameter B0 , we study the two popular methods of linear probing ( LP ) and fine-tuning ( FT ) to learn the final predictor ( See Section 2 ) . Both methods involve optimizing the training loss via gradient descent ( or variants ) . In order to effectively analyze these gradient based algorithms , we study vanishing step sizes leading to gradient flows . Gradient flows can be thought of as a continuous time analog of gradient based methods and have been extensively studied in recent years as a way to understand gradient based methods ( Gunasekar et al. , 2017 ; Arora et al. , 2018 ; Du et al. , 2018 ) . Formally , for training loss L̂ ( v , B ) = ∥XB⊤v − Y ∥22 , the gradient flow differential equations for LP and FT are as follows . ∂tvft ( t ) = −∇vL̂ ( vlp ( t ) , Bft ( t ) ) , ∂tBft ( t ) = −∇BL̂ ( vft ( t ) , Bft ( t ) ) ( 3.2 ) ∂tvlp ( t ) = −∇vL̂ ( vlp ( t ) , B0 ) , ∂tBlp ( t ) = 0 , ( 3.3 ) initialized with Bft ( 0 ) = Blp ( 0 ) = B0 and vft ( 0 ) = vlp ( 0 ) = v0 . In practice , the head parameter v0 is initialized randomly—our results hold for any standard random initialization Glorot & Bengio ( 2010 ) , for example v0 ∼ N ( 0 , σ2I ) for any σ2 . Recall that the initial value of the base parameter B0 is assumed to be available and obtained via pretraining . The final LP and FT solutions are the limit points of the corresponding gradient flows : v∞ft = lim t→∞ vft ( t ) and B∞ft = lim t→∞ Bft ( t ) ( 3.4 ) v∞lp = lim t→∞ vlp ( t ) and B∞lp = lim t→∞ Blp ( t ) = B0 ( 3.5 )

Summary: This paper discovers an interesting behavior of model fine-tuning: the performance is worse compared to linear probing on OOD data (i.e., data from other domains), especially when the distribution shift between inner distribution and out of distribution are big. The explanation provided in the paper is that fine-tuning distorts the feature representations, overfits on inner distributions, and thus has a higher error on OOD data. The authors also provide a simple solution to this issue by fine-tuning with a classification head initialized from linear probing and had better results in all the benchmarks they have in the paper.