Features are fate: a theory of transfer learning in high-dimensional regression

Dear Reviewer n1uD,

We are happy to hear that you found the work to be readable and well-written. Thank you for your thoughtful questions– we address them below:

1. We believe that this framework would also describe models trained with other optimization algorithms such as finite size gradient descent or SGD. The challenge lies in describing their dynamics precisely. For this reason, we focused on the analytically tractable setting of gradient flow. However, similar theorems on the global convergence of gradient descent for finite step size are proven in “Global Convergence of Gradient Descent for Deep Linear Residual Networks” (Wu 2019). The main takeaway is that with a sufficiently small learning rate the network converges to a global optimum, in which case our results would hold exactly as written in the paper.
2. We agree that Theorem 2.2 refutes prior work, which is why we find it so surprising! In Fig. 5 (Appendix E), we plot the KL divergence and Wasserstein Metric for the source and task distributions in our theory. If these metrics were predictive of transfer performance, the correlation would be negative, which is not the case. We have also performed similar experiments in 2-layer nonlinear networks. Here we construct two teacher networks with the same feature space but different output weights so that the Wasserstein distance grows, but transfer performance remains constant. We can include this plot in the appendix to address this concern
3. Since the focus of our paper is regression, the extension to classification is outside the scope of the current work. However, it is interesting to consider how the results may generalize. Consider a classification task defined by class conditional distributions that are gaussians with equal covariances and means $m_1 = \vec{0}$, $m_2 = \vec{\beta}$. Then the optimal classifier learns the discriminator vector $\vec{{\beta}}$. In this way, the problem is very closely related to the regression problem we consider. The features should be defined by the two vectors $\vec{\beta_s}$ and $\vec{\beta_t}$, which define the separation between the means of the gaussians in the source and target task respectively. Then the analog of the feature overlap is $\vec{\beta_s}^T \vec{\beta_t}$ as in our current work. We predict that, as in the paper, this overlap will control the transferability. This could be the subject of future work and we thank you for pointing out the extension.

Dear Reviewer n1uD,

Thank you for your insightful question. To clarify how Theorem 2.2 connects to transfer learning, consider a pretrained model that has been trained on the dataset $p_{f}(x,y)$ defined by the function $f(x)$. If we want a predictive metric for transfer learning based on the datasets themselves, it should correlate negatively with transfer performance, i.e. transfer is better for datasets that have smaller distance. Assume that pertaining is successful, in the sense that the trained model learns a representation of the function $f(x)$. Under the mild architecture assumption that the model contains a final linear layer that maps from $\mathbb{R}^{d_{model}}$ to$\mathbb{R}^{d_{labels}}$ , it must be true that the space spanned by the features $\phi_i$ is contained in the space spanned by the activations of the model before the final readout layer. In fact we prove in Theorem 3.4 that for linear networks this is the case and that the space spanned by the final hidden layer activations is equal to the space spanned by pretraining task function. Now consider the entire space of downstream tasks. In our framework this corresponds to the set of $\{p_g(x,y)\}$ for any $g \in L_2$. A subset of tasks for which transfer learning will be successful are those that live in the subspace spanned by the $\phi_i$, since the network only needs to update its final layer weights in order to represent $g(x)$. One can make a Rademacher complexity argument here to argue that with finite data, learning the output weights with the correct features will be more successful that learning new features and output weights. Therefore, the subspace of tasks that we consider in Theorem 2.2 will transfer well. But we prove in the theorem that one can always find a task in this subspace for which these dataset based metrics are arbitrarily large. Therefore, it is always possible, mathematically, to find a counterexample to the notion that distant distributions should not transfer well. We hope that this explanation clarifies any confusion.

Dear Reviewer vks9,

Thank you for your review of our paper. We agree that normally distributed data and linear source and target functions are a major simplification from realistic transfer learning scenarios. However, we chose this simplified setting in order to have an analytically solvable model. We note, however, that the results on gradient flow dynamics still hold for gaussian data with arbitrary covariance, and the generalization error in this setting is the subject of existing work (for example: “A Theory of High Dimensional Regression with Arbitrary Correlations”, (Mel Ganguli 2021)), so this assumption can likely be relaxed. We also explore nonlinear source and target functions in Section 4 when we discuss student teacher ReLU networks. Finally, we direct you to the paper “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution” (Liang 2022) which demonstrates the importance of the learned feature space for out of distribution generalization empirically in a number of realistic settings. Additionally, we answer your questions below:

1. Thank you for pointing this out— the notation is not clear. Equations 1 and 2 should be written with $\epsilon_s$ and $\epsilon_t$ to make it clear that these are independent noise realizations for the source and target tasks. In our proof we consider the label noise to be drawn independently from the same Gaussian distribution in the source and target task, but the extension to Gaussian distributions with different variances is straightforward and our proofs will still hold with only minor adjustments.
2. In Equation 6, $\Theta_i$ is the $i^{th}$ parameter in the full set of parameters
3. Thank you for pointing this out, it is a typo. The subscript $L$ should follow the norm, to denote the Lipschitz norm of the function.
4. Although these theorems are simply extensions of those found in Yun et. al., they are essential components of our argument. We felt that citing them as important building blocks and reproducing the proofs in the main text makes our paper accessible and self-contained for readers who may not be familiar with the literature on deep linear networks. In addition, we extend their results on the convergence of gradient descent to underparameterized networks in Appendix B.4.

2. Thank you for pointing this out. You are correct that in the regime of “anomalous positive transfer” the benefit of pre-training comes solely from having more data in the source task– we will make this point clear in our revised manuscript. While we do make the point that simple data-dependent metrics do not predict transfer, we do not foresee measuring feature space overlap directly as an “alternative” metric, but rather as a theoretical guidepost for future algorithm development. Although we believe that this quantity (defined as the norm of the projection of the target function into the orthogonal complement of the pretrained RKHS) would correlate with transfer performance, it is impossible to calculate before pretraining when the ground truth target function is not known. However, one can estimate this quantity given a pretrained model and samples of the data. For the two models we consider, this quantity can be computed exactly given the functions $f_s$ and $f_t$ (see Appendix C for the procedure in the 2-layer ReLU network). The main idea we are trying to advance is that similarity between tasks should be defined with respect to the features of the model, not with respect to the datasets themselves. We agree that the double descent phenomenon obscures this point a bit. Since the double descent occurs in the scratch-trained model, it could indeed be avoided by regularization. Whether the appropriate scratch-trained model to compare to should include regularization or not is an interesting philosophical question. In our work we adopt the perspective that the correct model to compare is one trained using the same algorithm as in pretraining. In our setting that is unregularized gradient flow, which for deep linear networks will converge to a representation of the minimum $l_2$ interpolator, hence the double descent. Our theoretical techniques to describe gradient flow rely on the absence of weight decay (the conserved quantity Eq. 35 is no longer conserved if we add weight decay) so we do not include it in order to make rigorous statements about the dynamics. However, in our revised draft, we will include a numerical simulation that demonstrates that anomalous positive transfer is eliminated if the scratch network is optimally regularized.

Dear Reviewer YzFe,

Thank you for your thoughtful review of our paper. We are very happy to hear that you felt our manuscript was well written and “logically flawless”. In addition, we thank you for pointing out sections of the work that felt misleading. We address these points in the following comments:

1. We agree that deep linear networks are a particularly simple setting, but believe that much of the phenomenology we rigorously demonstrate in these networks still holds for arbitrary architectures. The overall picture is the following. During pretraining, the network will learn a representation of the ground truth function $f_s(x)$ that generates the labels $y_s$. We can think of the function $f_s$ as a linear combination of basis functions, since it is an element of the vector space $L_2$. We argue that training in the feature learning regime causes the network to learn these (and only these) basis functions to build a representation of $f_s$. We demonstrate that this is the case analytically for deep linear networks in Theorem 3.4, which shows that after training, the features of the model (the activations at layer $L-1$) only span the one dimensional subspace along $\beta_s$. This should be compared with training in the lazy regime, where the function space of the model is fixed at initialization, and the best the model can do is to learn the projection of $f_s$ into the RKHS defined by the NTK. In the deep linear case, this would correspond to fixing the weight matrices to their initial random values and performing regression using features that are essentially a product of $L-1$ random matrices. The fact that the network only represents the features relevant to $f_s$ is the source of negative transfer for downstream tasks. During linear transfer, when only the output weights can be tuned on the target task, the best the network can do is learn the component of $\beta_t$ that lies along $\beta_s$, since its features are fixed to their pretrained values (see Theorem 3.7). You mention that we could instead consider a simple linear regression model and “perform linear evaluation on fresh new weights”. We would greatly appreciate it if you could clarify the experiment you have in mind so that we can effectively address the concern. For arbitrary nonlinear networks, we instead need to think in function space instead of the vector space $\mathbb{R}^d$. The relevant function space is the RKHS defined by the activations at layer $L-1$ after pretraining. We demonstrate that in 2-layer ReLU networks, the pretrained network encodes a basis for this space in Section 4 (see Fig 4a-b). We believe that in deeper nonlinear networks, the picture is the same, but identifying the features becomes more complicated. For this reason, we focus on deep linear networks in the student teacher setting, as the features are readily identifiable and we can prove rigorously that our feature learning hypothesis holds. To address this confusion, we plan to include an “Our Contributions” section to the introduction, where we describe this qualitative picture more clearly. Your observation that depth plays no role in generalization is important and we will comment on it in the main text. In our setting depth does not play a role since the ground truth function is linear, as you point out. In nonlinear networks with more complicated ground truth functions, depth will modify what features of $f_s$ are learnable, but will not fundamentally change the argument. Whatever features are learned define an RKHS which constrains how well the target function can be regressed. Finally, we’ll mention that in the full fine tuning case, the situation is more complicated. We are able to theoretically describe this setting in deep linear networks since these models enjoy a number of favorable analytical properties (see Eq. 35) that allow us to describe gradient flow from an arbitrary initial condition. This kind of analysis is notoriously difficult for an arbitrary network. We are unaware of existing theoretical techniques to describe optimization dynamics in deep nonlinear networks from non-random initializations.

Minor issues/questions

1. M is dimensionality of the subspace of $L_2$. It is defined through the notation {\phi_i\}_{i=1}^M in the statement of Assumption 2.1
2. The Dudley metric is an Integral Probability Metric, as mentioned in the text preceding the statement of the Theorem (line 150) and defined precisely in Appendix B Eq. 23
3. We will add a few lines about Equation 22 in the main text for clarity. This is mostly a technical assumption. It is the most general initialization we could find in the literature for which our results on gradient flow dynamics hold. We note that it holds with high probability for He initialization, making it a valid assumption in realistic settings.

3. We would like to emphasize that the regularization we consider in this section is for the transfer algorithm, not the scratch-trained model. Equation 15 describes the generalization error of transfer learning using ridge regularized tuning of the final layer weights. Therefore, the expression for the transferability is simply Eq. 10 minus Eq. 15. This corresponds to replacing the $\sin^2\theta$ term in Eq. 13 with Eq. 15. Since we show that Eq. 15 is monotonically increasing in the ridge parameter $\lambda$, the transferability is strictly worse (smaller) for any nonzero $\lambda$. Since the scratch training procedure error is the same in this setting, the double descent phenomenon is still present. Since this concern is related to Weakness 2, we feel that it will be adequately addressed by the numerical experiment proposed in our response. With that being said, we will adjust the wording of Section 3.3.1 to emphasize that the regularization we consider is on the transferred model, not scratch training.

Dear Reviewer NGae,

Thank you for your detailed review of our paper. We are happy to hear that you found it “well-written and easy to follow”. We also thank you for your feedback on its weaknesses, which we address below.

Weaknesses:

1. While we agree that empirical verification in deeper networks would be compelling, the focus of our paper is primarily theoretical, so we have chosen to use the allotted pages to clearly describe the theory and its implications. The goal of our theory is to explain the origin of a phenomenon already observed in large scale deep neural networks, as demonstrated empirically in:

• “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution” (Liang 2022). This paper demonstrates a failure mode in fine tuning related to feature overlap for a number of benchmark distribution shift experiments.
• “What is being transferred in transfer learning” (Zhang 2021). This paper shows empirically that feature similarity drives positive transfer in a number of benchmark tasks. Thus while we appreciate your request to carry out experiments in deeper networks, this has partially already been done in prior experimental work (though without any theoretical explanation), and we feel it is a strength of our theory that it quantitatively explains what is happening linear networks, and provides a good description of two layer nonlinear network experiments in our paper, as well as deeper network experiments in prior work.

2. While self supervised learning is beyond the scope of this work, we believe our theoretical results have implications for this setting, as well. Consider a self supervised technique such as SimCLR. The objective of the network is to collapse augmented pairs toward similar network representations and pull different examples away from each other in representation space. Such an unsupervised representation learning algorithm would preserve features that are important for a downstream classification task if and only if data augmentations transform data points from one downstream task class into data points that are recognizably from the same class. Thus the data augmentation should scramble features that are irrelevant for determining the downstream task’s class labels, while preserving features that are important for determining the downstream task’s class label. If it does so, then features learned through SimCLR should transfer positively to the downstream task. So the logic is the same as in our paper: if the pretraining tasks (whether an upstream classification task or SimCLR representation learning) learns features that are useful for computing class labels in the downstream classification task, then, and only then, will positive transfer occur. To see this explicitly in a toy example consider a downstream classification task defined by two class conditional distributions consisting of two two-dimensional Gaussians with isotropic covariances and means $m_1 = -b\hat{x}$ and $m_2 = b\hat{x}$. The optimal binary classifier is the vector $\hat{x}$. We can consider SimCLR data augmentations that create augmented pairs from a single example by adding a small perturbation in the direction orthogonal to $\hat{x}$. This would collapse irrelevant dimensions for the downstream task and therefore lead to positive transfer on the downstream classification task. On the other hand if the data augmentation consisted of adding noise in the direction of $\hat{x}$, SimCLR would then collapse a feature dimension important for the downstream task, leading to negative transfer. This is the analog of feature mismatch in our theory.

Questions:

1. Addressed above.
2. We do believe that our theory is actionable since it highlights the potential pitfalls of pretraining in the feature learning regime. That is, learning the source function too well may actually hinder transfer performance if the target task does not lie in the same feature space as the source function. This suggests a number of regularization approaches that will lead to more flexible pretrained models, which we plan to discuss in the conclusion of our revised draft. One such method is to generate a sweep of pretraining runs that span the lazy learning to feature learning axis. This is done by adding a regularization term $\lambda \sum_{i=1}^m \lVert w_i - w_i^{(0)}\rVert^2$ to the pretraining loss, where $w_i$ is weight vector $i$ and $w_i^{(0)}$ is its value at initialization. When $\lambda$ is very large the weights do not move and the model trains in the lazy regime. We can then transfer models with varying $\lambda$ to the target task and choose the one that minimizes the generalization error on the target task. We will include the results of this experiment in the appendix of our revised draft, but to summarize, choosing the optimally regularized pretrained network eliminates negative transfer.

Dear Reviewer Ct4H,

We thank you for your review and apologize for any confusion stemming from the presentation of Theorem 2.2. In words, Theorem 2.2 says the following: consider a task with labels defined by $f(x)$ as in Eq 1. This function induces a joint distribution on the input/label pairs $p_f(x, y)$. Since $f(x)$ lives in $L_2$ by the assumption in the Theorem, we can express it as a linear combination of basis vectors $f(x) = \sum_{i=1}^M c_i \phi_i(x)$. Therefore, $f(x)$ lives in a subspace of $L_2$ defined by the span of the basis vectors $\phi_i$. We prove in the Theorem that we can always find another function $g(x)$ in the same subspace for which these distance based measures between $p_f$ and $p_g$ are arbitrarily large (i.e. for any $\delta >0$ $d(p_f, p_g) \ge \delta$). We provide a constructive proof in Appendix B.1. Now consider defining a transfer learning scenario using the distributions $p_f$ and $p_g$ as source and target distributions respectively. If training is succesful on the source task, the network will build a representation of the function $f(x)$. Since $f(x)$ and $g(x)$ can be built from different linear combinations of the same basis vectors (or features in the language of our paper), transfer learning will work very well between the two tasks. For example, one could simply train the output weights on the new task defined by $g(x)$ to learn the correct new linear combination of features. This is a scenario in which transfer is efficient, but the distributions $p_f$ and $p_g$ are very far apart as measured by these metrics. Conceptually, the purpose of this Theorem is to demonstrate that these metrics cannot reliably predict transfer performance, since one can always construct these counterexamples.

Your understanding of the Theorem appears to be the opposite of what we proved. In fact, it is not possible to build two tasks with identical distributions, for which transfer performance is poor. The reason for this is the following: we expect transfer efficiency to be poor when the important features associated with the two tasks do not lie in the same subspace of $L_2$, since the network must learn additional features in order to express the target function. If two tasks have the same distribution, they must be defined by the same function and trivially lie in the same subspace.

Thank you for pointing out that this picture was not clear in the current manuscript. We hope that with this explanation, you will consider re-evaluating the paper. We will reword this section in our revised draft to address any confusion. In addition, we plan to include an “Our Contributions” section after the introduction in which we clearly map out our argument.

Thank you for your response. I added a few additional comments in the original review that I will repeat here.

Thank you for your discussion around Theorem 2.2. While reading your rebuttal plus re-reading the paper, I have a few additional concerns that should be discussed.

You initially say that x and y are jointly distributed and that the labels are real numbers y in R (second sentence in Section 2). This is fine although a slight abuse of notation since random variables are not numbers. Then you define noisy versions of y and call them ys and yt. They are essentially noisy transformations of x. So far that is fine. Now in Assumption 2.1, you say that x is a data point and that y is a function on x that derives its statistical properties from f(x) + e. This is incorrect, because, if y is a function and not a random variable, then it cannot have statistical properties. Essentially, you are simultaneously saying that y=f(x)+e AND that y is a function of x (ie y(x) =f(x)+e), which doesn't make sense. Additionally, if x is a point, then y is just Gaussian since e is Gaussian. So where did the randomness in x go? Also are you only considering Gaussian data? Actually, it seems to be univariate Gaussian, and so the most basic case possible.

Your explanation is helpful but not complete. For example, let's say that we train the source model and get f which is essentially a wavelet transform. Then the basis functions phi's are just some wavelet basis in L^2. Then the subspace is the entire space. Then you are saying that you can find another function using the wavelet basis that gives poor performance. Obviously that is true -- just pick random wavelet coefficients. And if you want larger error you can just scale the coefficient weights.

Point 2 seems especially true, since you do not say that f and g are similar. You just say that you can find a g in the same subspace that f is in that is different.

If x has a fixed distribution for source and target tasks but fs and ft are different then the joint distributions p(x,ys) and p(x,yt) will be different. So I am not sure how you can claim in your rebuttal that "it is impossible to build two tasks with identical distribution for which transferring is poor". It must be that you are defining these functions, somehow, by learning them. But it is not clear how. Maybe you are assuming that the networks converge to a global minima?