H OW TO PREPARE YOUR TASK HEAD FOR FINETUNING
A BSTRACT
In deep learning, transferring information from a pretrained network to a downstream task by finetuning has many benefits.
The choice of task head plays an important role in fine-tuning, as the pretrained and downstream tasks are usually different.
Although there exist many different designs for finetuning, a full understanding of when and why these algorithms work has been elusive.
We analyze how the choice of task head controls feature adaptation and hence influences the downstream performance.
By decomposing the learning dynamics of adaptation, we find that the key aspect is the training accuracy and loss at the beginning of finetuning, which determines the “energy” available for the feature’s adaptation.
We identify a significant trend in the effect of changes in this initial energy on the resulting features after finetuning.
Specifically, as the energy increases, the Euclidean and cosine distances between the resulting and original features increase, while their dot products (and the resulting features’ norm) first increase then decrease.
Inspired by this, we give several practical principles that lead to better downstream performance.
We analytically prove this trend in an overparamterized linear setting, and verify its applicability to different experimental settings.
I NTRODUCTION
In the era of deep learning, pretraining a model on a large dataset and adapting it to downstream tasks is a popular workflow.
With the help of large amount of data and huge computing resources, the pretrained model can usually provide beneficial features for the downstream tasks.
Such a framework is proven to be efficient and effective in many domains and tasks, e.g. natural language processing (Kenton and Toutanova, 2019), computer vision (Chen et al., 2020b), graph based learning (Liu et al., 2022), and so on.
Although different variants of pretraining and finetuning (FT) methods are widely applied – including direct finetuning, finetuning after linear probing (Kumar et al., 2022), side-tuning (Zhang et al., 2020a), using different learning rates for different layers (Zhang et al., 2021), and more – a detailed understanding of how features are adapted during finetuning under different settings remains elusive.
Our work builds significantly off the analysis of Kumar et al.
(2022), who study the interactions between the “task head” (the final layer of the network, usually randomly initialized) and the “backbone” (usually copied from the pretrained model).
Kumar et al.
claim that the standard finetuning method, randomly initializing a task head then updating all parameters of the whole network, can distort the pretrained features and hence can deteriorate the generalization ability if (as they assume)
the previous backbone features were optimal for the downstream task.
By analyzing an overparameterized linear model, they prove that linear probing (i.e., only updating the parameters of the task head) first, followed by finetuning the whole network, leads to better performance in their setting.
In this work, we consider less stringent assumptions than they made, and study more practical settings from a different perspective.
First, we consider scenarios where the pretrained features are not optimal for the downstream tasks, thus feature adaptation is indeed beneficial.
Unlike the two extreme cases studied by Kumar et al.
(2022), i.e. finetuning with fully random initialization and
Second, our task heads are not necessarily linear.
Inspired by the illustrations of Olah et al.
(2020), it is reasonable to only preserve the lower layers of the pretrained model and reinitialize the top layers, assuming that the low-level features are common across task.
That is, the probed task head is non-linear, and we refer to this more general process as “head probing” (HP) rather than linear probing.
Our analysis can also help to explain feature behavior in this setting.
Finally, following our analysis, we provide a user guide to conclude when and why specific methods should be considered.
Specifically, we have one basic method: stop head probing earlier, before convergence; and three advanced tricks: 1.) use label smoothing during head probing; 2.) use more complex task head design; 3.) merge and reinitialize some later layers of the backbone and attach them to the task head.
In summary, in this work:
• we formalize and explain feature adaptation by decomposing the learning dynamics; • we find a non-trivial trend in feature adaptation and verify it in many cases; • and we show how controlling feature adaptation can improve downstream performance.
M OTIVATION
Pretrain-then-finetune is a popular workflow for many tasks in deep learning.
One common practice is to 1) randomly initialize a task head, 2) attach it to a pretrained backbone, then 3) finetune the whole network together (Li et al., 2020).
However, the untrained task head may distort the pretrained features during finetuning.
To solve this problem, Kumar et al.
(2022) propose to train the head to fully converge before the finetuning.
However, suppose we train the head long enough and its training accuracy (HP-train-acc) converges to 100%, then the features won’t change during the finetuning stage.
To sum up from the above, we can see that neither probing the head to converge nor no probing is optimal, since the pretraining and downstream tasks (or datasets) are usually distinct.
To verify this argument, we HP various number of epochs before finetuning, and record the corresponding validation accuracy after finetuning (FT-valid-acc for short), and the results are shown in Figure 1.
It is surprising to see that stopping the head training earlier (before the convergence of HP-train-acc) brings more improvement.
As the only variable among these experiments is the parameters of the head before finetuning, the following two questions emerge:
How does the task head influence the pretrained features during finetuning?
How does the feature adaptation influence the generalization performance after finetuning?
B ACKGROUND
We first clarify the scope of our analysis.
We don’t consider the details of the pretraining procedure, instead just assuming that there are some well-trained checkpoints for a particular dataset or task.
Meanwhile, our formulation is not restricted to classification tasks; our use of the term “label” or “target” can be any form of supervisory signals.
3. P ROBLEM S ETTING AND T WO S TAGE T RAINING
When training a model for a downstream task, our goal is to find a predictor f ◦ g : R d → R k that maps a high-dimensional input signal x ∈ R d to a task-related prediction q ∈ R k .
As depicted by the left-bottom panel of Figure 1, we split the predictor into two parts: the backbone f ( x ; B ) : R d → R h which maps the input signal to intermediate representations z ∈ R h , and the task head g ( z ; v ) : R h
→ R k which gives the prediction vector q ∈ R k (e.g. logits in classification tasks using cross-entropy loss).
Usually, the backbone f is parameterized by B and initialized by copying from a pretrained model.
The task head g parameterized by v , on the other hand, is usually randomly initialized and might be a complex non-linear function.
The training has two distinct stages: 1) head probing (HP) where we fix f and only update the parameters v of g for τ epochs; 2) finetuning (FT) where the parameters { v , B } of f ◦ g are updated together until convergence.
In this work, we analyze how the FT stage is influenced by the architecture and v at the beginning of finetuning.
Following the general formulation above, we show a simple overparameterized linear regression (or equivalently, a binary classification) below as a case study to illustrate more insights.
Suppose the
N input signals are stacked, X =
[ x (1) , ..., x ( N ) ]
⊤ ∈ R N × d , the loss function can be written as
[Equation]
(1) where B ∈ R h × d , v ∈ R h , and Y ∈ R N .
That is, k
= 1 , z
=
Bx , q = z ⊤ v , and L ( q, y )
=
12 ( q − y ) 2 .
3. W HAT TO E XPECT D URING A DAPTATION
Compared with training a randomly initialized model, adaptation on downstream tasks needs more care.
One reason is that the pretrained parameters from f ( x ; B ) inherits all the information from the pretraining task, even any bias and noise.
Furthermore, as mentioned by Kumar et al.
and Du et al.
(2018), f ( x ; B ) is tied to g ( z ; v ) at each time step during FT, thus the bias and noise in f ( x ; B ) also influence the learning of g ( z ; v ) .
Hence, before conducting downstream adaptation, we might consider: to what extent do we want to change the feature extractor f
( x ; B ) ?
Here, we list three possible cases of how much we should update f ( x ; B ) :
• Strong: the pretrained features are far from the optimal ones for downstream task, so we need substantial feature adaptation.
• Mild: f ( x ; B ) is reasonably good, but adaptation to the downstream domain is helpful.
• Tiny: the pretrained f ( x ; B ) is near optimal and only need to be slightly adapted.
In the rest of this paper, we first analyze how f ( x ; B ) and g ( z ; v ) interact with each other, under both the general case and the simplified case.
Based on our observations, we propose several practical principles in Section 5.1 for better downstream performance.
I NTERACTION BETWEEN THE B ACKBONE AND H EAD
Rather than linking the choice of v 0 (the task head parameters at the beginning of FT) to the loss function, we sketch how f ( x ; B ) and g ( z ; v ) interact during FT, which depicts how z changes accordingly.
Although our analysis cannot provide any theoretical guarantee, knowing how z changes under different v 0 will lead us to better HP-FT design in practice.
4. A VERAGE I NITIAL E NERGY
We start from analyzing the behavior of z ( j ) t = f
( x ( j ) ; B t ) , i.e. the feature extractor at time t , when the network parameters are updated with samples x (1) , ..., x ( N ) using gradient descent. When using [Figure]
cross-entropy loss, we can have the following result:
[Equation]
where γ is the learning rate, κ ( j,n ) t [Equation]
∈ R h × h is the empirical neural tangent kernel (NTK) of the backbone between x ( j ) and x ( n ) at time t , 2 ∇ z q ( n ) t ∈
R k × h is the gradient of the task head prediction w.r.t the representation vector z at time t , 3 p ( n ) t =
Softmax ( q ( n ) t ) is the predicted probability vector for input x ( n ) , and e y n is the one-hot vector of the label y n .
In this decomposition, the first term is often understood to change only slowly during FT, i.e. the lazy-parameters setting used by Chizat et al.
; Yang and Hu (2020).
The second term determines the direction, and the last term provides “energy” for the update of z .
Formally, we can define the Average Initial Energy (AIE) and use it to bound the norm of z ( j ) T − z
( j ) 0 :
Proposition 1.
E x ( j ) ∥
z ( j ) T − z
( j )
∥ 2 ≤ c · E aie , where E aie ≜ E x ( n )
∥
e y n − p ( n )
∥ 2 is the Average
Initial Energy (AIE).
Here c is a constant, T is the FT epochs, and p ( n )is the model’s prediction of sample x ( n ) at the beginning of FT (or at the end of HP).
Proof.
See Appendix A.
Although this bound is loose and requires approximations based on assumptions, the proposition supports our intuition: if we prob the head for too long (i.e. a large τ ), the HP-train-acc can be very high, hence on average p 0 is close to e y and the features adapt less, and vice versa.
Our intuition is further verified in the right panel of Figure 2 where we plot the predicted probability of the correct class, i.e. [ p ( n ) 0 ] y n , for each sample.
Furthermore, we find that it is unlikely to obtain zero energy (not as assumed by Kumar et al.
in many practical applications: even with high accuracy there will always be some gap between p 0 and e y by definition, and moreover the training accuracy after HP is sometimes far from perfect (see Section 5).
4. N ON - TRIVIAL TREND OF FEATURE ADAPTATION
Following Proposition 1, we can link the adaptation of features to τ via AIE.
In the proof of Proposition 1, we disentangle the dependence of the direction and energy terms using Cauchy-Schwarz [Table] = inequality, i.e., ( ∇ z q ) ⊤ ( e y − p [-0 ) ≤ ∥∇ z q ∥ 2 · ∥ e y − p 
-][Equation]
{+) ≤ ∥∇ z q ∥ 2 · ∥ e y − p 
+}
∥ 2 .
However, recall that the direction term
Equation (2) in also plays an important role, especially at the beginning of finetuning.
To verify our hypothesis about the direction term, we depict the change of this term during finetuning in the right panel of Figure 2.
Suppose a linear task-head (non-linear heads have similar behavior in the NTK regimes), this quantity can be approximated by the norm of the gradients to v t , since ∥∇ [Equation]
[-As we find ∥ v t ∥ 2 changes little during the 
 In a linear model z-]∥∇
[Equation] L∥ 22 .
As we find ∥ v t ∥ 2 changes little during the finetuning stage, the large ∥∇ v L∥ 2 is more likely from a big direction change.
As illustrated by
Figure 2, when τ = 0 , ∇ z q t changes a lot at the beginning of FT, which can make z ( n ) change in inconsistent directions.
When τ = 1024 , the direction term changes only a little through finetuning.
This finding inspires us to look deeper at what is the difference between a strong adaptation (e.g., τ = 0 ) and a mild adaptation (e.g., τ = 4 ).
To get an overall picture of z ’s change, only ∥ z T − z 0 ∥is not enough.
Hence we analyze the following four quantities related to the similarity between the features before finetuning z
and afterwards z T : ∥
z T − z
∥ 2 , ∥ z T ∥ 2 , z ⊤ T z 0 , and cos( z T , z 0 ) .
With an overparameterized model, we can analytically calculate the expressions of them and make the following conclusion:
Proposition 2 (Informal) .
In an overparameterized two-layer linear model, when τ increases (the
AIE decreases),
∥ z T − z
∥ 22 monotonically decreases while ∥ z T ∥ 22 and z ⊤ T z 0 exhibit a quadratic trend.
The trend of cos( z T , z 0 ) is hard to predict, but there is a phase that this value increases fast.
Such a trend is illustrated in Figure 3, supported by various experiments in Appendix C, and strictly proved in Appendix D. With this information, we can infer the behavior of z T and sketch it in the second and the third panels.
In the “tiny” case, when v τ 0 is fully converged ( τ → ∞ ), the features are almost stable during finetuning; this only works well when the pretrained features are perfect for the downstream task.
In the “mild” case, when τ is reasonably large, the resulting z T will be stretched ( ∥ z T ∥ 2 increases) in a similar direction (cosine changes little).
This kind of z can make the original overlapped features become more separable without changing the manifold of the features too much, which is desirable in most transfer learning scenarios.
For the “strong” case, where we only HP for a few updates or simply use a random head, z T will change in an unpredictable way, especially for the early updates.
Although the fine-tuned model may generalize well, z T might be quite different from z 0 .
Thus, if we believe the pretrained features are too specific to the pretraining dataset, the “strong” case is a reasonable choice.
Note that even though the scatter plots of the mild and strong cases look similar, the corresponding feature manifold might be quite different.
To verify this, we first run HP τ -FT
(i.e., load the pretrained model, prob the head for τ epochs, then finetune to converge) for 7 different τ .
Then we save the converged backbone and task-head separately for each τ , and pair-wisely exchange the head and backbone to build 49 new models (without further tuning).
The training and validation accuracy of the 49 models are reported in Figure 3.
The results match well with our analysis: 1) the off-diagonal values are lower, which means the representations learned by experiments with different τ are different; 2) for the strong adapted backbones (the first two columns), the heads from other cases are not compatible, which means the features’ manifold are significantly changed; 3) for the mild and tiny cases (latter columns), the aforementioned incompatibility almost disappears, which means the features’ manifold of these cases are quite similar.
4. B ACKBONE D EPTH AND H EAD C APACITY
Beyond how long we train the task head, the structure and capacity of the task head also influence the training prediction after this stage, thus the adaptation of z t .
We briefly discuss the trend here and verify them in Section 5.
When using a low-capacity head (e.g. a linear head parameterized by a h × k matrix), it might be hard to preserve the pretrained z 0 even a very large τ is chosen, as the head cannot achieve a high enough training accuracy to decrease AIE.
On the other hand, if the capacity of the task head is much bigger than the backbone, the information from the pretrained network might be easily distorted.
For example, consider using only the first block of a pretrained ResNet18 as the backbone and concatenating the other 3 blocks with a 10-layer wide MLP.
If this huge task head is randomly initialized, the information contained in the backbone could be completely rewritten, as the random changing phase of z t can be very long (remember that in Equation (2), z t will change in random directions before ∇ v
q t becomes less noisy).
E XAMPLES ON R EAL T ASKS
In this section, we first provide a “user guide” of all the aforementioned methods.
Note that given the pretrained and downstream tasks, determining the optimum adaptation energy is rather heuristic: the user guide only provides some high-level suggestions for the practitioners.
We then provide abundant real examples of how to apply these principles.
Generally, we can first try sweeping the optimal τ using linear head and consider other advanced tricks if necessary.
5. U SER G UIDE
The goal of this paper is to provide a toolbox for preparing the task-head before finetuning.
Here, we suggest a “ phenomenon → hypothesis → solution ” workflow and use the validation performance for verification.
Recall how we mitigate the overfitting problem using dropout: validation accuracy decrease after convergence → model overfitted → add dropout .
Similarly, we can have: HP trainacc converge to 100% → no enough energy → use smaller τ .
However, in practice, it is unknown that how much energy is beneficial, as the neural network might not encode the input samples as we humans expected.
Hence, we suggest starting from the basic setting, i.e. using a linear head and sweeping τ to get a high-level understanding of how the energy influence the downstream generalization.
Usually, selecting the optimal τ ∗ using validation accuracy can ensure a reasonably good test performance, as verified in the next subsection.
The advanced tricks are only applicable to specific scenarios and need more consideration.
For example, if we really want tiny energy but using the linear head only achieves less than 50% training accuracy after head probing, we can consider an MLP head to increase the converged training accuracy (hence reduce the energy).
If we want a mild adaptation, but the training accuracy during
HP goes to 100% too fast, using label smoothing during HP can be considered.
If we believe the downstream task only needs low-level features of the backbone, partial-backbone is a good choice.
However, these advanced tricks also have side effects that are deleterious to the downstream performance.
We will analyze their advantage and limitations in Section 5.3 with concrete examples.
5. B ASIC M ETHOD : EARLIER STOPPING HP
Following our theory, when the learning rate for HP is fixed, the value of τ is positively correlated with the training accuracy, thus negatively correlated with energy.
Hence, given a pretrained backbone and a downstream task, we can always start from sweeping τ (see Figure 4).
As stated in the motivation part, during finetuning, we usually expect some feature adaptation while keeping some pretrained information.
Hence, neither τ = 0 nor τ = ∞ is optimal.
We find such a downstream data: [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
trend is consistent across various settings: from image input to graph input, from classification task to segmentation task, and from supervisory pre-training to unsupervised pre-training.
Other than this general trend, the nuance of experiments under different settings also supports our analysis well.
Specifically, in the image classification experiments shown in Figure 4, adapting a ResNet50 to STL10 behaves differently on different pretraining task.
In the first column, we see a large τ (small energy) hurts the downstream generalization performance, because the features pretrained on CIFAR10 might be far from optimal for the downstream task.
In contrast, the features pretrained on ImageNet-1K (IN1K for short) all provide good results.
Among those IN1K pretrained models, the model from a supervised classification task leads the best overall downstream performance, but large τ is still harmful.
In other words, the features pretrained this way might be somewhat too specific to the pretraining dataset, and hence mild adaptation is beneficial.
Regarding the unsupervised pretraining cases, BYOL (Grill et al., 2020) is less sensitive to the choice of τ , while MoCo (He et al., 2020) behaves more similarly to the supervised case.
Another interesting finding is that the HP training accuracy at the optimal τ ∗ is usually smaller than the converged value: we cannot select τ based on the standard early stopping criterion on HP train accuracy.
As the task head under τ ∗ usually has not converged on the pretraining dataset, we call this method “earlier stopping HP” (HP τ ∗ -FT for short).
As shown in Table 1, HP τ ∗ -FT can always bring improvements.
5. A DVANCED T RICKS : SUPERIORITY AND LIMITATION
MLP Head: Instead of using a linear head, authors of Chen et al.
(2020b) claim that using an MLP head can sometimes bring improvement.
Following the analysis of this paper, we can consider this trick when we want small energy while the linear head cannot reach a high HP training accuracy even after a long HP.
For example, in the first two columns in Table 2, the HP-train-acc in the linear head cases plateaued after reaching 92% or 78%, while a 2-layer MLP can reach 99% and 95%.
As the energy decreases, features adapt less during finetuning (see the decreased distance metrics) and the final models generalize better.
However, we should be careful when applying this trick if we want to use a small τ (i.e., large energy).
Recall our analysis that the inconsistent direction term (i.e., ∇ z q t defined in Equation (2)) makes the feature adaptation more unpredictable at the beginning of finetuning.
Increasing the head capacity in this case would make the head converge slower and hence prolong such a chaos phase.
[-Label Smoothing HP: Recall that the energy is upper bounded by ∥ e y − p 
 ∥ 2 2 , where p 0 is model’s prediction after HP for τ epochs. It converges to the labels used in HP when τ → ∞ .-][Table]
, instead changing τ , we can also manipulate HP to achieve a similar goal.
One simple yet effective way is label smoothing (M¨uller et al., 2019, e.g.).
By setting the labels during HP as η HP e y + (1 − η HP ) ∗ u , where u is a uniform K-class categorical distribution, the HP stage can always reserve at least (1 − η HP ) ∗ ∥ e y − u
∥ 2 energy for the following feature adaptation, even τ → ∞ .
Such a trick (lsHP for short) is quite helpful when the HP-train-acc converges to 90%+ very fast, yet we still want a mild adaptation, like the example shown in the second two columns in
Table 2.
With lsHP, we see that the features adapt more even the HP-train-accs are unchanged.
To verify that the aforementioned improvement comes from the reserved energy during lsHP, we further try using smoothed labels during finetuning (e.g., η FT = 0 . 9 ).
The results match our analysis: when η HP = η FT = 0 .
9 , the reserved energy disappears, as the labels of the two phases are the same again.
Hence, all the numbers under this condition are quite similar to the baseline case ( η HP =
η FT = 1 ).
For the “opposite energy” case, we observe a larger adaptation but a worse generalization performance.
That is because the reserved energy make the features adapt in opposite directions.
These results remind us that if we decide to use smooth label in both finetuning and head probing (e.g., we assume most of the samples in the downstream dataset are ambiguous), we need to set η HP ≤ η FT to ensure a correct direction.
In summary, the lsHP trick is suitable for scenarios where the pretrained features are pretty well and the standard HP converges to 90%+ very fast.
When the HP-train-acc is too low, the assumption used in HP, i.e. p 0 converges to the labels, no longer holds.
Hence, lsHP does not always bring enhancement, an example is given in the last row in Table 1.
[Table]
[Table]
This is a more intricate trick that requires a stronger and more heuristic hypothesis.
For example, there is a common belief in the deep-learning community that the lower layers extract fundamental features (e.g. edges, curves, and texture in vision tasks) while the higher layers learn more semantic features (e.g. dog heads or wheels)
(Baldock et al., 2021;
Olah et al., 2020).
Hence, if we believe the downstream task treats the low-level features as beneficial while the high-level features are harmful, reinitializing the higher layers (like the fortuitous forgetting mentioned by Zhou et al.
and incorporating them as part of the task head can be beneficial.
See the results in Table 3, where the Quick 4 dataset is likely to rely more on the low-level features learned during pretraining.
So the optimal setting for the Quick column is removing information from the last three layers (i.e., L4.3, L4.2, and L4.1) in ResNet50, while the optimal setting for the other two cases is to reinitialize only the last layer (i.e., L4.3).
Using this trick, the task-head capacity might increase significantly (see the HP-train-acc increase to 100% after reinitializing).
Hence, the principles discussed in the MLP-head trick also hold here.
However, as it is hard to figure out what are the beneficial features for downstream tasks, such a trick has the lowest priority in our toolbox
R ELATED W ORK AND D ISCUSSIONS
HP, FT and HP-FT.
Head probing and fine tuning are two fundamental algorithms in transfer learning, which have also attracted much discussion and debate.
Intuitively, by freezing the pretrained backbone, HP will re-combine the features without distorting them, and hence yields better performance than FT when features are perfect (Peters et al., 2019).
However, when pretrain and downstream tasks are very different, adapting the features is important and FT outperforms HP (Chen et al., 2020b; Zhai et al., 2019;
He et al., 2022).
Combining the strengths of HP and FT, authors of
Kumar et al.
(2022) demonstrate that HP-FT (i.e., first HP, then FT) yields the best performance on both in-distribution and out-of-distribution cases.
Although they provide some theoretical guarantees for the superiority of HP-FT, some of their assumptions are dubious.
To shed more lights on the pros and cons of these two methods, this paper analyze in detail how the features change under different task-head settings.
Specifically, by controlling the HP epochs τ or otherwise influencing the energy term, we can design different types of feature adaptations for different downstream tasks.
The relationship between the choice of task head and resulting z ’s adaptation is explained using
Equation (2), verified by various experiments, and proved in an overparameterized linear model.
Although the paper only analyze these fundamental settings, we believe the analysis provided here can also be combined with (or explain the benefits of) other more complex finetuning mechanisms, like those of Guo et al.
; Zhang et al.
(2020a);
Aghajanyan et al.
; Howard and Ruder (2018).
For example, a common practice in transfer learning is to use a small learning rate for the backbone and a large learning rate for task head.
From our perspective, this method can weaken the influence of the “noisy direction” term and make the features adapt in a low energy condition.
Backbone and head depth.
Besides the influence of the initial value of the head, this paper also discusses the influence of the relative capacity between backbone and head.
Chen et al.
(2020b) show that a simple MLP head sometimes also brings enhancement.
When the backbone only copies the early layers of the pretrained network, as the downstream task might need different levels of features (Olah et al., 2020; Baldock et al., 2021; Zhang et al., 2020b), the later layers and the original head can be combined as the new complex task-head, like the example in Figure 1.
Furthermore, in some encoder-decoder style models, like some the models for some language tasks (Peters et al., 2019;
Zhu et al., 2019), the head (decoder) might have comparable capacity with the backbone (encoder).
In another extreme case, if we use pretrained word2vec features (Mikolov et al., 2013) and plan to FT them in a downstream task, the task head is the whole network, which is much bigger than the embedding layer.
We believe the discussion in this paper might also help inspire the design of
HP-FT strategy for these practical scenarios.
C ONCLUSION
This paper studies how the choice of task heads influence the pretrained features z ’s adaptation and hence influence the downstream performance.
By decomposing the learning dynamics of z , we find the keys are the energy and direction terms, which highly correlate with the accuracy at the beginning of the FT process (i.e., the accuracy after HP).
Hence under a common HP-FT framework, we carefully analyze the relationship between the HP epochs ( τ ) and the features’ adaptation (discrepancy between z T and z 0 ).
Different from most existing works, which mainly focus on Euclidean or cosine distance, we further analyze z ⊤ T z 0 and ∥
z 0 ∥ 22 to depict a more comprehensive relationship between τ and z T , like the examples in Figure 3.
This non-trivial trend is strictly proved in an overparameterized linear model and experimentally verified under different practical settings.
Based on these observations, we speculate that a suitable adaptation of z T is beneficial when the pretrained features are not perfect for the downstream tasks.
Under different experimental settings, we illustrate how to achieve better downstream performance by controlling the adaptations of z T (using early-stopping HP, label smoothing during HP, or manipulating the head’s capacity).
Finally, there are still many open questions and phenomena not address by this study.
For example, methods to quantitatively or even adaptively analyze the discrepancy between the pretrain and downstream tasks would be very useful in knowing what kind of head probing to perform in advance.
However, we believe the methods proposed in this paper can help provide a new perspective on understanding feature adaptation, which we hope will aid future work in this area.
A D ECOMPOSITION OF CHANGE OF Z
Proof of Equation (2):
Recall z = f B ( x ) , q = g z
( z ) and p = Softmax ( q ) .
Use b ∈ R hd × 1 to represent the vector form of
B ∈ R h × d .
The MSE loss on one sample is L mse (cid:0) q ( n ) , e y n (cid:1)
=
12 ∥ q ( n )
− e y n ∥ 22 .
The cross-entropy loss on one sample is L ce (cid:0) q ( n ) , e y n (cid:1)
=
H (cid:0) e y n , q ( n ) (cid:1) , where H ( p , q ) = − (cid:80) k
p k log q k is the cross entropy for categorical distribution.
Using 1st Taylor expansion, we have:
[Equation]
We then calculate b t +1 − b t assuming the parameters are updated in batch-SGD:
[Equation]
For different loss functions
[Equation]
we have different expressions for the last term:
[Equation]
where q is the logits and p
= Softmax ( q ) is the predicting probability.
We see these two kinds of loss have similar expression on this term, so we only explain the cross-entropy version (as it is more common in practices) in the main content.
Using the above expressions, we can first bound the high-order term (cross-entropy version):
[Equation]
as long as the hyperparameters are appropriately chosen and the loss doesn’t blow up (or gradient clipping is applied) in FT stage.
Finally
[Equation]
combining all the expressions, Equation (10) can be rewritten as:
[Equation]
which is the same with Equation (2) in the main context
[Equation]
Proof of Proposition 1:
It is easy to get z ( j ) T − z
( j ) 0 by stacking the LHS of Equation (2) (or Equation (10)) under different t :
[Equation]
where the first equation follows definition.
The second equation is assuming a uniform training sample distribution, i.e., p ( x ) = 1 N .
The third equation follows the slow-change NTK assumption, hence κ ( j,n ) no longer depends on t .
The last equation follows the linear-head assumption, i.e., ∇ z q ( n ) t = v t .
We cannot go further as there are three matrices (or vectors) in the expectation.
Hence we instead analyze the F-norm (i.e., L2-norm for the vector z ) of the features’ change.
[Table]
[Equation]
Here the first equation is by definition.
The second inequality follows triangle inequality.
The third inequality follows Cauchy-Schwarz inequality.
The forth inequality is assuming the F-norm of NTK is bounded by C 1 .
The fifth equation is sweeping the summation order.
The sixth inequality also follows Cauchy-Schwarz inequality.
The seventh inequality assumes the norm of the last layer is bounded by C 2 .
The eighth inequality is assuming a stable learning process where the gap between e y n and p ( n )keeps decreasing during training.
Finally, by taking expectation of all the input samples, we can have:
[Equation]
23) where c = γ · C 1 · C 2 · T is a constant and E x ( n ) ∥
e y n − p ( n )
∥ is the AIE term ( E aie ) defined in
Proposition 1.
Verification of Slow Kernel Change Assumption:
pretrained.
We decompose the one-step dynamics of z ( j ) t in Equation (2) into three parts, which represent kernel, direction, and energy respectively.
In our analysis, we have a mild assumption that the kernel of the backbone, i.e., κ ( j,n ) t =
[Equation] ∇ B z [Equation]
n ) t (cid:17) ⊤ ∈ R h × h , changes slow during finetuning when the learning rate is small and the training is stable (the loss converges).
Although there are several works supporting this assumption (Yang and Hu, 2020; Geiger et al., 2020), to make the paper more self-contained, we directly observe how this term adapts during pretraining in Figure 6.
We measure the change of this term by calculating k norm =
∥ κ ( j,n ) t ∥ F and k gap =
∥ κ ( j,n ) t +1 − κ ( j,n ) t ∥ F for each t during finetuning.
As computing the empirical NTK on the whole dataset requires huge memory, we randomly sample 50 x as our “probing samples”, computing k norm and k gap on 2,500 different ( x ( j ) , x ( n ) ) pairs, and then report their mean values at each t .
To verify our assumption, we compare four different settings:
• In Rnd Backbone, we randomly initialize the whole network and train it on the downstream dataset; • In PT Backbone HP τ ep, we copy the pretrained backbone, attach a randomly initialized head, and then HP for τ epochs.
In the first panel in Figure 6, we see the training accuracy of all these settings converge to 100% (they also have similar validation accuracy).
In the second and third panel, we plot the change of k norm and k gap respectively.
It is obvious that compared with using random backbone, the NTK of a pretrained backbone indeed changes slow during the finetuning stage.
B E XPERIMENTAL S ETTINGS B. T OY CASE
The experiments of the toy setting appear in Figure 2 and Figure 3.
In the toy setting, we first pretrain a 4+1 (4 layers of backbone and 1 layer of task head) layers MLP on the full MNIST dataset (LeCun, 1998), with learning rate 10 − 3 , with cosine scheduler until convergence.
The hidden width is 128 and relu activation is applied for each layer in the backbone.
In the downstream task, we only consider the first three classes and randomly select 1000 samples for each class (i.e., 3*1000 samples in total).
We also apply random rotation and center crop augmentations to simulate the distribution shift in downstream task, like the right panel in Figure 7.
The downstream linear task head is then a 128*3 matrix (weights) and a 128*1 vector (bias).
To ensure a fine-grained observation of feature’s adaptation, we use a small constant learning rate ( 10 − 4 ) in all the experiments.
The FT epoch is 50, and the HP epoch ranges from 0 to 50,000 (50,000 is simulating the v ∞ hp case, usually train 0 to 1,024 epochs). [Figure]
B. P RACTICAL S ETTINGS – C LASSIFICATION
Besides the simple MNIST and MLP, we also verify our analysis in many real settings.
To demonstrate the generality of our findings, we conduct experiments on image classification, image seg- [Figure] [Figure]
𝑿
[Figure]
𝑻𝑿𝑿
[Figure]
mentation, and molecular graph classification tasks.
The qualitative trend is quite consistent across different tasks and networks.
Model Structure:
Specifically, we consider ResNet18/34/50 for the image classification task.
They all have 4 blocks, each containing 3 layers, as illustrated in Figure 9.
For the linear task head (or an MLP head), we directly concatenate a linear layer (or an MLP layer) to the backbone.
For the partial backbone method, several layers in the last block (e.g., L4.3, L4.2, L4.1) might be reinitialized and merged to the task head.
We also provide what is the backbone (all the blue squares) when self-supervisedlearning is considered, e.g., SimCLR, BYOL [Figure] and MoCo.
In short, after pre-training, the backbones under different PT tasks might have the same structure.
Pretraining tasks:
We pretrain the ResNet on CIFAR10 using common recipe ( 10 − 2 learning rate with a cosine scheduler (He et al., 2019), 5 ∗ 10 − 4 weight decay, simple augmentations like random flipping and cropping, etc.).
For the model pretrained on ImageNet, we directly download the checkpoints from open-source implementations.
Due to time and computing resource limitations, for vision tasks, we only consider MoCo (He et al., 2020), Byol (Grill et al., 2020), SimCLR (Chen et al., 2020a), and supervised classification tasks for pretraining.
For the graph dataset case, we consider a fundamental 5-layers graph convolutional network (GCN, (Kipf and Welling, 2017)), the hyper-parameters and other detailed settings can be found in our github repo.
The datasets applied in this paper are listed in Table 4.
Note that the backbone of image segmentation is the same as that in image classification.
The segmentation head is more complex, which will be discussed later.
Hyper-parameters for HP and FT:
In this paper, experiments with the same setting in different figures or tables share the same set of hyper-parameters.
Generally, for all the experiments, the batch size is 128, hidden layer width is 256 (in the MLP head case).
The input image size for CIFAR pretrained model is 32 × 32 while that for the IN1K pretrained model is 224 × 224 .
Hence in the corresponding experiments, we will first resize the input samples and then apply random cropping and random flipping augmentations.
For the HP phase, the τ ∗ is selected based on the validation performance, but the detailed learning rate and maximum τ might be different under different settings (as the dataset size are different).
For the
FT phase, we use a standard SGD with momentum ( β = 0 . 9 ).
The default learning rate is 10 − 3 and a cosine scheduler is applied (the maximum FT epoch also varies for different downstream datasets).
Note that we will early stop the FT process, hence the maximum FT epoch doesn’t influence the reported performance a lot.
For those dataset-dependent hyper-parameters, we summarize them as:
• For the insufficient pretrained backbone cases,
i
e., Res50-C10(Sup), Res18-C10(Sup), in Figure 4, Table 1 and other related experiments in the appendix, we set τ ∈ [0 , 200] , HP learning rate is 3 ∗ 10 − 2 , maximum FT epochs is 200 (usually converge less thanepochs), with usually a larger initial FT learning rate (e.g., 3 ∗ 10 − 3 ); • For the IN1K pretrained backbone (no matter by supervisory training, MoCo, Byol or SimCLR), which can be found in Figure 1, 2, 4, Table 1 and many figures in the appendix, we usually set τ ∈ [0 , 200] , HP learning rate as 10 − 2 , maximum FT epochs as 200, and a relative small FT learning rate (e.g., 3 ∗ 10 − 4 ); • For any experiments whose downstream tasks are Flowers, Cars (the dataset is small while number of classes is large), we consider to increase the HP learning rate to 5 ∗ 10 − 2 , and the maximum FT epochs to 1000 (usually converge less than 200 epochs); • For any experiments whose downstream task is CIFAR (the dataset is bigger than STL), we consider to increase the HP learning rate to 5 ∗ 10 − 3 , and the maximum FT epochs to(usually converge less than 50 epochs);
In summary, a general principle for these hyper-parameter selections is that the maximum τ should make the HP training accuracy converge (increase slowly for several consecutive epochs), while the finetuning epochs should make the training accuracy converge when τ = 0 .
For example, if the downstream task is CIFAR10, which has 50,000 training samples and only 10 classes, finetuneepochs is enough.
But CIFAR100, which also has 50,000 training samples but with more classes, 50 finetuning epochs are required (we set the maximum FT epochs as 100 for both CIFAR10 and
CIFAR100).
For Flowers102, which only has 6,149 training samples but 102 classes, we set the maximum FT epochs as 1,000 epochs.
Anyway, we find the proposed trend is quite robust to these hyper-parameters.
The detailed settings can be found in our code base.
B. P RACTICAL S ETTINGS – I MAGE S EGMENTATION
In addition to the image classification tasks, we also conduct experiments on image segmentation tasks to verify the robustness of our analysis on various tasks.
An image segmentation task is a task where a model is trained to make class predictions at each pixel of an image given an input image.
We train and test segmentation models on PASCAL VOC (Everingham et al., 2015) dataset, one of the most popular image segmentation datasets.
An example of input image and segmentation label are provided in Figure 10.
Model Structure:
For image segmentation tasks, we employ DeeplabV3 (Chen et al., 2017b) with ResNet50 as a backbone.
As with the classification tasks, ResNet50 consists of 4 blocks, each of which containslayers.
But unlike the classification tasks, kernel strides for convolutional layers are adjust to have # train # test # class Comments MNIST 60, 10,
Toy downstream use a subset, as in Figure 7.
STL 5, 8,
Coates et al.
(2011) CIFAR 50, 10,
Krizhevsky et al.
(2009) CIFAR 50, 10,
Krizhevsky et al.
Flowers 6, 2,
Nilsback and Zisserman (2008) StanfordCars 8, 8,
Krause et al.
Dom-real < 20, < 4, 200 (345)
A subset, original 345 classes and more samples.
Dom-sketch < 20, < 4, 200 (345)
Less sensitive to color and texture.
Dom-quick < 20, < 4, 200 (345)
Only contains simple curves.
(Peng et al., 2019)
ImageNet-1K 1,281, –
Deng et al.
(2009) ogbg-moltox 6, 1,
Wu et al.
(2018) ogbg-molhiv 32, 8, regression
Hu et al.
[Table]
Hu et al.
[Figure]
the higher spatial resolution for the output features of the backbone; the spatial resolution of the features is 7 × 7 for the classification tasks whereas it is 17 × 17 for the segmentation tasks.
The features extracted from the ResNet50 backbone are passed to a segmentation head, which consists of Atrous Spatial Pyramid Pooling (ASPP)
(Chen et al., 2017a) layers followed by a couple of convolutional layers.
ASPP is a stack of five regular and dilated convolutional layers which provide the features with various size of receptive fields.
The spatial resolution of the outputs of the segmentation head stays in 17 × 17 , which are then upsampled to the original input image size through bilinear interpolation.
Here, we use the whole segmentation head as the task head.
Hyper-parameters for HP and FT: The ResNet50 backbone is pretrained on ImageNet dataset.
We set τ ∈ [0 , 200] , HP learning rate as 0 . 3 , maximum FT epochs as 200 (usually converge less than 50 epochs).
Also, we use batch size of 16 , and a SGD optimizer with momention ( β = 0 . 9 ) but without weight decay nor learning rate scheduler.
C M ORE EXPERIMENTS C. V ERIFICATION ON THE CHANGE OF Z
In Equation (2), we decompose the learning dynamics of z t into three parts: relatively stable kernel, direction, and energy.
Based on this, we can expect a larger adaptation when the initial energy is large.
Furthermore, to get a more precise description of z t ’s adaptation, we analyze a simple overparameterized linear model in Appendix D.
We show that as the initial prediction gap (i.e., energy) increases, the resulting z t will first be stretched with a small direction change (compared with the original z 0 ).
If the initial energy keeps increasing, the resulting z t might be more and more dissimilar to the original one.
We use four distance related quantity to describe this trend, and verify it under many different settings.
In short, the trends are: 1.)
more energy leads to larger ∥ z t − z
∥and smaller cos( z t , z 0 ) ; 2.)
z
⊤ t z 0 and ∥
z t ∥ 22 has a quadratic shape when energy decrease.
Figureprovides an example of how z t changes when different τ is chosen. draw the
D projections (the first two components in PCA) of 100 randomly selected z ( n ) keeping the eigen-directions of this PCA, we project z ( n ) t of different t and encode t by transparency.
The converged z is represented by a cross marker with full transparency.
It is clear that FT using v rnd makes z change a lot, while v 50000 hp doesn’t change z too much.
[Figure]
Before diving deep into these examples, we provide some general commands.
First, we do not expect the experimental results to align perfectly with the examples in Figure 3, because too many designs (like dropout, data augmentation, SGD noise, locking in strange local minima, learning rate schedule, etc.) can influence the shape of the curve, hence we believe observing specific trends in many cases is already a strong support for our analysis.
Second, one might be curious about why in some cases, the ranges between z ⊤ t z 0 and cos( z t , z 0 ) are so small.
Hence in Figure 12, we provide an illustration of why this happens: as the dot product measures two long vectors (remember they have the same origin), a small change of z ⊤ t z 0 can make a big difference.
A large range of z ⊤ t zand cos( z t , z 0 ) is a sign that z t already moves to another basin, which is the strong-case in Figure 3.
Third, in Appendix D, the change of z is linked to q 0 − e y under an MSE loss, while in these experiments, we use τ as our x-axis and consider a cross-entropy loss.
Combining with the NTK approximation applied in this ideal model, the experimental trends might not be exactly the same with the theoretical ones.
Then, in Figure 13, Figure 14 and Figure 15, we demonstrate the trend when the model is pretrained on different datasets, on different tasks, and when the downstream tasks are different, respectively.
Besides the general trends, i.e., decreasing ∥ z t − z
∥ 22 , increasing cos( z t , z 0 ) , quadratic z ⊤ t z 0 and ∥ z t ∥ 22 , we can observe another interesting phenomenon when the model is pretrained using different tasks.
See Figure 14, in Byol and SimCLR, sometimes we might get a very big z ⊤ t z 0 (together with a very small cos( z t , z 0 ) ), which makes it hard to observe any quadratic trend in some metrics.
But in the supervised pretraining case (no matter what dataset we use in pretrain and downstream tasks), we never observe such a phenomenon – the ∥ z t − z
∥ 22 always changes in a relatively small range.
We speculate that in the first several updates during head probing, the features first adapt to the downstream task (i.e., classification), then gradually adapt to the downstream dataset distribution (e.g., STL, Flowers, etc.).
So a randomly initialized head might influence more when the pretrain task is different from the downstream one.
The detailed mechanism of adaptation of task and data distribution might be more complex than we expect, so we left this for the future work.
C. I NFLUENCE OF THE BACKBONE AND THE TASK - HEAD CAPACITY
In Section 4.3, we mention a case where the task head is more complex.
For example, the head can be a two-layers MLP rather than one linear layer, or we can only copy part of the pretrained model as our backbone.
However, these methods and tricks are quite complex and might not necessarily enhance the performance, hence we put some results and discussions here.
Remember that our analysis of how energy influences the features adaptation is still valid in these cases, as illustrated by the distance related metrics in the tables.
From Table 5, 6, 7, and 8, we can draw the following three conclusions.
First, the two-layers MLP design will surely increase the training accuracy after head probing, and hence make the features adapt less (but this might not enhance the downstream performance).
Second, we compare the downstream performance when some layers of the backbone are reinitialized as part of the task head.
The title “+L4.3” means we take the last layer (i.e., the 3-rd layer) of the last block (i.e., the 4-th block) in a ResNet50 out, and treat it as part of the task head.
The title “+L4.2” and “+L4.1” means we continue treating the 2-nd or the 1-st layer as the backbone.
Under such a setting, the “+L4.1” case will inherit the least amount of information from the pretrained model, and at the same time, have the biggest task head.
For the results in these tables, we train the head until convergence in head probing phase, then fine tune the whole network together.
Hence we see the “HP-train-acc” value all increase to roughly 100%.
However, the best setting differs when different downstream datasets are considered.
When the pretrained features are good (i.e., trained using IN1K), we see the
Domain-real dataset needs less adaptation and usually performs the best in the baseline case (i.e., copy all parameters from the pre-trained model).
But Domain-quick dataset, which only contains some black-and-white lines, prefers the “+L4.1” setting.
In other words, preserving the features captured by the earlier layers of the pretrained model is beneficial.
When the pretrained features are bad, like the CIFAR-pretraining case in Table 7, all these datasets prefer throwing away the later layers.
In summary, although the correlation between features adaptation and downstream task performance under different settings is quite complicated, our analysis of energy can still explain some phenomena well.
[Table]
Real
Sketch
Quick
Real
[Table]
Quick
[Table]
[Figure]
[Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure]
[Table]
D A NALYZE LINEAR OVERPARAMETERIZATION PROBLEM D. F ORMALIZE THE CHANGE OF REPRESENTATIONS
Section 4 provides some intuitive explanations of how feature extractor f ( x ; B ) changes given different g ( z ; v ) , which are well supported by the experimental results.
To provide more insights,
we
[Table]
analyze the change of features (i.e., z ) in a simplified overparameterization problem, i.e., the one provided in Section 3 and in Kumar et al.
Under some mild assumptions and approximations, we provide an overview of how the norm and direction of z changes under different choice of v 0 .
We show that the v 0 satisfying q
=
Y or q 0 =
Y are two critical points in feature adaptation.
We first rewrite Equation (1) in a non-matrix form:
[Table]
[Table]
[Equation]
(24) where y ∈ R as we are considering a regression problem (or a classification problem using MSE loss).
We use the subscript to represent the time step, e.g., q 0 , z 0 , v 0 and B 0 are output, feature, head parameters and backbone parameters before finetuning.
Similarly, q t , z t , v t and B t are the corresponding values after finetuned t steps.
Note that q ( X ) and z ( X ) are functions of N input samples X ∈ R N × d , but we omit it for simplicity.
We use lowercase letters to represent the nth element of the vector.
Before discussing a specific element, we will clarify whether we are discussing the initialized case or the finetuned case.
For example, specifying q t , we use q n = f ( x n ) =
( f ( X ))
n to represent the prediction of x n after finetuning.
Remember the goal of this paper is finding a suitable way to select task head, i.e., v 0 , given the pretrained feature extractor, i.e., B 0 .
Depending on the discrepancy between the pretraining task and the downstream task, we might expect z t change differently after finetuning.
In other words, we care about the expected change of z t compared to z 0 , i.e., E x ∼D [ d ( z t , z 0 )] , where D is the data distribution of the downstream task.
Depending on what distance measurement (i.e., d ( · , · ) ) we choose, there are three different metrics:
• Euclidean: ¯ d euc ≜ E x ∼D (cid:2) ∥
z t − z 0 ∥ 22 (cid:3) , • Dot product: ¯ d dot ≜ E x ∼D (cid:2)
z ⊤ t z 0 (cid:3) , • Cosine:
¯ d cos ≜ E x ∼D (cid:104)
[Equation]
⊤ t[-z 
-]
[Equation]
∥
[Equation]
t ∥ [Equation] ·∥ [Equation]
∥
[Equation]
We start from the Euclidean case:
[Equation]
where M ≜ tr ( xx ⊤ ) or M ≜ ∥ xx ⊤ ∥ op ) .
Similarly, we can get the expressions of ¯ d dot and ∥
z t ∥ 22 , which are building blocks of ¯ d cos :
[Equation]
D. C RITICAL P OINTS
Our goal is to find a good initialized task head, i.e., v 0 , which can lead to better downstream performance.
Instead of directly linking v 0 to the expected risk, which is a common practice for generalization analysis 5 , we consider this problem from another indirect way.
Specifically, we assume the SGD algorithm with appropriate regularization in FT stage can find a good optimum for the learning task.
What we care more about is whether the features learned from the pretraining stage adapts well to the downstream task.
We believe that if the features are properly adapted to the new environment, the model has more potential to generalize better.
Hence in this part, considering the aforementioned three distance metrics as the target functions, we formalize how v 0 influence them.
From Equation (25), (26) and (27), we find that the behavior of tr ( B ⊤ 0 B t ) and tr ( B ⊤ t B t ) are the keys.
In these two terms, B 0 is given and cannot change, while B t is determined by the choice of v 0 .
One way to link these two quantities is using Lemma A.4 in Kumar et al.
(2022) or theorem 2.
in Du et al.
[Equation]
[Equation] however, the expression of v t is still hard to obtain (but collecting and visualizing v t is much more cheaper than B t ).
We left this direction for our future work.
In this paper, we analyze the problem in the NTK regime (this is the main assumption of our analysis, which can cause discrepancies between the theory and experiments).
Closed-form of parameters under NTK approximation:
To get more insights, we approximate the behavior of this model in the NTK regime, in which the converged parameters can be analytically calculated.
Specifically, be applying Equation (8) in Lee et al.
and assuming t → ∞ , we can have:
[Equation]
where q 0 ∈ R N × 1 is the model’s prediction on N training samples, i.e., X , and θ t , θ 0 ∈ R ( d +1) ∗ h ×are the stacked parameters.
Without loss of generality, the first d ∗ h parameters in θ come from B and the last h parameters come from v .
The K 0 = ∇ θ q 0 · ( ∇ θ q 0 )
⊤ ∈ R N × N here is the empirical NTK on X .
Specifically, by stacking the paramters, we can calculate each elements in this kernel as:
[Equation]
where v i and x i is the i -th element in v and x respectively.
Then, the matrix form of emperical NTK is K [Equation]
=
[Equation]
With the help of Equation (29), we can get the closed-form expression of b t ≜ vec ( B t ) (i.e., the vectorization of matrix B t ):
[Equation]
In Equation (31), we know q 0 = X B ⊤ 0 v 0 .
The term ∇ b q 0 also depends on v 0 .
As controlling q 0 , i.e., the model’s prediction on training samples, is easier in practice (we can directly observe the training loss or training accuracy), we will consider q 0 as the optimizing variable in the rest of the paper.
The following lemma can link ∇ b q 0 to q 0 :
Lemma 1.
For b 0 ∈ R h ∗ d × 1 and ∇ b q 0 ∈ R N × h ∗ d , we have ∇ b q 0 · b 0 = q 0 .
Proof.
We check that equation elementwise.
For the n -th element in the RHS, we have:
[Equation]
) where q n is the n -th element of q 0 , x i is the i -th element of x n , and v j is the j -th element of v 0 .
For the LHS, the n -th value is e ⊤ n ∇ b q 0 · b 0 , where e ⊤ n is a one-hot row vector selecting the n -th row of ∇ b q 0 .
From the definition, we know e ⊤ n ∇ b q 0 = vec ( v 0 x ⊤ n ) ⊤ =
[ x 1 v ⊤ 0 , x 2 v ⊤ 0 , ..., x d v ⊤ 0 ] , which is a long row vector.
Then, the n -th element in the LHS should be:
[Equation]
where the last equation holds the rule of stacking elements in a matrix.
As each elements of the two sides are equal, the lemma holds.
Critical points of tr ( B ⊤ 0 B t ) and tr ( B ⊤ t B t ) :
Lemma 2.
q ∗ 0 =[-1 
-]
12 Y is a maximum of tr ( B ⊤ 0 B t ) .
Proof.
By definition, we have:
[Equation]
where we call C 1 = K − 1 0 for convenience.
By taking first derivative to q 0 and let is equal zero, we know q ∗ 0 =
12 Y is a critical point.
As the second derivative is − 2 C 1 and C 1 is positive definite, this critical point is a maximum.
Lemma 3.
q ∗
= αY, α ∈ (0 , 0 . 5) is a critical point (usually maximum) of tr ( B ⊤ t B t ) .
Here α = (cid:0) I N × N + ( C 2 − 2 C 1 ) − 1 C 1 (cid:1) Y , where
C 2 = K − 1 0 ˜ K 0 K − 1and ˜ K
= X B ⊤ 0 B 0 X T
=
( ∇ v q 0 )( ∇ v q 0 )
⊤ is the NTK when fixing the parameters of the backbone.
[Figure]
Proof.
Similar to Lemma [Equation] we can write tr ( B ⊤ t B t ) as b ⊤ t b t and substitute Equation (31):
[Equation]
By taking first derivative and letting it equal zero
[Equation] assuming ( C 2 − 2 C 1 ) is invertible, we know q ∗ {+
+} = (cid:0) I N × N + ( C 2 − 2 C 1 ) − 1 C 1 (cid:1)
Y is a critical point.
The second derivative is 2 C 2 − 4 C 1 , which is usually positive definite in our settings (explain later).
To get more insights, we can look deeper into the critical point mentioned in Lemma 3.
Following the definition of C 1 and C 2 , we can have:
[Equation]
where K − 1 0 ˜ K 0 is the key of understanding this term.
The exact form of this expression is hard to obtain, but as all the NTK or covariance matrices mentioned here are sysmetric, we can compare the trace of them to get some insights.
By definition, tr ( ˜ K 0 ) = tr ( X B ⊤ 0 B 0 X ⊤ ) = tr ( B ⊤ 0 B 0 X ⊤ X ) .
If each dimensions of the data samples are independent, then X ⊤ X
≈ I d × d , and tr ( ˜ K 0 ) ≈
tr ( b ⊤ 0 b 0 ) =
∥ b 0 ∥ 22 .
Then similarly, tr ( K 0 ) ≈ ∥ b 0 ∥ 22 + ∥ v 0 ∥ 22 .
Thus the behavior of term K − 1 0 ˜ K 0 can be described by
[Equation]
• if ∥ b 0 ∥ [-
-]22 ≫ ∥ v 0 ∥ 22 , which might happen as the backbone contains more parameters than the head, K − 1 0 ˜ K
≈ I N × N , • if ∥ b 0 ∥ 22 ≪ ∥ v 0 ∥ 22 , which might happen if we split the network in the earlier layer, K − 1 0 ˜ K 0 ≈ 0 · I N × N .
But in either case, the negative of the second derivative of tr ( B ⊤ t B t ) , i.e., Equation (36), is likely to be positive definite, which means the critical point in Lemma 3 is usually a maximum.
Another interesting fact of Lemma 3 is q ∗ 0 under the aforementioned two extreme conditions.
After some calculation, we
euc , ¯ d
[-[Figure] − 𝐳 
-]dot
and ¯ d
[-𝐳 𝐭T 𝐳 Y Y/2 𝛼Y𝐛 𝟎 22 +𝑎𝐛 𝟎 2𝐳 𝐭 
 Y Y/2 𝛼Y𝐛 𝟎 2 2 +𝑎𝐛 
-]cos :
Recall the definitions of different distance metrics we care about.
First, ¯ d dot is proportional to z ⊤ t z 0 , hence its shape would like the second panel in Figure 16.
For ¯ d euc , we know it is proportional to tr ( B ⊤ t B t ) − 2 tr ( B ⊤ 0 B t )
+ tr ( B ⊤ 0 B 0 ) .
Substituting results in Equation (34) and Equation (35), we have:
[Equation]
Obviously, q ∗
= Y is the minimum of ¯ d euc , as depicted in the first panel in Figure 16.
The shape of ¯ d cos is hard to obtain.
However in the last panel of Figure 16, we demonstrate ∥ z t ∥ 22 , its denominator, to assist our further analysis.
D. L EARNING DYNAMICS OF THE REPRESENTATION
[Equation]
[Table]
[Equation] [Table]
In this subsection, we will put everything together to provide an overview of how z t changes compared with z 0 .
Remember when q 0 =
Y , the gradient of B 0 is zero, hence z t =
z
(Kumar et al., 2022).
When q 0 moves linearly from Y to 0 , we have the following three phases (see Table 9):
• Y → 12 Y : z t lengthen its norm with a slightly change in the direction, hence ¯ d dot also increase; • 12 Y → αY : the norm of z t keeps increasing, but its direction drastically changes in this phase, which makes ¯ d dot decrease; • αY → 0 : the norm of z t begins to decrease and the angle between z t and z 0 keeps increasing, which makes z t changes a lot.