A Theoretical Analysis of Fine-tuning with Linear Teachers
Year: 2021

Abstract:
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.

text: 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.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Related work
text: 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.

heading: 3 Preliminaries and settings
text: 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.

heading: 4 Analyzing fine-tuning in linear regression
text: In this section we analyze fine-tuning for the case of linear teachers for linear regression when using gradient descent for optimization. We define Θ(t) = w(t) ∈ Rd and overload f(x,Θ(t)) , x>w(t). In what follows we denote the parameter learned in the fine-tuning process by γ , limt→∞w(t).

heading: 4.1 Results
text: The following known results (e.g., [24, 25, 10]) show the inductive bias of gradient descent with non-zero initialization in under-determined linear regression and the corresponding population loss. Theorem 4.1. [24, 25, 10] When f(x,Θ) is a linear function, fine-tuning with GD under Assumption 3.1, Assumption 3.2 and Assumption 3.3 results in the following model:
γ = P⊥θS + P‖θT , (1)
and
R(γ) = ∥∥∥Σ1/2P⊥ (θT − θS)∥∥∥2 . (2)
Theorem 4.1 provides two interesting observations: the first is that γ consists of two parts, one which is the projection of the initial weights θS into the null space of X, and the other which is the projection of θT into the span of X. The second observation is that the population risk depends solely on the difference θT − θS that is projected to the null space of the data. For completeness, the proof of Theorem 4.1 is given in the supplementary.
Theorem 4.1 depends on the data matrix X (via P‖,P⊥). However, to better understand the properties of fine-tuning, a high probability bound on R that does not depend on X is desirable. We provide such a bound, highlighting the dependence of the population risk on the source and target tasks, and the target covariance Σ. Theorem 4.2. Assume the conditions of Theorem 4.1 hold, and assume that the rows of X are i.i.d. subgaussian centered random vectors. Then, there exists a constant c > 0, such that, for all δ ≥ 1, and for all 1 ≤ m ≤ d such that λm > 0, with probability at least 1− e−δ over X, the population risk R(γ) is bounded by:
2g(λ, δ, n)3 ‖P≤m(θT − θS)‖2
λ2m + 2g(λ, δ, n)‖P>m(θT − θS)‖2, (3)
where g(λ, δ, n) = cλ1 max{ √∑
i λi nλ1
, ∑
i λi nλ1
, √
δ n , δ n} and ∥∥∥Σ̃−Σ∥∥∥ ≤ g(λ, δ, n). In the proof, we address the randomness of P⊥(θT − θS) in (2), by decomposing θT − θS into its top-k span and bottom-k span components, and then applying the Davis-Kahan sin(Θ) theorem [33] to bound the norm of the projection of the former to the null space of the data. The full proof is given in the supp.
The bound in Theorem 4.2 has two key components. The first is the function g(λ, δ, n) that captures how well the covariance Σ is estimated, and shows the dependence of the bound on the number of train samples used (as it depends on n−0.5). The second relates to the two matrix norms of θT − θS with respect to different parts of the covariance Σ. Notice that the term relating to the top-k span decreases like n−1.5, while the term relating to bottom-k span decreases like n−0.5.
This theorem highlights the conditions under which fine-tuning is expected to perform well. For small enough n s.t. g(λ, δ, n) > 1, the bound mainly depends on ‖P≤m(θT −θS)‖. In this case, the bound will be low if θT and θS are close in the span of the top eigenvectors of the target distribution. On the other hand, for large enough n s.t. g(λ, δ, n) < 1, the bound mainly depends on ‖P>m(θT − θS)‖. Thus, the bound will be low if θT and θS are close in the span of the bottom eigenvectors of the target distribution.
We conclude with a remark regarding the integer m appearing in the bound, in the case where g(λ, δ, n) < 1. While finding the exact m that minimizes the bound is not straightforward, the trade-off in selecting it suggests taking the largest m which holds λm+1 ≈ λm. This will “cover” more of P>m (θT − θS) without greatly increasing the left part of (3).

heading: 4.2 Experiments
text: In Figure 1 we empirically verify the conclusions from the bound in (3). We set d = 1000 and design the target covariance Σ s.t. the first m = 50 eigenvalues are significantly larger than the rest (1.5 vs. 0.3). We then consider two settings for θT − θS . In the first, which we call “Top Eigen Align”, we select θT and θS such that P≤m(θT − θS) = 0. In the second which we call “Bottom Eigen Align” we set P>m(θT − θS) = 0. In both settings we use the same norm ‖θT − θS‖2, to show that the bound is not affected by this norm.
As discussed above, our bound suggests better generalization performance of “Bottom Eigen Align” for large n and better performance of “Top Eigen Align” for small n. Indeed, we see that while for very few samples “Top Eigen Align” has a lower population loss than “Bottom Eigen Align”, the population loss of ”Bottom Eigen Align” drops significantly as n grows, and drops to zero well before n = d.
We next evaluate the bound on fine-tuning tasks taken from the MNIST dataset [34], and compare it to alternative bounds. Specifically, since we do not expect bounds to be numerically accurate, we calculate the correlation between the actual risk in the experiment and the risk predicted by the bounds. The task we consider (both source and target) is binary classification, which we model as regression to outputs {−1,+1}. We generate K source-target task pairs (e.g., source task is label 2 vs label 3 and target tasks is label 5 vs label 6). For each such pair we perform source training followed by fine-tuning to target. We then record both the 0-1 error on an independent test set and the value predicted by the bounds. This way we obtain K pairs of points (i.e., actual error vs bound), and calculate the R2 for these pairs, indicating the level to which the bound agrees with the actual error. In addition to our bound in (3), we consider the following: the norm of source-target difference ‖θT − θS‖2 and a bound adapted from [25] to the case of fine-tuning.1 The results in Table 1 show that there is a strong correlation between our bound and the actual error, and the correlation is weaker for the other bounds.

heading: 5 Analyzing fine-tuning in deep linear networks
text: In this section we focus on the setting of overparameterized deep linear networks. Although the resulting function is linear in its inputs, like in the previous section, we shall see that the effect of fine-tuning is markedly different. Previous works (e.g. [35, 36]) have shown that linear networks exhibit many interesting properties which make them a good study case towards more complex non-linear networks.
We consider networks with L layers, given by the following matrices: Θ(t) = {W1(t), · · · ,WL(t)} s.t. Wj(t) ∈ Rdj−1×dj , d0 = d, dL = 1 and for 1 ≤ j ≤ L− 1 : dj ≥ d. We also define:
β(t) = W1(t) ·W2(t) · · ·WL(t),
such that f (x; Θ(t)) (t) = x>β(t). From Assumption 3.2, we have that β(0) = θS .
We recall the condition of perfect balancedness (or 0-balancedness) [32]: Definition 5.1. The weights of a depth L deep linear network at time t are called 0-balanced if:
Wj(t) >Wj(t) = Wj+1(t)Wj+1(t) > for j ∈ [L− 1]. (4) 1The adaptation is straightforward: since the population loss for non-random initialization depends on θT − θS instead of θT , we can replace the ground-truth expression θ? in Theorem 4 from [25] with θT − θS .
Our analysis requires the initial random initialization (prior to pretraining) to be 0-balanced, which can be achieved with a near zero random initialization, as discussed in [32]. We provide three results on the effect of fine-tuning in this setting. The first result shows the inductive bias of fine-tuning a depth L deep linear network (Theorem 5.2), which holds for arbitrary L and generalizes known results for L = 1 (Theorem 4.1) and L = 2 [29]. The second result analyzes the population risk of such a predictor when L→∞ for certain settings (Theorem 5.3 and Theorem 5.4). The third result shows why fixing the first layer (or any set of layers containing the first layer) after pretraining can harm fine-tuning (Theorem 5.5).
The next theorem characterizes the model learned by fine-tuning in the above setting (it can thus be viewed as the deep-linear version of the L = 1 result in Theorem 4.1):
Theorem 5.2. Assume that before pretraining, the weights of the model were 0-balanced and that Assumption 3.1, Assumption 3.2 and Assumption 3.3 hold. Then:
lim t→∞ β(t) =
( ‖ limt→∞ β(t)‖
‖θS‖
)L−1 L
P⊥θS + P‖θT (5)
and:
lim L→∞ lim t→∞ β(t) = ‖P‖θT ‖ ‖P‖θS‖ P⊥θS + P‖θT . (6)
To prove this, we focus on W1, and notice that the gradients Ẇ1(t) are in the span of X, and hence P⊥W1(0) and its norm remain static during the GF optimization ([30]). We then analyze the norm of the fine-tuned model by using the 0-balancedness property of the weights and the min-norm solution to the equivalent linear regression problem, and achieve (5). (6) is achieved by calculating the limit w.r.t. L. The proof of Theorem 5.2 is given in the supplementary.
Although the expression in (5) is not a closed form expression for limt→∞ β(t) (because ‖limt→∞ β(t)‖ appears on the RHS), taking L to infinity (6) does result in a closed form expression and demonstrates the effect of increasing model depth. As in (1), we see that the end-to-end equivalent has two components: one which is parallel to the data and one which is orthogonal to it. However, while in (1) the orthogonal component has the original norm of the orthogonal projection of θS , the expression in (6) offers a re-scaling of the norm of this component by some ratio that also depends on θT . Presenting this phenomenon for the infinity depth limit might look impractical, but the empirical results given in this section show that the effect of depth is apparent even for models of relatively small depth.

heading: 5.1 When Does Depth Help Fine-Tuning?
text: In this subsection we wish to understand the effect of depth on the population risk of the fine-tuned model. For simplicity we focus on the limit in (6), and denote β = limL→∞ limt→∞ β(t).
Since the linear network is a linear function of x, we can derive an expression for the population risk of the network, similar to (2):
R(β) = ∥∥∥∥Σ 12 P⊥(θT − ‖P‖θT ‖‖P‖θS‖θS )∥∥∥∥2 . (7)
However, since P‖ depends on the random matrix X, without further assumptions this expression by itself is not enough to understand the behaviour of R(β). Theorem 5.3 and Theorem 5.4 analyze cases for which a bound on (7) can be achieved, showing that it depends on ‖θT ‖ (θ̂T − θ̂S), i.e. the product of the norm of θT and the difference of the normalized θT and θS , compared to (2) which depends on the difference between the un-normalized vectors. This observation further highlights the fact that the distance between source and target vectors is not a good predictor of fine-tuning accuracy for some architectures, as fine-tuning can still succeed even if the source and target are very far as long as they are aligned.
We formalize this in the following result, where θT is identical to θS in direction, but not in norm.
Theorem 5.3. Assume that the conditions of Theorem 5.2 hold, and that θ̂T = θ̂S . Namely: θT = αθS , for α > 0,
then for L→∞ the risk of the end-to-end solution β is R(β) = 0,
while for the L = 1 solution γ, the risk is:
R(γ) = ( α− 1 α )2 ‖Σ1/2P⊥θT ‖2 6= 0 for α 6= 1, α > 0. (8)
This setting highlights our conclusion on the role of alignment in deep linear models: if the tasks are aligned, the deep linear predictor achieves zero generalization even with a single sample, while the population risk of the L = 1 predictor still depends on n.
Another example for this behaviour can be seen when X is i.i.d Gaussian (i.e., D = N (0, 1)d). Theorem 5.4. Assume that the conditions of Theorem 5.2 hold, and let X ∼ N (0, 1)d. Suppose n ≤ d, then there exists a constant c > 0 such that for any > 0 with probability at least 1− 4 exp(−c 2n)− 4 exp ( −c 2(d− n) ) the population risk for the L→∞ end-to-end predictor β is bounded as follows:
R(β) ≤ d− n d
(1 + )2 ‖θT ‖2 ∥∥∥θ̂T − θ̂S∥∥∥2 + d− n
d ζ(‖θT ‖)2, (9)
for ζ(‖θT ‖) ≈ ‖θT ‖. For the L = 1 linear regression solution γ this risk is bounded by
R(γ) ≤ d− n d (1 + )2 ‖θT − θS‖2 . (10)
The above result is a direct analysis of (7) when Σ = I by using Lemma 5.3.2 from [37] to analyze the effects of P‖,P⊥. Comparing (9) and (10), we see that while (10) depends on the distance between the two un-normalized tasks, (9) depends on the norm of the target task and the alignment of the tasks, but not at all on the norm of the source task. The proofs of Theorem 5.3 and Theorem 5.4 are given in the supp.

heading: 5.2 Deep linear fine-tuning with fixing the first layer(s)
text: A common trick when performing fine-tuning is to fix, or “freeze” (i.e. not train), the first k layers of a model during the optimization on the target task. This method reduces the risk of over-fitting these layers to the small training set.2 The next theorem shows that for deep linear networks this method degenerates the training process.
2This over-fitting is sometimes referred to as “catastrophic forgetting” of the source task.
Theorem 5.5. Assume the setting of Theorem 5.2. Then, if we freeze the first layer (or any number k of first layers) during fine-tuning, the fine-tuned model will be given by 〈β(t),x〉 = c〈x,θS〉, for some constant c.
The key idea in the proof is to show that the product of the k first layers is equal to θS up to a scaling factor, which is a result of [30]. The result implies that after fine-tuning the model is still equal to the source task, independently of the target task. Thus, fine-tuning essentially fails completely, and its error cannot be reduced with additional target data.
This result is achieved under the assumption of 0-balancedness prior to pretraining, which happens e.g. when initializing the weights with an infinitesimally small variance, as this property leads to the degeneracy of the output of the frozen k-layers. Though the proof of Theorem 5.5 depends on this 0-balancedness property of the network, the experiments shown in Figure 3 were conducted with a small initialization scale, that is not guaranteed to result in 0-balancedness, but rather in δ-approximate balancedness [32] when δ is small. These experiments show empirically that the phenomenon of learning failure is observed even when δ > 0. Intuitively, this is because the effective rank of the weight matrices is close to one, and thus learning the second layer is an ill-conditioned problem, which leads to slower convergence and can prevent the model from fine-tuning on the target data with a constant gradient step.
A possible workaround to this failure of learning would be to initialize the weights prior to pretraining with a larger scale of initialization (e.g. with Xavier [38]), thus increasing the rank of each layer and preventing degeneracy. Pre-training with multiple source tasks (as suggested in e.g. [22]) may also help the fine-tuning optimization.

heading: 5.3 Experiments
text: We next describe experiments that support the results in this section. Theorem 5.3 predicts that deeper nets will successfully learn a case where source and target vectors are aligned, but with different norms. This is demonstrated in Figure 2a where source and target tasks are related via θT = αθS + , where is a standard Gaussian vector whose norm is approximately 0.5 ‖θS‖. It can be seen that when α ≈ 1, there is no difference between models of different depth. However, as α increases, adding depth has a positive effect on fine-tuning accuracy. Theorem 5.4 predicts that the test loss for a deep linear model would depend only on the alignment of θS and θT (i.e.
∥∥∥θ̂T − θ̂S∥∥∥) and on the ‖θT ‖, but not on ‖θS‖. This is demonstrated in Figure 2b where source and target task are initialized s.t.
∥∥∥θ̂T − θ̂S∥∥∥ ≈ 0.1. In each experiment, either θT = αθ̂T or θS = αθ̂S , where α is the “Scaling Factor”, and the other has norm of 1. It can be seen that increasing the norm of the target vector harms generalization much more than increasing the norm of the source vector, as the theorem predicts, even for a relatively shallow model.
Theorem 5.5 states that fixing the first layer in deep linear nets can result in failure to fine-tune. We illustrate this empirically in Figure 3, where we compare three two-layer linear models on the same target task: 1) A “Frozen” model that fixes the first layer after pretraining. 2) A “Vanilla” model that trains the network from scratch on the target, ignoring the source pre-training. 3) A “Finetune” model that first trains on source and fine-tunes to target. As predicted by theory, the ”frozen” model’s performance is poor, and fine-tuning has better sample complexity.

heading: 6 Analyzing fine-tuning in shallow ReLU networks
text: Analyzing optimization and generalization in non-linear networks is challenging. However, analysis in the Neural Tangent Kernel (NTK) regime is sometimes simpler [39, 31]. Thus, here we take a first step towards understanding fine-tuning in non-linear networks by analyzing this problem in the NTK regime. Specifically, we consider the setting of a two-layer ReLU network with m neurons in the hidden layer. Hence, we consider Θ(t) = {W(t),a} and f (x; Θ(t)) = 1√
m
∑m r=1 arσ(x >wr(t))
where σ is the ReLU function, w1(t), . . . ,wm(t) ∈ Rd, the rows of W(t), are vectors in the first layer, and a ∈ {−1, 1}m is the vector of weights in the second layer. We initialize a uniformly and fix it during optimization as in [39]. Before pretraining, the first layer parameters are initialized from a standard Gaussian with variance κ2. We also assume that ‖x‖ = 1 for all x samples from D. We let f (X,Θ) ∈ Rn be the vector of predictions of f on the data X. For the next theorem we do not assume linear teachers, and instead assume an arbitrary labeling function gS such that yS = gS(XS), for XS ∈ RnS×d,yS ∈ RnS the pretraining data and labels, respectively. We also assume that y = gT (X) for some arbitrary function gT . For simplicity, we assume |y|i ≤ 1 for i ∈ [n]. We consider a setting where the pretraining phase is done using a two-layer network in the NTK regime, under the assumptions of Theorem 4.1 from [31] with respect to the variables m, κ, η and sufficiently many iterations.3 Next, in the fine-tuning phase, we train a network initialized with the weights given by the pretraining phase. We use the same value of m for the fine-tuning phase. We rely on the analysis given in [39, 31] and achieve an upper bound on the population risk of the fine-tuned model:
Theorem 6.1. Fix a failure probability δ ∈ (0, 1). We assume that Assumption 3.1 holds. Suppose κ = O ( λ0δ n ) , m ≥ κ−2 poly ( n, nS , λ −1 0 , δ −1). Consider any loss function ` : R× R→ [0, 1] that is 1-Lipschitz in the first argument such that `(y, y) = 0. Then with probability at least 1− δ,4 the two-layer neural network f(·,Θ(t)) fine-tuned by GD for t ≥ Ω ( 1 ηλ0 log ‖ỹ‖−12 )
iterations has population loss:
R (Θ(t)) ≤ 2
√ ỹ> (H∞) −1 ỹ
n +O √ log nλ0δ n  , (11) for ỹ ≡ y − f (X,Θ(0)).
3See the supp for a bound on the number of iterations. 4Over the random initialization of the pretraining network.
The above result shows that the true risk of the fine-tuned model is related to the distance of learned outputs y from the outputs after pretraining f (X,Θ(0)). The proof of Theorem 6.1 is given in the supp.
As in previous NTK regime analyses, this result holds when the weights of the fine-tuned model do not “move” too far away from the weights at random initialization. Thus, the proof approach is to bound the distance between the Gram matrix H(t) and the infinite-width gram matrix H∞ with a decreasing function in m. The main challenge is that the weights W(0) are not initialized i.i.d as described above. To address this we provide a careful analysis of the dynamics and show that H(t) is close to H at random initialization, even when considering the pretraining phase, which in turn is close to H∞.
We next apply our results to the case of linear source and target tasks. We thus assume that gS , gT are linear functions with parameters θS ,θT . For simplicity of exposition we assume f (x,Θ(0)) = x>θS exactly (Assumption 3.2). Before bounding the risk of fine-tuning we bound the RHS of (11) in the linear case:
Corollary 6.2. Suppose that gS(X) , X>θS , gT (X) , X>θT , and assume Assumption 3.2 holds. Then, √ ỹ>(H∞)−1ỹ ≤ 3 ‖θT − θS‖2 .
This is a direct corollary of Theorem 6.1 from [31] on ỹ defined above. Theorem 6.1 and Corollary 6.2 result in the a bound on the risk of the fine-tuned model:
Corollary 6.3. Under the conditions of Theorem 6.1 and Corollary 6.2, it holds that
R(Θ(t)) ≤ 6 ‖θT − θS‖2√
n +O √ log nλ0δ n  . We note that fine-tuning is improved as the distance between source and target decreases. In our analysis of linear networks (Theorem 4.2 and Theorem 5.4) we obtained a more fine-grained result depending on the covariance structure. We conjecture that the non-linear case will have similar results, which will likely involve the covariance structure in the NTK feature space.

heading: 7 Discussion
text: This paper gives a fine-grained analysis of the process of fine-tuning with linear teachers in several different architectures. It offers insights into the inductive bias of gradient-descent and the implied relation between the source task, the target task and the target covariance that is needed for this process to succeed. We believe our conclusions pave a way towards understanding why some pretrained models work better than others and what biases are transferred from those models during fine-tuning.
A limitation of our work is the simplicity of the models analyzed, and it would certainly be interesting to extend these. Our setting deals only with linear teachers, and assumes the label noise to be zero. Furthermore, we only show upper bounds on the population risk, and not matching lower bounds. For deep linear networks we assume a certain initialization which is less standard than normalized initializers such as Xavier. For non-linear models, we analyze the simple model of a shallow ReLU network, and only in the NTK regime.
An interesting direction to explore is formulating a bound similar to Theorem 4.2 for regression in the RKHS space given by the NTK, where the covariance is now over the RKHS space and thus more challenging to analyze. Another interesting setting is classification with exponential losses. Since the classifier learned by GD in this case has diverging norm, it is not clear how fine-tuning is beneficial, although in practice it often is. We leave these questions for future work.
Acknowledgments and Disclosure of Funding
This work has been supported by the Israeli Science Foundation research grant 1186/18 and the Yandex Initiative for Machine Learning. AB is supported by the Google Doctoral Fellowship in Machine Learning.

Review summaries:
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. 

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.

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.