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: This paper presents a theoretical analysis of fine-tuning for linear regression problems. The main results in the paper are as follows.  1. Theorem 4.2 develops an upper bound on the population risk on the target task in terms of the number of samples of n, the training time t, the data matrix of the target task and the difference between the true weights of the target task (the teacher) and the pertained weights on the source. There is technical novelty in how this bound is derived (matrix perturbation methods). The conclusion of this theorem is that if the teacher weights and pertained weights are close, as measured by the top few eigenvectors of the target data, then one gets a better population risk by training for fewer iterations. 2. This analysis is repeated for deep linear networks in Theorem 5.1 with similar conclusions. 3. Theorems 5.2-5.3 show an interesting phenomenon where if the student is a deep linear network, then the weights of the student need only be aligned with the teacher (as opposed to being close in ell-2 norm). 4. Theorem 6.1 is an analogue of the above result in the NTK regime. The technical novelty in the proof here is that one has to analyze the NTK dynamics in two stages.   Some experimental results are provided on synthetic data and MNIST.