Efficient Transfer Learning from Arbitrary Pre-Trained Models

We appreciate your constructive feedback and recognition for the importance of our work. To address your questions, we run several new experiments to show that AFT indeed learns to downweight less useful features and is the best-performing method even under data scarcity. The new results are included in the updated version of the paper. We hope you can consider raising your score in light of our response.

On learning to downweight less useful features. By allowing the model to freely weigh the pre-trained features via $\rho$ and ignore an arbitrary set of features without incurring any penalty, AFT incentivizes the model to prioritize learning the most useful pre-trained features, as they are most helpful for fitting the downstream data, and downweight the less useful ones. For this reweighting to be useful, the learned $\rho$ need not be sparse as long as it assigns sufficiently higher weights to the most relevant features.

We provide two pieces of evidence that AFT indeed learns to downweight less useful features. First, we show that AFT learns to downweight the features from a pre-trained model that is less useful for the downstream task. In Figure 4(a), when we use BiT, DINO, and CLIP as the pre-trained models, AFT successfully learns to upweight CLIP and DINO features and downweight BiT features, which has considerably lower linear probe accuracy. Second, we run a new experiment where we append random noise features to DINO features and find that AFT learns to significantly downweights the noise features (Figure 4b), and thereby maintains its performance as we increase the number of noise features (Figure 4c), whereas KD performance quickly degrades.

On data scarcity. Many of our datasets are already very small in size: Oxford flowers (1K), Oxford pets (3.6K), RTE (2.5K), MRPC (3.7K), and AFT leads to significant gains on these datasets (Figure 2(b) and Figure 3(b)). Furthermore, in Figure 5(b), we present an additional experiment where we subsample the training data down to 10 examples per class on CIFAR-100, and 100 examples per class on BoolQ. We show that AFT continues to outperform competing methods even in extreme data-scarce settings.

On LoRA. LoRA is a method to efficiently fine-tune the original pre-trained model, but does not transfer to a smaller model and therefore still requires running expensive inferences with the original pre-trained model. As a result, LoRA is not a viable alternative to AFT or KD as a method to efficiently transfer from larger models into smaller ones. On the other hand, LoRA is entirely complementary to AFT as it can be used together during AFT fine-tuning to further reduce training cost.

On arriving at the final objective. We included a detailed derivation of the objective starting with the bound on the mutual information to highlight the conceptual difference between AFT and KD. While the final objective is indeed a kernelized version of the L2 objective in Eq 9 (now Eq 4), it is fundamentally different from the similar-looking L2 objective in Eq 10 (now Eq 5) which KD optimizes. However, we appreciate the suggestion for a more concise presentation and have compressed the derivation.

On clarifying our notations The $d_\psi \times d_\psi$ matrix $\rho$ is parameterized by the $d_\psi$ dimensional vector $s$ via $\rho = \mathrm{diag}(\sigma(s)),$ i.e. $\rho_{ii} = \sigma(s_i),$ where $\sigma$ is the sigmoid function. We have clarified this in the text.

We did not assume that the labels $Y$ are linear regression labels. The equation $Y=W \Phi$ states that the output of the model is a linear transformation of its last layer features $\Phi,$ which holds for all neural networks and in both regression and classification tasks (where $Y$ would be the logits). Similarly, Algorithm 1 is applicable to both regression and classification tasks and we have updated the text to show that the loss $L$ is any loss function specified for the downstream task.

Thank you for responding! We’d like to clarify our answer to this question. If multiple features are simultaneously relevant but correlated, the weights learned by AFT indeed may not be sparse, nor have we claimed that they would be.

However, assigning non-zero weights to correlated features should not harm AFT’s performance. In fact, we can show that if non-zero weights $\rho_a$ and $\rho_b$ are assigned to two duplicated features $\psi_a$ and $\psi_b$ for $a\neq b,$ the effect is exactly identical to having assigned a sparse $\rho'$ where $\rho'_a = \sqrt{\rho^2_a + \rho^2_b}$, $\rho'_b = 0$, and $\rho'_i = \rho_i$ for $i\notin\{a,b\}.$

To prove this, it suffices to show that the target kernel $K_{\rho\Psi}(x,x') = (\rho \psi(x))^\top (\rho \psi(x'))$ in Eq 6 constructed by AFT with non-sparse weights $\rho$ is identical to the kernel constructed by AFT with sparse weights $\rho':$ $K_{\rho'\Psi}(x,x') = (\rho' \psi(x))^\top (\rho' \psi(x')).$

This follows from direct calculation: $K_{\rho\Psi}(x,x') = (\rho \psi(x))^\top (\rho \psi(x')) = \sum_{i} \rho^2_i \psi_i(x)\psi_i(x') = \sum_{i\notin\{a,b\}} \rho^2_i \psi_i(x)\psi_i(x') + \rho^2_a \psi_a(x)\psi_a(x') + \rho^2_b \psi_b(x)\psi_b(x') = \sum_{i\notin\{a,b\}} \rho^2_i \psi_i(x)\psi_i(x') + (\rho^2_a + \rho^2_b) \psi_a(x)\psi_a(x')$

By definition of $\rho'$, the above is equal to $\sum_{i\notin\{a,b\}} \rho^2_i \psi_i(x)\psi_i(x') + \rho'^2_a \psi_a(x)\psi_a(x') = (\rho' \psi(x))^\top (\rho' \psi(x')) = K_{\rho'\Psi}(x,x').$ The argument can be straightforwardly extended to an arbitrary number of duplicated features. Therefore, due to using the kernelized objective, assigning non-zero weights to duplicated features is automatically equivalent to assigning sparse weights in AFT.

We hope you find this answer satisfactory and will consider revising your score!

We appreciate your constructive feedback. To address your questions, we have significantly expanded our evaluation to both verify the claimed properties of AFT and evaluate its performance on a wider range of datasets. The new results are included in the updated version of the paper. We hope you can consider raising your score in light of our response.

On further evaluation. Following your suggestion, we now include 2 additional vision datasets used in CLIP evaluation (DTD and Food-101) and 6 additional language datasets from the GLUE benchmark, bringing the total number of datasets to 15. We also report the averaged results in Figure 2(a) and Figure 3(a), showing that AFT significantly outperforms competing methods on average.

On the difference between AFT and KD. Conceptually, the main difference between AFT and KD is that KD aims to transfer all pre-trained features to the downstream model, while AFT aims to transfer only the features relevant to the downstream task and is therefore more suitable for transfer learning. As shown in Ahn et al. (2019), the KD objective minimizes a bound on $I(\Psi; X|\Phi)$ (equivalently $H(\Psi|\Phi)$) where $\Psi$ is the pre-trained feature, $\Phi$ is the downstream feature, and $X$ is the input. This objective penalizes any information in the pre-trained features that is not in the downstream features. By contrast, AFT optimizes a bound on $I(\Phi; X|\Psi),$ which only penalizes new information in the downstream features that is not contained in the pre-trained features. In particular, the latter is zero as long as the downstream features are a subset of the pre-trained features, where the opposite is true for the former. As a result, this conceptual difference manifests in the final objectives as whether we train the pre-trained features to be predictable from the downstream features (KD) or train the downstream features to be predictable from the pre-trained features (AFT), as you correctly pointed out.

Following your suggestion, we have added Figure 4 to show that AFT indeed learns to weigh the pre-trained features differently. For example, when we use BiT, DINO, and CLIP as the pre-trained models, AFT upweights CLIP and DINO features and downweights BiT features, which has considerably lower linear probe accuracy (Figure 2d). We also run an experiment where we append random noise features to DINO features and find that AFT learns to significantly downweights the noise features (Figure 4b), and thereby maintains its performance as we increase the number of noise features (Figure 4c), whereas KD performance quickly degrades.

On choosing the starting point for AFT. As mentioned above, Eq (7) and Eq (8) (now Eq (4) and Eq (5)) are fundamentally different despite their similar forms. The former is conceptually more well-suited for transfer learning and allows selectively transferring the most relevant features. While the ablation study shows that directly optimizing Eq (7) (AFT w/o kernel) is not effective, we show that its kernelized version (AFT) is the best performing method. By contrast, kernelizing Eq (8) is similar to AFT w/o $\rho$ (not allowing weighting the pre-trained features), which we find is less effective.

On clarifying our definitions. (1) Thank you for catching the typo in the kernel definition. We have fixed it. In Eq 9 (now Eq 6), $k_\Phi(X, X')$ follows the same definition given in the paragraph directly above, namely $\phi(X)\phi(X'),$ where $X$ and $X'$ each denote a single input.

(2) Under the kernel formulation, $\rho$ no longer needs to map the pre-trained features to the shape of the downstream features because the kernel enables comparison between features of different shapes. We have added a sentence to clarify this.

(3) We do not need to actually optimize for an orthogonal $U,$ as a result of using the kernel. Two feature maps $\phi(\cdot)$ and $\phi'(\cdot)$ have the exact same kernel even if they differ by an orthogonal transformation $U^*.$ By contrast, we do need to optimize for the right $U^*$ if we measure their $\ell_2$ distance as in Eq 4.

(4) We updated the text to clarify that the mini-batch loss $\hat{L}(\theta)$ is computed as $L(\theta, Y_{\mathrm{batch}}, \hat{Y}_{\mathrm{batch}})$ where $L$ is some given loss function on the downstream task (e.g. cross-entropy loss).

[1] Ahn et al. (2019). Variational Information Distillation for Knowledge Transfer

Thank you for your thoughtful and constructive feedback. To address your questions, we have run several new experiments to verify and investigate the properties of AFT, including ablation studies on the optimization of $\rho$ and visualization of learned $\rho$. We have also significantly expanded the scope of our evaluation to include 8 more datasets. The new results are included in the updated version of the paper. We hope you can consider raising your score in light of our response.

On further evaluation. Following your suggestion, we include 2 additional vision datasets and 6 additional language datasets from the GLUE benchmark, bringing the total number of datasets to 15. We also report the averaged results in Figure 2(a) and Figure 3(a), showing that AFT significantly outperforms competing methods on average.

On rejecting less useful features. In Figure 2(c) (previously Figure 2(b)), transferring convolutional features from a BiT into a ViT downstream model did in fact improve its performance compared to the ViT only standard transfer learning baseline (STL). It's true that transferring from BiT + DINO + CLIP is slightly worse than DINO + CLIP only. However, AFT is successful at upweighting the more useful features even in this case since transfer from BiT + DINO + CLIP is much better than from BiT alone. Indeed, we verify this in Figure 4(a) by visualizing the learned $\rho,$ which are much smaller for BiT features than for DINO and CLIP features when transferring from all three models. To further test AFT's ability to downweight useless features, we run a new experiment where we append random noise features to DINO features and find that AFT learns to significantly downweights the noise features (Figure 4b), and thereby maintains its performance as we increase the number of noise features (Figure 4c) to up to 2048, whereas KD performance quickly degrades.

On optimizing $\rho$ more often than $\theta$. By using the kernel formulation and only learning a diagonal $\rho,$ we have considerably reduced the difficulty of optimizing $\rho.$ Instead of optimizing a dense $d_\psi \times d_\psi$ matrix, we only optimize a $d_\psi$ dimensional vector. Furthermore, the bound in Eq 2 is valid even if $\rho$ is not optimal for a given $\theta.$ Therefore, to maximize the number of updates to $\theta$ given a fixed compute budget, we optimize them jointly, exactly analogous to how the encoder and decoder are jointly optimized in VAEs. However, following your suggestion, we added an ablation experiment where we update $\rho$ 5 times per $\theta$ update, and find that it does not improve performance even with the number of updates to $\theta$ fixed (Figure 5(a)).

On the benefit of invariance under orthogonal transformation. Invariance of the kernel under orthogonal transformation of the features is helpful because it reduces the number of parameters to learn. By only learning a diagonal $\rho$ with $d_\psi$ parameters, we can effectively search for all transformations $\rho' = U\rho$ for any orthogonal matrix $U$ without having to learn $d_\psi^2$ parameters.

On using a different kernel. We have added an ablation experiment where we use a Gaussian (RBF) kernel, and find that it achieves similar but slightly worse performance (Figure 5(a)).

On AFT performance vs linear probe performance. We plotted the linear probe accuracy vs test accuracy on different scales for the convenience of visualization because the former has a much larger range than the latter. It’s true that for sufficiently large and strong pre-trained models, linear probe can be better than AFT (and all other methods). However, the linear probe is expensive to deploy, as it requires inference with the original pre-trained model, the very problem that AFT aims to solve, and is therefore not a viable alternative. For example, the linear probe accuracy of ViT-L CLIP roughly matches AFT accuracy when transferred to ViT-S on CIFAR-100. But ViT-L CLIP has 428M parameters, 20 times larger than ViT-S. We have added this discussion to the paper.

On combining with model/feature selection techniques. We agree that combining AFT with model selection techniques is a promising direction, and have included a discussion and provided references to the works you mentioned. However, not explicitly performing model selection is not a weakness of AFT. Model selection and AFT are entirely complementary, as AFT solves the problem of how to transfer from any pre-trained model(s), while model selection techniques solve the problem of how to select the best pre-trained models to perform transfer learning from. Similarly, AFT can be combined with feature selection techniques, but not doing so is not necessarily a weakness of AFT. First, transferring only a sparse subset of features is a strong inductive bias that may not be appropriate for all tasks. Second, AFT already learns to downweight the less useful features, as shown in Figure 4.

Thank you for your thoughtful review of our work. Following your suggestion, we have added comparisons to a stronger baseline (B-Tuning in ref [4] you provided) and significantly expanded the scope of our evaluation to include 8 more datasets. The new results are included in the updated paper. We put a significant effort into these new experiments, inspired by your comments. We hope you can consider raising your score in light of our response. We respond to your questions below.

On comparison with related work. Following your suggestion, we have added B-Tuning, the method proposed in [4], as a stronger baseline throughout our experiments. In Figure 2(a) and 3(a), we show that AFT significantly outperforms B-Tuning on average across all 6 vision datasets and 8 language datasets. Similar to KD, we find B-Tuning benefits considerably less than AFT from better pre-trained models, as shown in Figure 2(c) and Figure 3(c). We also added a discussion of other related works to the paper.

On expanding our evaluation. Following your suggestion, we include 2 additional vision datasets and 6 additional language datasets, bringing the total number of datasets to 15, covering a comprehensive set of standard benchmarks for transfer learning. We report the averaged results in Figure 2(a) and Figure 3(a), showing that AFT significantly outperforms competing methods.

On using larger downstream datasets. We've evaluated on 5 new downstream datasets that are larger than any dataset we previously used, including QQP (364K), MNLI (393K), QNLI (105K), Food101 (76K), and SST-2 (67.3K). AFT is the best performing method for all downstream models on 4 of these 5 datasets, verifying that AFT is effective even when the downstream datasets are large.

On novelty. We would like to clarify that AFT has significant novelty relative to KD methods. AFT and KD have fundamentally different goals: KD aims to transfer all information in the pre-trained (teacher) features into the downstream (student) features, irrespective of whether it is relevant to the downstream task. By contrast, AFT aims to transfer only a subset of the information in the pre-trained features and allows the model to freely select which features to transfer, reflecting the fact that not all information in the pre-trained features is expected to be relevant to the downstream task. To illustrate the extent of this difference, consider an extreme example where the downstream features are identically zero and thus contain no information from the pre-trained features. The regularization term in Eq 6 which AFT optimizes is at its minimum and exactly zero, by setting $\rho = 0,$ since the downstream features are perfectly predictable from pre-trained features. By contrast, the regularization term in Eq 5 which KD optimizes will be at its maximum.

The main novelty of AFT is in the specification of a novel prior for transfer learning that accurately reflects our belief that not all the pre-trained features will be relevant to the downstream task -- they are only a collection of features in which some task-relevant features are likely to be included. The specific act of transforming the pre-trained features instead of downstream features is a consequence of this distinct prior, and not the fundamental reason why AFT is novel.

Furthermore, while Equation 4 in the reference you provided [1] presents a generic KD loss where both student and teacher features are allowed to be transformed, none of the specific methods mentioned there optimize for the selective transfer of knowledge by allowing the student to learn to freely ignore some of the teacher features in the manner that AFT does. Therefore, we believe that AFT is a novel method that is distinct from existing feature-based KD methods.

On AFR vs linear probe on pre-trained features. Simply using a linear model with the pre-trained features is not a viable alternative, as it requires running expensive inference with the original pre-trained model, the very problem that AFT (and KD) aims to solve by transferring to much smaller models. For example, the linear probe accuracy of ViT-L CLIP roughly matches AFT accuracy when transferred to ViT-S on CIFAR-100. But ViT-L CLIP has 428M parameters, 20 times larger than ViT-S. We have added this discussion to the paper.