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

Dear Reviewer jfTK,

Thank you for your thoughtful review of our work. We are grateful you found the manuscript easy to follow and we appreciate your feedback on areas that can be improved. We respond to your smaller comments below and leave our response to finite source datasets to the end, as this is more involved.

First, thank you for pointing out these two references we have not cited. They will be included in the Related Works section of the introduction While we agree that deep linear networks are particularly simple, we chose to prioritize analytic tractability in this work. However, as we mention in our response to Reviewer 65EM, a lot of the take aways are the same in the case of linear transfer with nonlinear networks. In this case we can describe the feature space using Reproducing Kernel Hilbert Space theory as in Appendix D, where we carry out a calculation for two layer ReLu networks. You are correct that Theorem A.2 is proven only for the Dudley Metric and KL divergence. However, we note that there is a hierarchy of integral probability metrics. For example, the theorem also holds for The Wasserstein metric since it is lower bounded by the Dudley Metric. Although we expect it to be the case, we were not able to prove that the relation holds for any IPM or $\phi-$divergence. You are correct that we have called our model label shift, but it appears that the term “concept drift” is more common in the literature. We will adjust our terminology accordingly in the revised manuscript As for the question of finite source datasets, we agree that this is a richer setting to study transfer learning. To this end, we have carried out a calculation for the full fine tuning transferability surface in the case of a finite source dataset, the results of which are described in the following. Let $\gamma_s = n_s/d$ and $\gamma_t = \n_t/d$ where $n_s$ and $n_t$ are the number of source and target data points respectively. Let $\sigma_s$ be the standard deviation of the Gaussian label noise in the source task (defined analogously to Assumption 3.1). Then the transferability is

The outline of the proof, which we will include in the final version of the paper, is similar in spirit to that of Theorem 3.9. In particular we know that gradient flow will converge to the minimum norm least square solution during source training, and we reason about the norm of its projection in the space orthogonal to the row space of the target data. There are a few interesting aspects of this expression. First, note that as long as the target task is overparameterized ($\gamma_t < 1$), the negative transfer boundary is completely determined by $\gamma_s$. Second, in the case of no label noise in the source task, the phase diagram is the same as in Fig 2. That is to say that fine tuning does not depend on the amount of source data if there is no label noise. The case of general noise level and source dataset size is more interesting: for some values of $\sigma_s$ there are disconnected regions of positive transfer. We will include the analysis of these equations, as well as plots of the phase diagram in the final version of the manuscript. We thank the reviewer for this suggestion.

Dear Reviewer AyJw,

Thank you for your review of our work. We are glad you found the paper well written and we thank you for your feedback on how to improve the manuscript. First, we appreciate your pointing out the additional references. We will discuss their relevance to our work in the Related Works section of our updated draft. Thank you for pointing out moments that felt unclear or overstated. In regards to Lines 104-109 we primarily want to highlight that training dynamics, implicit bias, and generalization error are all relevant to a theory of transfer learning and that deep linear models are an ideal test bed to explore the interplay of these phenomena since their manifestations are analytically solvable. We agree that many works have explored deep linear networks, and we will change these lines to highlight that our results build on prior work to tackle these multifaceted aspects of the transfer learning problem. Thank you for also pointing out the lack of clarity in Lines 271-273. We will change this sentence to: “We can view $\gamma$ as a dimensionless measure of the amount of target data, and $\cos^2 \theta$ as the amount of power that the target function has in the subspace of the pretraining task. The condition for negative transfer is satisfied when there is more target data than there is power in the pretrained subspace”. In response to your final question, there indeed is a regime where linear probing will outperform fine tuning. We can solve for this condition by subtracting Equation 15 from Equation 11 and looking for points in the $(\theta, \gamma)$ plane where this function is positive. In the noiseless case, the expression simplifies a bit. In the overparameterized regime ($\gamma < 1$) linear probing out performs fine tuning as long as $\gamma < \sin^2(\theta/2)$. The intersection of this condition with that for positive transfer ($\gamma < \cos^2{\theta}$) is satisfied when there is limited target data. We will include a plot of this region in the appendix of the final draft. For the underparameterized case ($\gamma > 1$) linear probing always induces negative transfer whereas fine tuning has zero transferability. So, although fine tuning has better transferability than linear probing, it’s not worth doing any pretraining in the underparameterized regime if there is no label noise.

Dear Reviewer UEjP,

Thank you for taking the time to review our paper. We are happy to hear that you found our results interesting and the presentation clear. We agree that we could be more detailed in our presentation of the simulations. In the final draft of the manuscript we plan to make the code publicly available, including a jupyter notebook that readers can use to explore the concepts of the paper for themselves. In addition, we will mention Appendix E in the main text, per your recommendation.

Dear Reviewer 65EM,

Thank you for your thoughtful review and feedback. You are correct that the rigorous proofs are limited to deep linear networks. This architecture exhibits certain symmetries which generate a conserved quantity in gradient flow that we exploit heavily to prove convergence. However, we believe the main takeaway of the work is applicable to other architectures, as demonstrated by the experiments in shallow nonlinear networks. Namely, during gradient descent in the feature learning regime, the model will learn features present in the source task. This defines a subspace of functions that the model can represent. When transferring to a target task with linear transfer, the model can only fit the portion of the target function that lives in the space spanned by the learned features of the target task. Since we can show this effect precisely in the deep linear setting we focus on this as the primary running example in the manuscript. The main practical take-away is that feature space overlap is the relevant object for predicting transferability. The focus of this paper is to demonstrate this phenomenon precisely in a solvable model, but we believe that interesting directions for future work include designing feature-level metrics that are predictive of transfer performance, or designing optimal sampling protocols for the source task when source sampling is readily available, but target data is scarce. Per your request of practical implications, we will add these ideas to our discussion section. While the space constraints keep us from including a simple example in the main body of the paper, we plan to make our code publicly available, including a Jupyter notebook with a simple example to help readers explore the concepts of our work. Additionally, we will include an additional figure with a cartoon depiction of the feature space overlap concept in the main body of the paper. We also thank you for raising the question of finite training time and learning rate. We show in Appendix C2 that in deep linear networks that convergence to a global minimum is subexponential in time, so at long times the results of our paper hold with very small error. However, optimal early stopping is an interesting and likely useful regularization technique that would prevent the model from sparsifying to the source features. We predict that the results of source task early stopping would be similar to those for weight decay, which we demonstrate in the appendix (Figure 5). That is, keeping the model from completely learning the source task may actually help transfer if the source and target features are very different. As for the learning rate, 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. Finally, we respond to your specific questions below:

We define the feature space overlap in the following way: during the source training the model will learn a function $f = \sum_{i} c_i \phi_i{x}$. We call span($\{\phi_i\}$) the feature space of the source task. The feature space overlap is the norm of the orthogonal projection (in the $L_2$ sense) of the target function $f_t$ into this space. In the case of our model, the learned feature space is $\beta_{s}$ (see Theorem 3.4), the projection is simply $\beta_{s}^T \beta{t}$. Since both vectors have unit norm, this can be viewed as the cosine angle between the two tasks, so we plot the overlap in radians. In the case of two layer ReLU networks, computing the overlap is also analytically tractable. We describe how this is done in Appendix D. That is the correct intuition. We measure the success of transfer learning by its generalization relative to training from scratch on the target dataset. If the target dataset size is very large, there is little benefit to pretraining. Yes we believe our results extend beyond two layer ReLU networks. Equation 18 holds generally for any Reproducing Kernel Hilbert Space. We choose two layer ReLU networks with Gaussian data since the projection in Equation 18 can be computed exactly.