Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution

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). Achieving high in-distribution accuracy is not enough: high-stakes applications such as poverty mapping in under-resourced countries (Jean et al., 2016), self-driving 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). OOD accuracy can be much lower than ID accuracy; for example, an ImageNet pretrained ResNet-50 fine-tuned on CIFAR-10 gets 98% accuracy on CIFAR-10 (ID) but 82% on STL (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),Probing is commonly used but usually for interpretability or assessing feature quality. 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 adapting them to the specific task, while linear probing simply inherits the frozen pretrained features.

In this work, we investigate the OOD accuracy of fine-tuning and linear probing and find that surprisingly, fine-tuning can do worse than linear probing in the presence of large distribution shift. We experiment on ten distribution shift benchmarks (Breeds Living17, Breeds Entity30, DomainNet, CIFAR $\to$ STL, CIFAR10.1, FMoW geo-shift, ImageNetV2, ImageNet-R, ImageNet-A, ImageNet-Sketch), 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 relative to linear probing from $83\%$ to $85\%$ but brings down the OOD accuracy from $66\%$ to $59\%$ (Figure 1).

Under what conditions 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 directions outside the span of the training data when using “good” pretrained features. Even with an infinitesimally small learning rate, fine-tuning distorts pretrained features—the features of ID training data are updated while those of OOD data change less. 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 (relative to linear probing) on the less changed features of OOD points. Interestingly, we show that this feature distortion issue cannot be simply fixed by early stopping—throughout the entire process of fine-tuning, we never pass through parameters that do well OOD (relative to linear probing). On the other hand, given “good” features, linear-probing extrapolates better OOD because it preserves pretrained features, but does not do as well as fine-tuning ID because linear probing cannot 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, analyses 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 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 this is a non-convex optimization problem and there is no known closed form for this trajectory. Two-layer linear networks are widely studied in the literature on implicit regularization (Saxe et al., 2014; Gunasekar et al., 2017; Gidel et al., 2019; Arora et al., 2018). However, they analyze random and often small initializations, which don’t capture pretraining.

Algorithmic implications. Our theory shows that fine-tuning underpeforms 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 a simple two-step strategy of first linear-probing 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. Prior work has used LP-FT (Levine et al., 2016; Kanavati & Tsuneki, 2021) (or variants such as layerwise fine-tuning (Howard & Ruder, 2018) or larger learning rates for the head layer (Prabhu et al., 2021))—however it has not been used for robustness / OOD accuracy, and we show that it addresses the ID-OOD tradeoff theoretically and empirically. Note that LP-FT is not meant to be a SOTA method but a simple, principled way to get good ID and OOD accuracy—we hope our analysis inspires even better methods for robust fine-tuning.

Empirical validation. Finally, we check whether fine-tuning underperforms 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 (if the pretrained features are good, and OOD shift is large); (2) fine-tuning changes the features for ID examples more than for OOD examples, leading to distortions; (3) LP-FT indeed changes both ID and OOD features $10\times-100\times$ less than fine-tuning does; (4) 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 (5) LP-FT gets the best of both worlds, better accuracies than fine-tuning and linear probing, both ID and OOD (Figure 1).

Setup

We assume access to some initial pretrained feature extractor $B_{0}$ that is obtained by training on potentially large amounts of data from a distribution that contains unlabeled or weakly supervised $x$ inputs from $P_{\mathsf{id}}$ and $P_{\mathsf{ood}}$. We focus on two popular methods to learn a predictor $f_{v,B}$ given training data from $P_{\mathsf{id}}$: (i) linear probing where $B=B_{0}$ 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 $B_{0}$.

Theory: fine-tuning distorts pretrained features

Our goal is to understand under what conditions fine-tuning does worse than linear probing out-of-distribution (OOD).For example, without additional assumptions we can have $P_{\mathsf{id}}=P_{\mathsf{ood}}$ and so the same method will do better both ID and OOD. We consider a linear setting (feature extractor $g_{B}$ is linear) where the pretrained features are “good” and the OOD shift is large (Section 3.1). We prove our main result: that fine-tuning, in which all model parameters are updated, distorts features and gets suboptimal OOD error (Section 3.2, Theorem 3.3). We use this result to show that linear probing gets better OOD error but worse ID error than fine-tuning (Section 3.3). Finally, we explain why linear probing then fine-tuning can mitigate this ID-OOD tradeoff (Section 3.4).

Our analysis handles two key challenges which distinguishes 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. We also study overparameterized models where the training loss alone does not determine test performance—this captures the fact that both training neural networks from scratch and fine-tuning them have the same training loss but very different test performance. However, it also makes the analysis challenging because we need to reason about the trajectory of gradient methods starting from a pretrained initialization, which has no known closed form.

Pretraining coverage intuition: Intuitively, the existence of $B_{\star}$ corresponds to assuming that there exists a shared set of useful features for ID ($P_{\mathsf{id}}$) and OOD ($P_{\mathsf{ood}}$). We also assume that $B_{0}$ is close to $B_{\star}$—one way this can happen is if pretraining is done on large scale data and has seen unlabeled or weakly supervised $x$ inputs that cover the support of $P_{\mathsf{id}}$ and $P_{\mathsf{ood}}$. Formally, the task diversity assumption in Tripuraneni et al. (2020) is sufficient (but not necessary) for obtaining a good $B_{0}$. In our paper we show that even if we have these good features, fine-tuning can distort them and lead to low OOD accuracy.

Training methods. Given training data and a pretrained feature extractor $B_{0}$, we study the two popular methods of linear probing (LP) and fine-tuning (FT) to learn the final predictor. 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 analogue 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 $\widehat{L}(v,B)=\|XB^{\top}v-Y\|_{2}^{2}$, the gradient flow differential equations for LP and FT are as follows:

initialized with ${B_{\mathsf{ft}}}(0)={B_{\mathsf{lp}}}(0)=B_{0}$ and ${v_{\mathsf{ft}}}(0)={v_{\mathsf{lp}}}(0)=v_{0}$. In practice, the head parameter $v_{0}$ is initialized randomly—our results hold for any standard random initialization (Glorot & Bengio, 2010), for example $v_{0}\sim\mathcal{N}(0,\sigma^{2}I)$ for any $\sigma^{2}$, or zero initialization where $v_{0}=0$. Recall that the initial value of the feature extractor $B_{0}$ is obtained via pretraining.

The final LP and FT solutions are the limit points of the corresponding gradient flows:

2 Fine-tuning distorts pretrained features

The more common method of using a pretrained feature extractor is fine-tuning (FT) which typically improves ID performance relative to linear probing (LP). In this section, we show theoretically that FT can distort features leading to poor OOD performance. We first present the key intuitions demonstrating potential issues of FT and then present our formal theorem lower bounding the OOD error of FT (Section 3.2.2).

There are two main observations that we use to characterize when and why FT has higher OOD error than linear probing.

1. Features get distorted: representations change only in the ID subspace (i.e., subspace spanned by the training data) and are unchanged in the orthogonal subspace. To see this, we take the derivative of the training loss $\widehat{L}(v,B)=\|XB^{\top}v-Y\|_{2}^{2}$ with respect to the feature extractor parameter $B$:

By definition, if $u$ is a direction orthogonal to the training subspace $S=\mbox{rowspace}(X)$, then $\nabla_{B}\widehat{L}(v,B)u=0$, that is the gradient updates to $B$ do not modify $Bu$ for $u\in S^{\perp}$. However, the gradient is non-zero for directions $u$ in the ID subspace and the corresponding features $Bu$ change across the fine-tuning process. We call this feature distortion: the features in some directions are changed but not others. Next, we explain why this can lead to high OOD error.

2. Distorted features can lead to higher OOD error. Consider a toy example (Figure 2) where $d=2$ and the dimensionality of the representations $k=1$. The linear head $v$ is a scalar quantity that denotes how much the feature extractor $B$ has to be scaled by. Suppose the ID-subspace is the $x$-axis. There are different ways of fitting the ID subspace depending on the feature extractors $B$ as shown in the Figure—both fine-tuned and linear probed estimators match the true parameter in the ID subspace (since ${w_{\mathsf{lp}}},{w_{\mathsf{ft}}},w_{\star}$ have the same projection on the $x$-axis). If the feature extractor were optimal or scaled versions of the optimal, good performance on the ID subspace would translate to good performance everywhere, even in directions orthogonal to the ID subspace. However, in FT, the features change only for inputs in the ID subspace (see (1)) and thus the updated features are not simply scaled but distorted. In Figure 2, this corresponds to the feature extractor $B_{0}$ changing along the $x$-axis. In this case even if the ID error is low, error in directions orthogonal to the ID subspace can be high, leading to high OOD error.

The only way the pretrained features are not distorted and only scaled during FT is if the initial feature extractor $B_{0}$ is exactly aligned with the ID subspace. In Figure 2, if $B_{0}$ is along the $x$-axis (the ID subspace), then updating the features exclusively along the $x$-axis would simply scale the initial features. In this case linear probing and fine-tuning will have identical behavior. If the angle between $B_{0}$ and the $x$-axis is non-zero—which occurs with probability $1$ if the training data $X$ or pretrained feature extractor $B_{0}$ involves even a tiny amount of randomness e.g., from SGD in pretraining—the updates would lead to distortions. In high dimensions, we measure the alignment between $B_{0}$ and the ID subspace with the largest principal angle:

Let $A$ and $B$ be arbitrary subspaces, and $E$ and $F$ be matrices with orthonormal columns than span $A$ and $B$ respectively, with $r=\min(\dim(A),\dim(B))$. Then $\cos\theta_{\mathsf{max}}(A,B)=\sigma_{r}(E^{\top}F)$, which is the $r$-th largest singular value of $E^{\top}F$.

Note that $E,F$ are not unique in Definition 3.2, but $\sigma_{r}(E^{\top}F)$ is the same for every valid choice of $E$ and $F$. See Appendix A.1 for more information on principal angles.

2.2 General result on the OOD error of fine-tuning

Our main theorem lower bounds the OOD error of fine-tuning outside the span of the training data. In Section 3.3 we compare this lower bound with an upper bound on the OOD error of linear probing.

In the overparameterized linear setting, let $S^{\perp}=\mbox{rowspace}(X)^{\perp}$, $R_{0}=\mbox{rowspace}(B_{0})$, and $v_{\star},B_{\star}$ be the optimal parameters with $w_{\star}=B_{\star}v_{\star}$. If $\cos\theta_{\mathsf{max}}(R_{0},S^{\perp})>0$, then for all time steps $t$, the OOD error of the fine-tuning iterates $({B_{\mathsf{ft}}}(t),{v_{\mathsf{ft}}}(t))$ is lower bounded:

where $\varphi^{2}=\lvert(v_{0}^{\top}v_{\star})^{2}-(v_{\star}^{\top}v_{\star})^{2}\rvert$ is defined to be inital head alignment error and $\epsilon\geq d(B_{0},B_{\star})$ is the error in the pretrained feature extractor.

Proof sketch. Since the features do not change for examples in $S^{\perp}$ (perpendicular to the training data), we show that in order to achieve low error on $S^{\perp}$ the linear head ${v_{\mathsf{ft}}}(t)$ would have to become very similar to the optimal $v_{\star}$ at some time $t$. The head initialization $v_{0}$ is random (or zero) and likely to be far from $v_{\star}$ (measured by the alignment error $\varphi$), so the head would have to change a lot to get close to $v_{\star}$. As we see from the fine-tuning gradient flow (3.2), ${v_{\mathsf{ft}}}(t)$ and ${B_{\mathsf{ft}}}(t)$ change in a “coupled” manner, and a “’balancedness” invariant in Du et al. (2018) holds across the fine-tuning trajectory. Correspondingly, if ${v_{\mathsf{ft}}}(t)$ changes a lot and gets close to $v_{\star}$, the features ${B_{\mathsf{ft}}}(t)$ also change a lot for examples in $S$—we show that this would lead to high error on examples in $S$. Either way, fine-tuning would get some subspace ($S$ or $S^{\perp}$) of examples wrong, leading to high OOD error. The full proof appears in Appendix A.

Interpretations of various quantities. Quality of pretrained features ($\epsilon$). To unpack the bound consider a special case where the pretrained features are perfect ($\epsilon=0$). With perfect features, Proposition A.20 shows that linear probing gets zero OOD error. Theorem 3.3 shows that $L_{\mathsf{ood}}({v_{\mathsf{ft}}}(t),{B_{\mathsf{ft}}}(t))>0$ at all times $t$—so fine-tuning underperforms when the features are perfect. The $\epsilon>0$ case just captures the fact that even if the features are not perfect, fine-tuning can still get positive error. Ideally we would like the lower bound to increase if we have worse features (so “$+\epsilon$” instead of “$-\epsilon$” in the bound)—the reason we do not is that the errors of the pretrained feature extractor $d(B_{0},B_{\star})$ and the fine-tuning step can potentially cancel out.Intuitively this cancelation is very “unlikely” to happen, and we hope future work can capture this intuition.

Alignment error of random head initialization ($\varphi^{2}$). The lower bound (Equation A.14) increases as $\varphi^{2}$ increases i.e. alignment error increases because the gradient updates to the head and feature extractor are coupled. If the head were somehow initialized perfectly at $v_{\star}$, then fine-tuning updates may not increase the OOD error. However, when the head is randomly initialized (or initialized to zero) as is standard in fine-tuning, the alignment error is high, leading to high OOD error. We use this insight in Section 3.4 to show that better head initialization (namely via linear probing) improves OOD performance of fine-tuning.

Span of Training data ($S$). Theorem 3.3 lower bounds the error outside the span of the training data. If the training dataset is very small, then even the support of the ID distribution $P_{\mathsf{id}}$ may not be spanned by the training data, and the ID error can be large. Indeed, even in the ID setting Kornblith et al. (2019) show that linear probing can do better than fine-tuning if we have very few training examples, but fine-tuning does better on all 11 of their datasets once we have more than just 30 examples per class.

Conjectures for improved bounds. We believe it may be possible to improve $\cos\theta_{\mathsf{max}}(R_{0},S^{\perp})$ to the cosine of the minimum principal angle.Which would be a larger/better lower bound since the cosine of a smaller quantity is larger. This may look like a technicality but would be a substantial improvement, because it would imply that fine-tuning has error in every direction outside the training span, whereas we show that it would have errors in some directions. Our proof strategy requires the maximum principal angle (a crucial step is a variational characterization of the maximal principal angle in Lemma A.2—we use this in Step 1 of the proof in Appendix A to show that to get low OOD error ${v_{\mathsf{ft}}}(t)$ must become similar to $v_{\star}$).

3 Linear probing vs. fine-tuning

In this section, we use our main theorem on fine-tuning (Theorem 3.3) and adapt prior work on linear probing to show that linear probing is better than fine-tuning OOD, but worse ID, when the ID distribution has density on a lower $m<d$ dimensional subspace $S$, and $B_{0}$ is close to $B_{\star}$ (so we have “good” pretrained features).

Recall that the ID error is the expected mean-squared error over the ID distribution $P_{\mathsf{id}}$:

OOD comparison: Under mild non-degeneracy conditions, we show that as the feature extractor error $\epsilon$ goes to , linear probing does much better than fine-tuning OOD: the ratio of the losses goes to . The non-degeneracy conditions are similar to Section 3.2—we require that the training data cannot be exactly in the same direction or orthogonal to the pretrained features, formally that $\cos\theta_{\mathsf{max}}(R_{*},S)$ and $\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})$ are not where $R_{*}=\mbox{rowspace}(B_{\star})$. In the toy example in Figure 2, this means that $x_{\mathsf{id}}$ cannot be exactly in the same direction or orthogonal to $B_{0}^{\top}$—in these cases fine-tuning and linear probing get the same loss but in all other cases in the toy example in Figure 2 linear probing does better OOD.

In the linear overparameterized setting, under the ID subspace assumption (Assumption 3.4), if $\cos\theta_{\mathsf{max}}(R_{*},S)\neq 0$ and $\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})\neq 0$ where $R_{*}=\mbox{rowspace}(B_{\star})$, then,

This holds for all times $t$ for FT (and therefore also for the limit ${v_{\mathsf{ft}}^{\infty}},{B_{\mathsf{ft}}^{\infty}}$) and the LP iterates converge to ${v_{\mathsf{lp}}^{\infty}},B_{0}$ as a result of the gradient flow on a convex problem.

Intuitively, if the pretrained features are good, LP learns a near optimal linear head which has small OOD error (Lemma A.14) but fine-tuning has high OOD error (Theorem 3.3). We give a more formal version of Theorem 3.5 and a proof in Appendix A.3. If $P_{z}$ is isotropic Gaussian, we can get a better result: Theorem A.15 derives a threshold $T$ (in terms of $d,n,k$) where LP does better than FT if $\epsilon<T$, instead of just the asymptotic result ($B_{0}\rightarrow B_{\star}$). Theorem 3.5 requires that $\cos\theta_{\mathsf{max}}(R_{*},S)\neq 0$ and $\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})\neq 0$—intuitively, for any subspace a small perturbation would make these angles non-zero and the assumption would hold. To illustrate that these assumptions typically hold, Lemma A.16 in Appendix A proves that if $S$ is a random $m$-dimensional subspace then these angles are non-zero almost surely.

ID comparison: When the pretrained features have some error, we show that fine-tuning does better than linear probing ID because fine-tuning can update the features to fit the ID data.

If the pretrained features are perfect so that the optimal predictor can be written as a linear combination of the pretrained features ($w_{\star}=B_{\star}^{\top}v_{\star}\in\text{rowspace}(B_{0})$), then both linear probing and fine-tuning get zero ID error. However, if the pretrained representation has some error, and the training data satisfies a mild non-degeneracy condition, then LP has high ID error because there is no linear head on $B_{0}$ that fits the training data perfectly. FT, on the other hand, can update the features to find a new ${B_{\mathsf{ft}}^{\infty}}$ that can fit the training data perfectly with a linear head ${v_{\mathsf{ft}}^{\infty}}$.

The non-degeneracy condition is similar to our previous results, and holds with probability 1 if the ID subspace is chosen randomly, from Lemma A.16. Formally, let ${R_{\mathsf{aug}}}$ be a $k+1$ dimensional subspace spanning $R_{0}\cup\{w_{\star}\}$, where we recall that $R_{0}=\mbox{rowspace}(B_{0})$. Then we just require that the ID subspace $S$ and ${R_{\mathsf{aug}}}$ are not orthogonal: $\cos\theta_{\mathsf{max}}(S,{R_{\mathsf{aug}}})\neq 0$. We state the formal proposition below and give a proof in Appendix A.

In the linear overparameterized setting, under the ID subspace assumption (Assumption 3.4), let $R_{0}=\mbox{rowspace}(B_{0})$, and ${R_{\mathsf{aug}}}=\mbox{Span}(\{w_{\star}\}\cup R_{0})$. Suppose $w_{\star}\not\in R_{0}$, $\cos\theta_{\mathsf{max}}(S,{R_{\mathsf{aug}}})\neq 0$, and that fine-tuning converges to a local minimum of its loss, then fine-tuning does better ID almost surely: $L_{\mathsf{id}}({v_{\mathsf{ft}}^{\infty}},{B_{\mathsf{ft}}^{\infty}})<L_{\mathsf{id}}({v_{\mathsf{lp}}^{\infty}},B_{0})$ with probability 1 (over the randomness of the training examples).

To summarize, we proved that there are tradeoffs between ID and OOD error: FT has lower ID error but higher OOD error than LP. In the next section, we extend our theoretical insights to illustrate why a simple variant of FT may mitigate such tradeoffs.

4 Linear probing then fine-tuning: a simple variant to mitigate tradeoffs

The advantage of fine-tuning is it can adapt both the feature extractor and head to fit the downstream task. Can we keep this benefit while ensuring that our OOD error is low when we have good pretrained features?

Going back to Theorem 3.3, we see that the alignment error in the head initialization $\varphi^{2}=(v_{0}^{\top}v_{\star})^{2}-(v_{\star}^{\top}v_{\star})^{2}$ plays an important role. The issue with FT was that under random or zero initialization, $\varphi^{2}$ is usually large and since the gradient updates to the feature extractor parameter are coupled with that of the head parameter, the features get distorted in a manner that increases the OOD error. This suggests that we should use a better head initialization—one obtained from linear probing. If the pretrained features are decent, a linear probed head would be much better aligned with $v_{\star}$ allowing the features to be updated in a manner that does not increase the OOD error much.

We formally prove this intuition in a simple setting where we have perfect pretrained features. Of course, if we have perfect pretrained features, linear probing alone gets zero OOD error—so Proposition 3.7 is just a first cut result to illustrate that if initialized well, full fine-tuning does not distort features.

Suppose we have perfect pretrained features $B_{0}=UB_{\star}$ for some rotation $U$. Let $R_{0}=\mbox{rowspace}(B_{0})$. Under the non-degeneracy conditions $\cos\theta_{\mathsf{max}}(R_{0},S)\neq 0,\cos\theta_{\mathsf{max}}(R_{0},S^{\perp})\neq 0$:

The case where we do not have perfect features ($d(B_{0},B_{\star})>0$) is challenging to analyze because except in very special cases, there is no closed form for the fine-tuning iterates $({v_{\mathsf{ft}}}(t),{B_{\mathsf{ft}}}(t))$. Our proof of Theorem 3.3 leveraged invariants to show a lower bound on the error of fine-tuning when $v_{0}$ and $v_{\star}$ are different, but we were not able to show an upper bound.

Experiments

We run experiments on ten distribution shifts to see if our theoretical predictions on the relative performance of linear probing (LP), fine-tuning (FT), and LP-FT, generalize to deep neural networks on real datasets. As expected, given good pretrained features, fine-tuning (FT) does better ID but worse on large OOD shifts than linear probing (LP). In particular, ID and OOD accuracy are not correlated, unlike Recht et al. (2018) but like Xie et al. (2021a). As predicted by the theory, we find that LP-FT does better than both methods ID and OOD and gets around this tradeoff. Our theory also predicts that the reason for these trends is that fine-tuning distorts features, and we see that this distortion indeed happens in practice. For more details on datasets, pretraining models, and experiment protocols, see Appendix B. The datasets we use are:

DomainNet (Peng et al., 2019) is a standard domain adaptation dataset. Here, our ID dataset contains “sketch” images (e.g., drawings of apples, elephants, etc), and the OOD dataset contains “real”, “clipart”, and “painting” images of the same categories. We use the version of the dataset from Tan et al. (2020).

Living-17 and Entity-30 are sub-population shift datasets from the BREEDS benchmark (Santurkar et al., 2020). In Living-17 the goal is to classify an image as one of 17 animal categories such as “bear”—for example, the ID dataset contains images of black bears and sloth bears and the OOD dataset has images of brown bears and polar bears. In Entity-30 the goal is to classify an image as one of 30 entities such as “fruit” or “insect”.

FMoW Geo-shift is adapted from the satellite remote sensing dataset Functional Map of the World (Christie et al., 2018; Koh et al., 2021). The goal is to classify a satellite image into one of 62 categories such as “impoverished settlement” or “hospital”. Our ID dataset contains images from North America, and the OOD dataset contains images from Africa and Europe.

CIFAR-10 $\to$ STL is a standard domain adaptation dataset (French et al., 2018), where the ID is CIFAR-10 (Krizhevsky, 2009), and the OOD is STL (Coates et al., 2011). The task is to classify an image into one of 10 categories such as “dog”, “cat”, or “airplane”—as usual, we remove the “monkey” class in STL since CIFAR-10 has no “monkey” images.

CIFAR-10 $\rightarrow$ CIFAR-10.1 (Recht et al., 2018) is a dataset collected using a very similar protocol to CIFAR-10, and the authors describe it as “a minute distributional shift”. The hope is that a classifier trained on CIFAR-10 gets high accuracy on CIFAR-10.1.

ImageNet-1K (Russakovsky et al., 2015) is a large scale dataset containing over a million images, where the goal is to classify an image into one of 1000 categories such as “Yorkshire terrier”, “Labrador retriever”, “acoustic guitar”, “library”, “school bus”, etc. We fine-tune on ImageNet as the ID dataset, and evaluate on four standard OOD datasets: ImageNetV2 (Recht et al., 2019), ImageNet-R (Hendrycks et al., 2020), ImageNet-A (Hendrycks et al., 2019b), and ImageNet-Sketch (Wang et al., 2019).

Pretraining and models. We use a CLIP pretrained ViT-B/16 for ImageNet. For the other datasets we use a ResNet-50 architecture and consider a diverse range of pretraining methods and datasets: MoCo-v2 (Chen et al., 2020b), CLIP (Radford et al., 2021), and MoCo-TP (Ayush et al., 2020). In Appendix B, we also show results for a CLIP-ViT-B/16 and more fine-tuning baselines on Living-17.

Results. Fine-tuning does better than linear probing on 5 out of 6 ID datasets (average accuracy of 85.1% for fine-tuning vs. 82.9% for linear probing, see Table 1). This is consistent with prior work and intuitions. However, linear-probing does better on 8 out of 10 OOD datasets (average accuracy of 66.2% for linear probing vs. 59.3% for fine-tuning, see Table 2)—linear probing does better on all datasets except CIFAR-10.1 and ImageNetV2, where the OOD is designed to closely replicate the ID dataset. This matches our theoretical predictions, which says that linear probing does better than fine-tuning when the ID and OOD are very different (and the pretrained features are “good”). Our training datasets vary in size from 20K examples to over a million examples, so linear probing does not appear to perform better than fine-tuning simply because of a small training set.

2 Linear probing then fine-tuning (LP-FT)

Experiment protocols. For LP-FT, we initialize the neural network head using the linear probed solution, and then fine-tune the model. LP-FT and fine-tuning use similar compute because the linear probing step is much faster than fine-tuning. As with fine-tuning, we swept over 6 learning rates, early stopping using ID validation accuracy. For the ImageNet experiments we swept over 3 learning rates, and explicitly ensured that LP-FT and fine-tuning use exactly the same compute (we ran each stage of LP-FT for half as many epochs as we ran vanilla fine-tuning).

Results. We find that LP-FT gets the best accuracy ID (average: 85.7%) and OOD (average: 68.9%). This is true for 5/6 ID and 10/10 OOD datasets—every dataset except FMoW ID, where LP-FT is better than linear probing but worse than fine-tuning. Since the ID accuracy on FMoW is low (56.5%), this could be because the pretrained features are not good.

3 Examining the feature distortion theory

Early stopping does not mitigate feature distortion. One might think that fine-tuning is simply overfitting ID, and so early stopping on OOD data (if it were available) might match linear probing OOD. However, our theory predicts that fine-tuning can do worse OOD (than linear probing) throughout the process of fine-tuning, and not just at the end. To test this, we early stop each fine-tuning method and choose the best learning rate based on OOD test accuracy (OOD data was not used except for this ablation). As expected, fine-tuning does improve a little, but linear probing (average accuracy: 67.1%) is still better than fine-tuning (average accuracy: 61.3%). See Appendix B for per-dataset results.

ID-OOD features get distorted from fine-tuning. The feature distortion theory predicts that fine-tuning changes features for ID examples more than for OOD examples, which is why fitting a head on ID examples performs poorly OOD. To test this, for each example $x$ in Living-17 (results for other datasets are in Appendix B), we took the Euclidean distance of the ResNet-50 features before and after fine-tuning: $\|g_{B}(x)-g_{B_{0}}(x)\|_{2}$. As expected, the average distance for ID examples ($0.0188\pm 0.0001$) is more than for OOD examples ($0.0167\pm 0.0001$). The theory also predicts that LP-FT changes features less than fine-tuning does. As expected, the average distance changed by LP-FT both ID ($0.0011\pm 0.0001$) and OOD ($0.0009\pm 0.0001$) is $20\times$ smaller than for fine-tuning.

Pretrained features must be good, ID-OOD far apart. Our theory gives conditions under which linear probing can do better than fine-tuning OOD. Specifically, we require that the ID distribution $P_{\mathsf{id}}$ and OOD distribution $P_{\mathsf{ood}}$ are quite different, and the pretrained features are good ($B_{0}$ is close to $B_{\star}$)—otherwise fine-tuning can do better OOD by adjusting the feature extractor ID. Here we test that these conditions are essential—when they are violated fine-tuning can do better than linear probing OOD.

Feature quality: We use a checkpoint of MoCo-v1 that got 10% worse accuracy (on ImageNet) and compare linear probing and fine-tuning on Living-17. With worse features, both methods do worse, but fine-tuning (96% ID, 71% OOD) does better than linear probing (92% ID, 66% OOD).

ID $\approx$ OOD: We fine-tune / linear probe on CIFAR-10, and test on CIFAR-10.1, a dataset collected using a similar protocol to CIFAR-10. As expected, fine-tuning (92.3%) outperforms linear probing OOD (82.7%). Even in this case, where we have no tradeoffs, LP-FT does the best (93.5%).

Related work and discussion

Fine-tuning vs. linear probing. Fine-tuning (FT) and linear probing (LP) are popular transfer learning algorithms. There is substantial evidence of FT outperforming LP in-distribution (ID) including recent large-scale investigations (Kornblith et al., 2019; Chen et al., 2021a; Zhai et al., 2020; Chen et al., 2020b) (the only notable exception is in Peters et al. (2019) where LP performs better than FT when using ELMo representations, but worse using BERT). This is not intended to be a comprehensive list. There is a large body of past work across different domains that have reported a similar observation. FT is therefore the method of choice for improving accuracy, while LP is used to analyze properties of representations (Peters et al., 2018; Belinkov et al., 2017; Hewitt & Manning, 2019). In our work, we find that FT can underperform LP especially when using high quality pretrained features in the presence of a large distribution shift. There are a variety of other fine-tuning heuristics (Ge & Yu, 2017; Guo et al., 2019; Zhang et al., 2020; Zhu et al., 2020; Jiang et al., 2021; Aghajanyan et al., 2021)—combining our insights with these ideas might lead to better methods.

The benefit of preserving pretrained features. Our work adds to growing evidence that lightweight fine-tuning, where only a small part of a pretrained model are updated, performs better under distribution shifts—and we give a theoretical grounding to why this might be the case. Zero-shot language prompting in vision (Radford et al., 2021) and other lightweight fine-tuning approaches in NLP (Houlsby et al., 2019; Li & Liang, 2021; Xie et al., 2021b; Lester et al., 2021; Utama et al., 2021; Zhou et al., 2021) have been shown to improve OOD performance. In independent and concurrent work, Andreassen et al. (2021) observe that through the course of fine-tuning, ID accuracy continues to increase but OOD accuracy plateaus. Our work shows something stronger: at no point in the fine-tuning process does FT outperform LP.

Mitigating ID-OOD tradeoffs. While LP-FT has sometimes been used as a fine-tuning heuristic (Levine et al., 2016; Kanavati & Tsuneki, 2021; fastai, ), it has not been used for robustness / OOD accuracy, and we show that it addresses the ID-OOD tradeoff theoretically and empirically. Tradeoffs between ID and OOD accuracy are widely studied and prior work self-trains on large amounts of unlabeled data to mitigate such tradeoffs (Raghunathan et al., 2020; Xie et al., 2021a; Khani & Liang, 2021). In contrast, LP-FT uses no extra unlabeled data and is a simple variant of fine-tuning. In concurrent and independent work, Wortsman et al. (2021) show that ensembling the weights of a zero-shot and fine-tuned model mitigates the ID-OOD tradeoff between these approaches, and this method could be promising for our datasets as well.

Theoretical analysis of transfer learning. Prior works on transfer learning mainly analyze linear probing (Wu et al., 2020; Tripuraneni et al., 2020; Du et al., 2020). In recent work, (Chua et al., 2021) study regularized fine-tuning in an underparameterized regime where there is a unique global optimum. In contrast, our analysis studies the overparameterized regime (mirroring modern settings of zero train loss) where we need to analyze the trajectory of fine-tuning from the pretrained initialization because there is no unique optimizer of the objective function. Prior works also focus on ID error, while we analyze OOD error. See Section C for additional related work on theory of overparameterized models.

Conclusion.

There is a strong trend towards leveraging pretrained models to improve downstream performance, and whenever feasible, it is common to fine-tune all model parameters. In this work, we show theoretically and empirically that preserving features might be important for robustness, and simpler approaches like linear-probing can improve out-of-distribution (OOD) performance. This OOD gap between fine-tuning and linear probing grows as the quality of pretrained features improve, so we believe our results are likely to gain significance over time with growing innovations and scale of pretraining.

Theoretical understanding of modern deep learning remains limited, especially the effect of pretraining and transfer learning. In addition to our specific results on fine-tuning, our work introduces some tools and ideas for dealing with the main challenge of characterizing properties of the trajectory from a specific initialization in the presence of multiple global optima (implicit regularization effect of initialization). There are several open questions and extensions such as dealing with non-linear activations, different layerwise learning rates, and the effect of explicit regularization. We found that LP-FT outperforms explicit regularization and using a higher learning rate for the linear layer on Living-17 (Appendix B.4), but a more extensive theoretical and empirical study on this is important.

Finally, we showed LP-FT can mitigate tradeoffs between ID and OOD accuracy in our context. LP-FT could be useful in other situations, for example in CLIP we could initialize the final layer with the zero-shot classifier and then fine-tune the entire model, as done in concurrent work (Wortsman et al., 2021). LP-FT is just a first step in leveraging the intuition from our theoretical analysis and we hope that this work inspires new methods of leveraging powerful pretrained models.

Proofs and Reproducibility: We include proofs for our theoretical results in Appendix A and additional experiment details in Appendix B.

Acknowledgements: We would like to thank Kumar Ayush and Burak Uzkent for MoCo checkpoints pretrained on unlabeled FMoW images, Nilesh Tripuraneni for clarifications on his work and references on principal angles, Daniel Levy for useful suggestions on experiments to run, Niladri Chatterji, Jeff Z. HaoChen, and Colin Wei for useful papers and comments on figures, Niladri Chatterji and Kaidi Cao for reviewing the paper at ML paper swap, Kevin Yang for his help with analyzing differential equations, Tri Dao and Pang Wei Koh for help with writing, Suriya Gunasekar, Adam Kalai, Simon Kornblith, Ting Chen, Sang Michael Xie, Albert Gu, and Kendrick Shen for useful discussions, and Pang Wei Koh, Niladri Chatterji, and Tri Dao for suggestions on framing our results better.

Ananya Kumar was supported by the Rambus Corporation Stanford Graduate Fellowship. Percy Liang was supported by the Open Philantropy Project and NSF Award Grant No. 1805310. Aditi Raghunathan was supported by a Google PhD Fellowship and Open Philanthropy Project AI Fellowship. Tengyu Ma acknowledges support of a Google Faculty Award, NSF IIS 2045685, the Sloan Fellowship, JD.com, SAIL, and SDSI.

References

Appendix A Proofs for Section 3

Big-Oh Notation: For convenience, we use big-oh notation in a way that differs from standard theoretical computer science texts. When we say $O(\mbox{<expr1>})$ we mean that this can be replaced by $c\mbox{ <expr1>}$ for some universal constant such that the statement holds. As an example, we can say $5x^{2}\leq O(x^{2})$ because there exists some universal constant ($c=5$) such that $5x^{2}\leq 5x^{2}$. More examples: we can also say $5x^{2}\geq O(x^{2})$ or if $x\geq 1$ then $7x^{2}\leq O(x^{3})$ and $0.1x^{2}\geq O(x)$.

Working with minimum singular values requires more care than maximum singular vectors. In particular, when we have rectangular matrices some bounds depend on whether the matrix is ‘fat’ (has more columns than rows) or ‘tall’ (has more rows than columns).

Given a matrix $A$, the operator norm $\|A\|_{2}$ is the maximum singular value: $\|A\|_{2}=\sigma_{\max}(A)$.

From this we can easily check that $\Pi_{R}^{2}=\Pi_{R}$ and $\Pi_{R}^{\top}=\Pi_{R}$. See e.g., Chapter 2.5.1 Golub & Loan (2013) for more information.

Principal Angles: Given two non-zero vectors $x$ and $y$, the cosine of the angle between them, $\cos{\theta}$, is:

If we consider the 1-dimensional subspaces (so basically lines) $S_{x}$ and $S_{y}$ spanned by $x$ and $y$ respectively, then the angle between them, $\cos{\theta^{\prime}}$ is given by the absolute value (since lines are undirected):

Principal angles generalize this notion to higher dimensions. See e.g., Chapter 6.4.3 in Golub & Loan (2013) for more information on principal angles.

The directions of the inequalities swap when we take the cosine of the principal angles:

where $\sigma_{i}$ denotes the $i$-th largest singular value. In this paper, we are interested in the cosine of the largest angle between them, given by:

We can massage this into a variational characterization of the maximum principal angle, which is important for lower bounding the error of fine-tuning outside the span of the training data.

The result now follows from some algebra:

A.2 Feature distortion theorem

We first prove our core theorem, that fine-tuning distorts pretrained features.

In the overparameterized linear setting, let $S^{\perp}=\mbox{rowspace}(X)^{\perp}$, $R_{0}=\mbox{rowspace}(B_{0})$, and $v_{\star},B_{\star}$ be the optimal parameters with $w_{\star}=B_{\star}v_{\star}$. If $\cos\theta_{\mathsf{max}}(R_{0},S^{\perp})>0$, then for all time steps $t$, the OOD error of the fine-tuning iterates $({B_{\mathsf{ft}}}(t),{v_{\mathsf{ft}}}(t))$ is lower bounded:

where $\varphi^{2}=\lvert(v_{0}^{\top}v_{\star})^{2}-(v_{\star}^{\top}v_{\star})^{2}\rvert$ is defined to be inital head alignment error and $\epsilon\geq d(B_{0},B_{\star})$ is the error in the pretrained feature extractor.

We follow the sketch in the main paper. We begin with a few lemmas, showing that certain quantities are preserved throughout the fine-tuning process.

Our first lemma says that the representations ${B_{ft}^{t}}x$ do not change for examples perpendicular the span of the training examples. Note that the final output ${v_{ft}^{t}}^{\top}{B_{ft}^{t}}x$ still changes, because ${v_{ft}^{t}}$ changes.

For all times $t$ and all $x\in S^{\perp}$, we have:

We initialized fine-tuning with the feature extractor ${B_{\mathsf{ft}}}(0)=B_{0}$. It suffices to show that $\partial_{t}{B_{ft}^{t}}x=0$ for all $x\in S^{\perp}$. Recall that $\partial_{t}{B_{ft}^{t}}$ is given by the gradient flow update equation:

Computing the RHS explicitly using multivariable chain rule, we get:

But $Xx=0$ for $x\in S^{\perp}$, since $x\in S^{\perp}$ is defined as $x$ is perpendicular to the rowspace of $X$ (i.e., perpendicular to the rows of $X$). So the RHS is —that is, $\partial_{t}{B_{ft}^{t}}x=0$, as desired. ∎

Next, we show that the change in the head and feature extractor are ‘coupled’. So if the head changes in a certain way, then the feature extractor cannot just stay the same. In the literature, this is sometimes called the “balancedness" lemma, and has been proved in prior work on two layer linear networks.

This follows by showing that the derivative is :

Which can be verified by direct calculation. See Theorem 2.2 in Du et al. (2018) and the proof of Theorem 1 in Arora et al. (2018). ∎

For our proof we will require that every feature $r\in R$ can be generated from some OOD direction, that is $r=B_{0}u$ for some $u\in S^{\perp}$. We will show that this is implied by the condition on the principal angle: $\cos\theta_{\mathsf{max}}(R,S^{\perp})>0$ where $R=\mbox{rowspace}(B_{0})$, which we assumed in Theorem 3.3. The following lemma shows this (and also quantifies that the norm of $u$ does not shrink too much when projected onto $R$).

Let $c=\cos\theta_{\mathsf{max}}(R,S)$. Firt, we get rid of an easy case—if $c=0$, then we need to show the claim for all $r\in R$ with $\|r\|_{2}=c=0$, which means $r=0$. Then we can just pick $s=0$, and $\Pi_{R}(s)=0=r$ and $\|s\|_{2}=0\leq 1$. So for the rest of the proof we assume $c>0$.

Step 2: Bounding norm of $s$: It suffices to show that $\|s\|_{2}\leq 1$. Since $F$ has orthonormal columns, $\|s\|_{2}=\|Fw\|_{2}=\|w\|_{2}$, so it suffices to show that $\|w\|_{2}\leq 1$. Since $E$ has orthonormal columns, $\|r\|_{2}=\|z\|_{2}$. Recall that $z=E^{\top}Fw$—since $w\in\mbox{rowspace}(E^{\top}F)$, from Lemma A.6 we have:

Rearranging, we get $\|w\|_{2}\leq\|z\|_{2}/c=1$, as desired.

In the lemma above, we used a standard linear algebraic result that we include for completeness. This says that $A$ cannot shrink vectors in its rowspace too much, where the shrinkage factor is given by the minimum singular value of $A$.

We bound the norm of $x$ using the SVD. Consider the singular value decomposition (SVD) of $A$:

Where for the fourth step, we used the fact that if $x\in\text{rowspace}(V^{\top})$ and the rows of $V^{\top}$ are orthonormal, then $\|V^{\top}x\|_{2}=\|x\|_{2}$. One way to see this is by writing $x=\sum_{i}\alpha_{i}v_{i}$, where $v_{i}$ are rows of $V^{\top}$, and then noting that $V^{\top}x=(\alpha_{1},\ldots,\alpha_{r})$ and so $x$ and $V^{\top}x$ have the same norm. ∎

We now prove Theorem 3.3, following the 3 steps outlined in the main text.

Let $c=\cos\theta_{\mathsf{max}}(R,S^{\perp})$. From Lemma A.7, we have $L_{\mathsf{ood}}({v_{ft}^{t}},{B_{ft}^{t}})\leq\sigma_{\min}(\Sigma)\|B_{\star}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}{v_{ft}^{t}}\|_{2}^{2}$ so it suffices to bound $\|B_{\star}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}{v_{ft}^{t}}\|_{2}$.

Because it makes the proof much easier, we will prove the contrapositive, and then convert back to the original theorem statement. We assume $\|B_{\star}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}{v_{ft}^{t}}\|_{2}\leq\Delta$, and will show that:

Where $g_{1}$ and $g_{2}$ are non-negative polynomials we will bound in the proof.

We gave a basic outline of the proof in the main paper, and here we are just trying to be careful about capturing all the dependencies. We also give intuition for each step before diving into algebra (which we include for completeness).

Recall that in the overparameterized linear setting we assumed we have orthonormal $B_{0}$ with $\|B_{0}-UB_{\star}\|_{2}\leq\epsilon$ for some $U$. We note that the setup is rotationally symmetric so without loss of generality we can suppose $\|B_{0}-B_{\star}\|_{2}\leq\epsilon$. This is because we can let $B_{\star}^{\prime}=UB_{\star}$ and $v_{\star}^{\prime}=Uv_{\star}$, and we have $w_{\star}=B_{\star}^{\top}v_{\star}=(UB_{\star})^{\top}(Uv_{\star})$, where $w_{\star}$ is the optimal classifier—so we can now write the entire proof in terms of $B_{\star}^{\prime}$ and $v_{\star}^{\prime}$.

Step 1: Show that $\|{v_{ft}^{t}}-v_{\star}\|_{2}\leq\Delta/c$: We first give intuition and then dive into the math. The key insight is to use the fact that in ‘many’ directions ${B_{ft}^{t}}$ and $B_{0}$ are the same (formally, for all $x\in S^{\perp}$, ${B_{ft}^{t}}x=B_{0}x$). But $B_{0}$ and $B_{\star}$ are close by assumption, which means that ${B_{ft}^{t}}$ and $B_{\star}$ are close in ‘many’ directions. Then since we assumed in the contrapositive that ${v_{ft}^{t}}^{\top}{B_{ft}^{t}}$ and $v_{\star}^{\top}B_{\star}$ are close, we get that ${v_{ft}^{t}}$ and $v_{\star}$ are close in ‘many’ directions. Because $S^{\perp}$ covers the rowspace of $B_{0}$, we get that ‘many’ is $k$, which is precisely the dimensionality of $v_{\star}$, so the two vectors ${v_{ft}^{t}}$ and $v_{\star}$ must be close.

We now dive into the math. Since $B_{0}$ has orthogonal rows, $B_{0}$ has full column rank.

We note that $\|z\|_{2}=c$. Then, we can find $y\in R=\text{rowspace}(B_{0})$ such that $B_{0}y=z$ (since $B_{0}$ has full column-rank) and then $\|y\|_{2}=\|z\|_{2}=c$ (since $B_{0}$ has orthonormal rows).

Since $c=\cos\theta_{\mathsf{max}}(R,S^{\perp})>0$, and $y\in R$ with $\|y\|=c$, from Lemma A.5 we can choose $x\in S^{\perp}$ with $\|x\|_{2}\leq 1$ and $\Pi_{R}(x)=y$. Then, we have $B_{0}x=z$.

From Proposition A.3, since $x\in S^{\perp}$, $B_{0}$ does not change in directions of $x$ when fine-tuning so we have: $B_{0}x={B_{ft}^{t}}x$.

The claim now follows from simple algebraic manipulation, following the intuition we described. The algebra just captures what ‘close’ means and adds up the error terms.

Which shows that $\|{v_{ft}^{t}}-v_{\star}\|_{2}\leq(\Delta+\epsilon\|v_{\star}\|_{2})/c$.

Step 2A: Show that $\|{B_{ft}^{t}}\|_{F}^{2}$ is small: The key insight is to take the trace on both sides of Proposition A.4, which bounds the Frobenius norm of ${B_{ft}^{t}}$ and therefore the operator norm.

Squares are non-negative, so we get the inequality:

Step 2B: Show that $\|B_{0}^{\top}v_{\star}\|_{2}^{2}-\|{B_{ft}^{t}}^{\top}v_{\star}\|_{2}^{2}$ is small: This step doesn’t involve much insight, and is standard peturbation analysis—we simply factor the difference of squares and bound each term.

First, we bound $\|{B_{ft}^{t}}^{\top}{v_{ft}^{t}}-{B_{ft}^{t}}^{\top}v_{\star}\|_{2}$:

Next, we bound $\|B_{0}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}v_{\star}\|_{2}$:

Finally, we bound $\lvert\|B_{0}^{\top}v_{\star}\|_{2}^{2}-\|{B_{ft}^{t}}^{\top}v_{\star}\|_{2}^{2}\rvert$, using the identity:

Step 3: Use Proposition A.4 to show $v_{0}$ and $v_{\star}$ must be close: The key insight is that we start from Proposition A.4, and left and right multiply by $v_{\star}$, after that we use the previous steps and do some some standard perturbation analysis.

The key step is to left multiply both sides by $v_{\star}^{\top}$ and right multiply both sides by $v_{\star}$ to get:

Rearranging, and then using Equation A.66, we get:

This is close to what we want, except we have $({v_{ft}^{t}}^{\top}v_{\star})^{2}$ on the LHS instead of $(v_{\star}^{\top}v_{\star})^{2}$. We previously showed that ${v_{ft}^{t}}$ and $v_{\star}$ are close, in Step 1, so with some algebra we can bound the difference between $({v_{ft}^{t}}^{\top}v_{\star})^{2}$ and $(v_{\star}^{\top}v_{\star})^{2}$:

Above, from the third line to the fourth line, we used triangle inequality and Cauchy-Schwarz.

So finally, by triangle-inequality we can now bound $\lvert(v_{\star}^{\top}v_{\star})^{2}-(v_{0}^{\top}v_{\star})^{2}\rvert$:

Wrap up i.e., writing out $\Delta_{4}+\Delta_{3}$ explicitly: This is basically the bound we want, but we would like to express $\Delta_{3},\Delta_{4}$ in terms of $\Delta$ and $\epsilon$. Note that this step has no insight, and is just algebra—we include the details for reference and verifiability. We recall:

Since $B_{0}$ has orthogonal rows (by assumption), $B_{0}^{\top}$ has orthogonal columns, so $\|w_{\star}\|_{2}=\|B_{0}^{\top}v_{\star}\|_{2}=\|v_{\star}\|_{2}$. In addition, since $B_{0}$ has $k$ orthogonal rows, $\|B_{0}\|_{F}=\sqrt{k}$. We also note that $\sqrt{\|B_{0}\|_{F}^{2}+\|v_{\star}\|_{2}^{2}}\leq\|B_{0}\|_{F}+\|v_{\star}\|_{2}=\sqrt{k}+\|w_{\star}\|_{2}$. Since $c\leq 1$, we have:

So for $\Delta_{2}$, up to constant factors we can ignore the $\epsilon\|v_{\star}\|_{2}+\Delta$ term—this means we get:

Using the fact that $\sqrt{k}+\|w_{\star}\|_{2}\leq\sqrt{k}(1+\|w_{\star}\|)$ we get:

Then since $\Delta+\epsilon\|w_{\star}\|_{2}\leq(1+\|w_{\star}\|_{2})(\Delta+\epsilon)$, we get:

Now for $\Delta_{3}$, first note that $\epsilon\leq 2$, since $B_{\star}$ and $B_{0}$ have orthogonormal rows so $\|B_{\star}-B_{0}\|_{2}\leq 2$. This means that $\epsilon\|w_{\star}\|_{2}\leq\|w_{\star}\|_{2}$, so $\Delta_{3}$ simplifies to:

Substituting the bound for $\Delta_{2}$ into $\Delta_{3}$, we get:

Since $\Delta/c\leq(\Delta+\epsilon)/c$ and $\|w_{\star}\|_{2}^{2}\leq(1+\|w_{\star}\|_{2})^{2}$ we have for the final error $\Delta_{3}+\Delta_{4}$:

Wrap up i.e., taking the contrapositive: So we’ve shown that if $\|B_{\star}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}{v_{ft}^{t}}\|_{2}^{2}\leq\Delta$, then:

We’d like to flip this around: suppose $\lvert(v_{\star}^{\top}v_{\star})^{2}-(v_{0}^{\top}v_{\star})^{2}\rvert\geq\varphi^{2}$ for some $\varphi$. To lower bound $\|B_{\star}^{\top}v_{\star}-{B_{ft}^{t}}^{\top}{v_{ft}^{t}}\|_{2}^{2}$, we simply take the contrapositive of what we have proved. Let $\Delta$ be given by:

In this case with some algebra, we can show that:

To see this, we bound each of the terms in the RHS separately using our definition of $\Delta$. Then, from the contrapositive of what we proved (compare with Equation A.92, we get:

Finally, we can massage $\Delta$ to combine terms and make it look slightly nicer:

Then applying Lemma A.7 we get the desired result. For even more interpretability, if $\|w\|_{2}=1$ and $\varphi$ is bounded above by some constant, then you can think of $\Delta$ as approximately $\frac{c}{\sqrt{k}}\varphi^{2}-\epsilon$. This completes the proof. ∎

A.3 LP vs. FT (OOD)

We now prove Theorem 3.5, which compares linear probing and fine-tuning in the linear overparameterized setting, when the ID data lies in a lower dimensional subspace.

We first state a more precise version of Theorem 3.5—basically we fix all problem parameters except $B_{0}$ (which limits to $B_{\star}$). To define the limit, we consider a sequence of pretrained feature extractors: $\{B_{0}^{i}\}_{i=1}^{\infty}$. We define the corresponding limit points of fine-tuning and linear probing when we start from the $i$-th pretrained feature extractor. That is, let ${v_{\mathsf{ft}}}^{i}(t),{B_{\mathsf{ft}}}^{i}(t)$ denote the parameters at time $t$ of fine-tuning if we initialize with $v_{0},B_{0}^{i}$ (see Equation 3.2 for the fine-tuning updates). Let ${v_{\mathsf{lp}}^{\infty}}^{i},B_{0}^{i}$ be the linear probing solution when initialized with $v_{0},B_{0}^{i}$ (see Equation 3.5 for the linear probing updates). We note that the LP iterates converge to ${v_{\mathsf{lp}}^{\infty}}^{i},B_{0}^{i}$ as a result of gradient flow on a convex problem.

Finally, Theorem 3.5 says that as the pretrained representations get better, linear probing does much better than fine-tuning OOD:

In the linear overparameterized setting, under the ID subspace assumption, fix the dimensions of the setting $d,k,m$, number of examples $n$, the ID subspace $S$, ID distribution $P_{\mathsf{id}}$, the distribution over the head $v_{0}$, and the ground truth parameters $v_{\star},B_{\star}$. Assume the non-degeneracy conditions $\cos\theta_{\mathsf{max}}(R_{*},S)>0$ and $\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})>0$ where $R_{*}=\mbox{rowspace}(B_{\star})$. Given a sequence of pretrained feature extractors $\{B_{0}^{i}\}_{i=1}^{\infty}$ with $B_{0}^{i}\to B_{\star}$, where the limit is in the pseudometric given by Definition 3.1, the ratio of OOD errors of linear probing and fine-tuning converges in probability to :

The purpose of the infimum is to capture the fact that the bound holds for all times $t$ for fine-tuning (and therefore also for the limit ${v_{\mathsf{ft}}^{\infty}},{B_{\mathsf{ft}}^{\infty}}$ when it exists). Note that the ratio is a random variable because the training data is sampled from $P_{\mathsf{id}}$ and the head is sampled ($v_{0}\sim\mathcal{N}(0,\sigma^{2}I)$ for some $\sigma^{2}$).

Recall that we say a sequence of real-valued random variables converges in probability to (written as $X_{i}\overset{p}{\to}0$) if for every $\epsilon^{\prime},\delta>0$, for all large enough $i$ (that is, for all $i\geq N_{i}$ for some $N_{i}$), we have:

Accordingly, fix arbitrary $\epsilon^{\prime},\delta>0$, and we will show that the ratio of errors is eventually smaller than $\epsilon^{\prime}$ with probability at least $1-\delta$.

Lower bounding fine-tuning error: Since $B_{0}^{i}\to B_{\star}$, from Lemma A.10 we have that $\cos\theta_{\mathsf{max}}(R^{i},S^{\perp})\to\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})$ where $R^{i}=\mbox{rowspace}(B_{0}^{i})$. Since $\cos\theta_{\mathsf{max}}(R_{*},S^{\perp})>0$, this means that for all large enough $i$ we have:

Next, from Lemma A.12, we have that with probability at least $1-\delta/2$, $\mbox{Head-Error}(v_{0},v_{\star})=\lvert(v_{0}^{\top}v_{\star})^{2}-(v_{\star}^{\top}v_{\star})^{2}\rvert\geq c_{\delta}$ for some $c_{\delta}>0$. Plugging this into the fine-tuning bound in Theorem 3.3, this means that for all large enough $i$ with probability at least $1-\delta/2$:

for some $c_{\delta}^{\prime}>0$. But since $B_{0}^{i}\to B_{\star}$ we have $d(B_{0}^{i},B_{\star})\to 0$ as $i\to\infty$. So this means that for all large enough $i$ with probability at least $1-\delta/2$:

Upper bounding the linear probing error: Since $B_{0}^{i}\to B_{\star}$, from Lemma A.10 we have that $\cos\theta_{\mathsf{max}}(R^{i},S)\to\cos\theta_{\mathsf{max}}(R_{*},S)$ and so since $\cos\theta_{\mathsf{max}}(R_{*},S)>0$, for all large enough $i$ we have:

Plugging this into the RHS of Lemma A.14, Equation A.132, which upper bounds the OOD error of linear probing, we get that for all large enough $i$, with probability at least $1-\delta/2$:

for some $u_{\delta}>0$. Again since $d(B_{0}^{i},B_{\star})\to 0$ as $i\to\infty$, this means for all large enough $i$, with probability at least $1-\delta/2$, $d(B_{0}^{i},B_{\star})$ will be small enough so that:

Taking the ratio: So taking the ratio of the lower bound for fine-tuning, and upper bound for linear probing, we get with with probability at least $1-\delta$:

We now prove the Lemmas that we used in the above proof.

Theorem 3.5 assumes conditions on the angle between the perfect feature extractor $B_{\star}$ and the ID subspace $S$. However, fine-tuning and linear probing start from features $B_{0}$ with some error, and do not get access to $B_{\star}$. We show that if $B_{0}$ and $B_{\star}$ are close, then the angles between their rowspaces to a third subspace $T$ (which could be the the ID subspace $S$) is similar.

Here in the first step we used the definition of $\cos\theta_{\mathsf{max}}$ (Definition 3.2), and the fact that $B_{0}^{\top}$ has orthonormal columns which form a basis for $R_{0}$ (the rowspace of $B_{0}$), so in Definition 3.2 we can subtitute $E=B_{0}^{\top}$. To get Equation A.110 we used Weyl’s theorem, which bounds the singular value under perturbations: $\sigma_{r}(A+B)\geq\sigma_{r}(A)-\sigma_{1}(B)$. To get Equation A.111 we used the fact that $\|Fv\|_{2}=\|v\|$ since $F$ has orthonormal columns.

Since this holds for all rotation matrices $U$, we can take the minimum over $U$ to get:

Since the relationship between $B_{0}$ and $B_{\star}$ are symmetric (and the distance $d$ is symmetric), this gives us the desired result:

This follows directly from Lemma A.9. $B_{0}^{i}\to B_{\star}$ means $d(B_{0}^{i},B_{\star})\to 0$. Then from Lemma A.9:

This means $\cos\theta_{\mathsf{max}}(R^{i},T)\to\cos\theta_{\mathsf{max}}(R_{*},T)$ as $i\to\infty$ ∎

A.3.2 Bounding the head error

We prove a lower bound on $\mbox{Head-Error}(v_{0},v_{\star})=\lvert(v_{0}^{\top}v_{\star})^{2}-(v_{\star}^{\top}v_{\star})^{2}\rvert$, which was a key term in the fine-tuning lower bound (Theorem 3.3). Note that if the head is initialized as $v_{0}=0$, then $\mbox{Head-Error}(v_{0},v_{\star})=\|v_{\star}\|_{2}^{2}=\|w_{\star}\|_{2}^{2}$. In practice, the head is usually initialized randomly, for example normally distributed. Intuitively, the head error is still high because we do not know which direction the head is pointing in, so most of the time the initial (randomly sampled) head will be pointing in the wrong direction. If $v_{0}\sim N(0,\sigma^{2}I)$ can show that for any $\sigma^{2}$, the head error will still typically be at least $\Omega(\|v_{\star}\|_{2})$ This is an illustrative result, one can show similar results for other random initializations as well.

We first prove an anti-concentration lemma, which says that if $u$ is univariate Gaussian, then it cannot be too close to any particular constant $a$, no matter how the variance of the Gaussian is chosen.

For some universal constant $c$, given $a>0$, for all $\nu^{2}$ if $u\sim N(0,\nu^{2})$ then for all $0\leq\delta\leq 1$:

Consider $\delta$ such that $\delta\leq 1/10$. Then for all $u$ with $\lvert u-a\rvert\leq\delta a$, we have $u\geq 9a/10$. For all $u\geq 9a/10$, the density $f(u)$ is upper bounded (from the formula for the density of a Gaussian random variable) by:

We can maximize this explicitly (e.g., use Mathematica or by taking the logarithm and then setting the derivative to 0) and we get for some universal constant $c^{\prime}\geq 10$ (it is OK to choose a larger universal constant than needed):

Since the density is less than $c^{\prime}/a$ and if $\lvert u-a\rvert\leq\delta a$ the size of the interval is $2\delta a$, we get for all $\delta\leq 1/10$:

Now, we substitute $\delta^{\prime}=2c^{\prime}\delta$. We get for all $\delta^{\prime}\leq 2c^{\prime}/10$:

Since $c^{\prime}\geq 10$, $2c^{\prime}/10\geq 1$, so the statement is true for all $0\leq\delta^{\prime}\leq 1$. ∎

We now bound the error in the head if the initialization is Gaussian. This bound holds for all initialization variances $\sigma^{2}$. Similar bounds can be shown for other (non-Gaussian) head initializations using similar anti-concentration arguments.

First note that $\mbox{Head-Error}(v_{0},v_{\star})=\mbox{Head-Error}(-v_{0},v_{\star})$ and $v_{0}$ is symmetric around ($v_{0}$ and $-v_{0}$ have the same probability), and is almost surely not exactly . So without loss of generality, we can suppose that $v_{0}^{\top}v_{\star}\geq 0$.

Suffices to bound $\lvert v_{0}^{\top}v_{\star}-v_{\star}^{\top}v_{\star}\rvert$: We decompose the error:

So we bound $\lvert v_{0}^{\top}v_{\star}-v_{\star}^{\top}v_{\star}\rvert$.

$v_{0}^{\top}v_{\star}$ is normally distributed: We note that $v_{0}^{\top}v_{\star}$ is distributed as:

In other words, a normal with mean , and variance $\sigma_{1}^{2}=\sigma^{2}v_{\star}^{\top}v_{\star}$, and therefore standard deviation $\sigma_{1}=\sigma\sqrt{v_{\star}^{\top}v_{\star}}$.

Apply Gaussian anti-concentration lemma: Then, from Lemma A.11, we have for some universal constant $c$ that with probability at least $1-\delta$:

So substituting this back into Equation A.124, we get the desired result:

A.3.3 Upper bounding linear probing error

We showed a lower bound for the OOD error of fine-tuning in Theorem 3.3. To compare this with linear probing, we prove an upper bound on the OOD error of linear probing.

For completeness we include an elementary lemma (note that the condition that the matrices are tall is important for composing $\sigma_{\min}$, unlike for $\sigma_{\max}$, and we included this lemma to be careful about these conditions):

Suppose we have two matrices $A$, $B$ of shape $(r,s)$ and $(s,t)$ respectively, and they are tall matrices so $r\geq s\geq t$. Then we have:

And $\min_{\|x\|_{2}\leq 1}\|x\|_{2}=1$ which completes the proof. ∎

In the linear overparameterized setting, under the ID subspace assumption, fix arbitrary $P_{z}$. Then there exists $c_{\delta}$ such that with probability at least $1-\delta$, for all $d,n,m,k,w_{\star}$, feature extractors $B_{\star},B_{0}$, and ID subspaces $S$ with corresponding $F$ (whose columns are orthonormal and form a basis for $S$), if $\cos\theta_{\mathsf{max}}(S,R)>0$, we have:

If $P_{z}$ is isotropic Gaussian so $\mathcal{N}(0,I_{m})$, then we derive a bound for $c_{\delta}$ analytically: if $n\geq 5m$ and $n\geq 10\log{\frac{1}{\delta}}$ then with probability at least $1-\delta$, the linear probing OOD error is upper bounded by:

From the ID subspace assumption, the data matrix $X$ of shape $(n,d)$ can be written as $X=ZF^{\top}$ where $Z$ be a matrix of shape $(n,m)$ with each row $Z_{i}$ sampled iid from $P_{z}$, and $F$ is a matrix of shape $(d,m)$ whose columns are orthonormal and form a basis for the ID subspace $S$.

Let $\epsilon=\|B_{\star}-B_{0}\|_{2}\leq$. We first prove the bounds for $\epsilon$, in terms of $d(B_{0},B_{\star})$ and we later handle the fact that the feature extractor distance involves the min over rotation matrices $U$: $d(B_{0},B_{\star})=\min_{U}\|UB_{0}-B_{\star}\|_{2}$.

Bounding key singular values: Before proceeding with the proof, we examine a key quantity $XB_{0}^{\top}=ZF^{\top}B_{0}^{\top}$ which comes up in the Hessian of the loss function. We will show that this is invertible almost surely, and get a lower bound on its min singular value.

First, we examine the shapes of the matrices. $ZF^{\top}B_{0}^{\top}$ has shape $(n,d)$ where $Z$ has shape $(n,m)$ and $F^{\top}B_{0}^{\top}$ has shape $(m,k)$. Since $n\geq m>k$ we have that $Z$ and $F^{\top}B_{0}^{\top}$ are tall matrices, and so from Lemma A.13 we can write the min singular value of $ZF^{\top}B_{0}^{\top}$ as:

Now from the definion of the principal angle (Definition 3.2), we have:

Since we assumed $P_{z}$ has density in the ID subspace assumption, from Lemma 3 in Xie et al. (2021a) we get that for some $c_{\delta}^{\prime}>0$ that depends on $\delta$ and $P_{z}$, with probability at least $1-\delta$: