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 explores how different strategies for fine-tuning affect in- and out-of-distribution performance. The authors contrast linear probing (updating only the parameters of the final linear layer), end-to-end fine-tuning (updating all parameters of the model) and a two-stage approach, where linear probing is followed by end-to-end fine-tuning. While end-to-end fine-tuning typically improves in-distribution performance, the authors show that it can also underperform linear probing out-of-distribution. The paper theoretically analyzes the tradeoffs in a simplified scenario with two-layer networks, finding that end-to-end fine-tuning can "distort" pre-trained features. Their experiments on a number of datasets including CIFAR, WILDS-FMoW and others, confirm the intuitions from their theory. The proposed mitigation strategy, a two-stage fine-tuning approach where end-to-end fine-tuning follows linear probing is found to be beneficial, especially out-of-distribution.   **Update:** The authors addressed most of the concerns raised by this and other reviews, and I am raising my score accordingly.