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: The authors perform an analysis of fine-tuning a parameter learned from a source task for a target task in the few-shot setting (i.e. $n \ll d$). They provide:   1. A fine-grained analysis of the linear setting, where both source and target tasks involve linear functions of the input $x$, with risk bounds that incorporate the data covariance of the target task. 2. A similar analysis for the overparameterized linear setting, providing risk bounds that are not a function of $\lVert \theta_S - \theta_T\rVert$ in the infinite depth limit. Furthermore, in the same limit, they show that fixing the initial layers provably forces the predictor to remain in the span of $\theta_S$. 3. An analysis of fine-tuning for shallow ReLU networks in the NTK regime.  The theoretical results are corroborated by simulations.