A Theoretical Analysis of Fine-tuning with Linear Teachers

Fine-tuning is a common practice in deep learning , achieving excellent generalization results on downstream tasks using relatively little training data . Although widely used in practice , it is lacking strong theoretical understanding . Here we analyze the sample complexity of this scheme for regression with linear teachers in several architectures . Intuitively , the success of fine-tuning depends on the similarity between the source tasks and the target task , however measuring this similarity is non trivial . We show that generalization is related to a measure that considers the relation between the source task , target task and covariance structure of the target data . In the setting of linear regression , we show that under realistic settings a substantial sample complexity reduction is plausible when the above measure is low . For deep linear regression , we present a novel result regarding the inductive bias of gradient-based training when the network is initialized with pretrained weights . Using this result we show that the similarity measure for this setting is also affected by the depth of the network . We further present results on shallow ReLU models , and analyze the dependence of sample complexity on source and target tasks in this setting . 1 Introduction . In recent years fine-tuning has emerged as an effective approach to learning tasks with relatively little labeled data . In this setting , a model is first trained on a source task where much data is available ( e.g. , masked language modeling for BERT ) , and then it is further tuned using gradient descent methods on labeled data of a target task [ 1 , 2 , 3 , 4 ] . Furthermore , it has been observed that fine-tuning can outperform the strategy of fixing the representation learned on the source task , mainly in natural language processing [ 1 , 5 ] . Despite its empirical success , fine-tuning is poorly understood from a theoretical perspective . One apparent conundrum is that fine-tuned models can be much larger than the number of target training points , resulting in a heavily overparameterized model that is prone to overfitting and poor generalization . Thus , the answer must lie in the fact that fine-tuning is performed with gradient descent and not an arbitrary algorithm that could potentially “ ignore ” the source task [ 6 ] . Here we set out to formalize this problem and understand the factors that determine whether fine-tuning will succeed . We note that this question can be viewed as part of the general quest to understand the implicit bias of gradient based methods [ 6 , 7 , 8 , 9 , 10 , 11 , 12 , 13 ] , but in the particular context of fine-tuning . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We begin by highlighting the obvious link between fine-tuning and initialization . Namely , the only difference between “ standard ” training of a target task and fine-tuning on it , is the initial value of the model weights before beginning the gradient updates . Our goal is to understand the interplay between the model parameters at initialization ( namely the source task ) , the target distribution , and the accuracy of the fine-tuned model . A natural hypothesis is that the distance between the pretrained and fine-tuned model weights is what governs the success of fine-tuning . Indeed , some argue that this is both the key to bound the generalization error of a model and the implicit regularization of gradient-based methods [ 14 , 15 , 16 , 17 ] . However , this approach has been discouraged both by empirical testing of the generalization bounds inspired by it [ 18 ] and by theoretical works showing this can not be the inductive bias in deep neural networks [ 19 ] . Our results further establish the hypothesis that the success of fine-tuning is affected by other factors . In this paper we focus on the case in which both source and target regression tasks are linear functions of the input . We start by considering one layer linear networks , and derive novel sample complexity results for fine-tuning . We then proceed to the more complex case of deep linear networks , and prove a novel result characterizing the fine-tuned model as a function of both the weights after pretraining and the depth of the network , and use it to derive corresponding generalization results . Our results provide several surprising insights . First , we show that the covariance structure of the target data has a significant effect on the success of fine-tuning . In particular , sample complexity is affected by the degree of alignment between the source-target weight difference and the eigenvectors of the target covariance . Second , we find a strong connection between the depth of the network and the results of the fine-tuning process , since deeper networks will serve to cancel the effect of scale differences between source and target tasks . Our results are corroborated by empirical evaluations . We conclude with results on ReLU networks , providing the first sample complexity result for finetuning . For the case of linear teachers , this asserts a simple connection between the source and target models and the test error of fine-tuning . Taken together , our results demonstrate that fine-tuning is affected not only by some notion of distance between the source and target tasks , but also by the target covariance and the architecture of the model . These results can potentially lead to improved accuracy in this setting via appropriate design of the tasks used for pretraining and the choice of the model architecture . 2 Related work . Empirical work [ 20 ] has shown that two instances of models initialized from pre-trained weights are more similar in features space than those initialized randomly . Other works [ 21 , 22 , 23 ] have shown that fine-tuned models generalize well when the representation used by the target task is similar to the one used by the source tasks . In linear regression , [ 24 ] showed that gradient descent finds the solution with minimal distance to the initial weights . More recently , attention has turned towards the phenomenon of “ benign overfitting ” [ 25 , 26 ] in high dimensional linear regression , where despite fitting noise in training data , population risk may be low . Theoretical analysis of this setting [ 25 ] studied how it is affected by the data covariance structure . Benign overfitting was also recently analyzed in the context of ridge-regression [ 27 ] and online stochastic gradient descent [ 28 ] . Our work continues this line of work on high dimensional regression , but differs from the above papers as we start from a source task , then train on a fixed training set from a target task and consider the global optimum of the this training loss ( unlike online SGD ) . Furthermore , we go beyond the linear regression framework , and obtain surprising characteristics of fine-tuning in deep linear networks . For linear regression with deep linear models , [ 29 ] have recently shown an implicit bias for a twolayer network with deterministic initialization , and [ 30 ] have shown an implicit bias for a network with arbitrary depth and near-zero random initialization . Our work generalizes the inductive bias found by [ 29 ] to a network of arbitrary depth , and analyses the generalization error of such networks for infinite depth . For linear regression with shallow linear networks [ 31 ] have shown a generalization bound that depends only on the norm of the target task , which we use in Section 6 . 3 Preliminaries and settings . Notations Let ‖ · ‖ be the L2 norm for vectors and the spectral norm for matrices . For a vector v we denote v̂ , v‖v‖ . For a matrix M ∈ R d×d and some 0 ≤ m ≤ d , we define M≤m ∈ Rd×m to be the matrix containing the first m columns of M. Similarly , we let M > m denote the matrix containing the columns from m+ 1 to d in M. LetD be a distribution over Rd . Let Σ be the covariance matrix ofD and let VΛV > be its eigenvalue decomposition such that λ1 ≥ . . . ≥ λd . We define the projection matrices : P≤k , V≤kV > ≤k ; P > k , V > kV > > k , projecting onto the span of the top k eigenvectors of Σ , onto the span of the d−k bottom eigenvectors of Σ , respectively . We will refer to the former as the “ top-k span ” of Σ , and to the latter as the “ bottom-k span ” of Σ . Let X ∈ Rn×d be the row matrix of n < d samples drawn from D , and denote the empirical covariance matrix 1nX TX by Σ̃ . Define P‖ to be the projection matrix into the row space of X , and P⊥ to be the projection matrix into its orthogonal complement , i.e . : P‖ , X > ( XX > ) −1X , P⊥ , I−P‖ . Consider a set of parameters Θ , and let Θ ( t ) denote the set of parameters at time t. We denote the output of a model whose weights are Θ ( t ) on a vector x by f ( x ; Θ ( t ) ) ∈ R. In the different sections of this work we will overload f with different architectures . We consider the problem of fine-tuning based transfer learning in regression tasks with linear teachers . Let θT ∈ Rd be the ground-truth parameters of the target task , i.e . the linear teacher which we wish to learn , and y ∈ Rn be the target labels of X , s.t . y = XθT . We define L ( Θ ) to be the empirical MSE loss on X , y and define R ( Θ ) as the D population loss : L ( Θ ) , 1 n ‖f ( X , Θ ) − y‖22 , R ( Θ ) , Ex∼D [ ( x > θT − f ( x , Θ ) ) 2 ] . We separate the training procedure into two parts . In the first “ pretraining ” part , we train a model on nS pretraining samples XS ∈ RnS×d labeled by a linear teacher θS ( i.e. , yS = XSθS ∈ RnS ) , resulting in the set of model weights ΘS . In the second part , which we call fine-tuning , we initialize a model with the pretrained weights Θ ( 0 ) = ΘS and learn the target task by optimizing L ( Θ ( t ) ) . Optimization is done by either gradient descent ( GD ) or gradient flow ( GF ) . Let θ ( t ) be some weight vector or weight matrix in Θ ( t ) . The dynamics for gradient descent optimization with some learning rate η > 0 are θ ( t+1 ) = θ ( t ) −η ∂L ( Θ ( t ) ) ∂θ ( t ) , and the dynamics for gradient flow are θ̇ ( t ) = − ∂L ( Θ ( t ) ) ∂θ ( t ) . Next we state several assumptions about our setup . Assumption 3.1 . XXT is non-singular . i.e . the rows of X are linearly-independent . This assumption holds with high probability for , e.g. , a continuous distribution with support over a non-zero measure set . This assumption is only used for simplicity , as the high probability can be incorporated into the analysis . Assumption 3.2 ( Perfect pretraining ) . The pretraining optimization process learns the linear teacher perfectly , e.g . for linear regression we assume that f ( x , ΘS ) = x > θS , for x ∼ D. Notice that for linear and deep linear models , perfect pretraining can be achieved when nS ≥ d. Our results can be easily extended to the case where the equality f ( x , ΘS ) = x > θS holds approximately and with high probability , but for simplicity we assume equality . Assumption 3.3 ( Zero train loss ) . The fine-tuning converges , i.e . limt→∞ L ( Θ ( t ) ) = 0 . We note that when f is standard linear regression , arbitrarily small train loss can be obtained via gradient descent . For deep linear networks , it can be shown [ 32 ] that under suitable initialization a global optimum can be reached , and thus Assumption 3.3 holds for this framework as well .

Summary: Fine-tuning is a common technique in training large networks on relatively small datasets. It achieves excellent generalization results. However, it is not well understood theoretically. This paper attempts to analyze the sample complexity of fine-tuning for regression with linear teachers in several settings. Intutitvely, the success of fine-tuning depends on the similarity between the source tasks and the target tasks. This paper shows that the relevant measure has to do with the relation between the source task, the target task and the covariance structure of the target data. The analysis offers the insights into the inductive bias of gradient-descent optimization and the implied relation between the source task, the target task and the target covariance that is needed for this process to succeed.  This work focuses on the case where both source and target regression tasks are linear functions of the input. It first considers the case where the model architecture is one layer linear networks, and derive sample complexity results for fine-tuning. It shows that under certain assumptions there can be substantial sample complexity reduction when certain measure of source task and target task is low.   Further, it analyzes deep linear networks. There it derives a novel result providing the learned model as a function of initialization, and uses it to derive corresponding sample complexity bounds. The results provide several surprising insights. The first is that the covariance structure of the target data has a significant effect on the success of fine-tuning. In particular the relation between the vector of the source-target difference and the eigenvectors of the target covariance will affect how well fine-tuning will work. Second, it finds that the depth of the network can dramatically affect the results of the fine-tuning process, since deeper networks will serve to cancel the effect of scale differences between source and target tasks. Further, it corroborates the results by empirical evaluations.  A limitation of the work is the simplicity of the models analysed. Their setting deals only with linear teachers, and assumes the label noise to be zero. Furthermore, it only shows upper bounds on the population loss, and not matching lower bounds. For deep linear networks it assumes a certain initialization which is less standard than normalized initializers like Xavier. For non-linear models, it analyses the simple model of a Shallow-ReLU network, and only in the NTK regime.