Knowledge Composition using Task Vectors with Learned Anisotropic Scaling

Authors: frederic z zhang
paul  albert
cristian  rodriguez-opazo
anton  van den hengel
ehsan  abbasnejad
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-02 05:26:17

Abstract: Pre-trained models produce strong generic representations that can be adapted via fine-tuning on specialised datasets. The learned weight difference relative to the pre-trained model, known as a task vector, characterises the direction and stride of fine-tuning that enables the model to capture these specialised representations. The significance of task vectors is such that simple arithmetic operations on them can be used to combine diverse representations from different domains. This paper builds on these properties of task vectors and aims to answer (1) whether components of task vectors, particularly parameter blocks, exhibit similar characteristics, and (2) how such blocks can be used to enhance knowledge composition and transfer. To this end, we introduce aTLAS, an algorithm that linearly combines parameter blocks with different learned coefficients, resulting in anisotropic scaling at the task vector level. We show that such linear combinations explicitly exploit the low intrinsic dimensionality of pre-trained models, with only a few coefficients being the learnable parameters. Furthermore, composition of parameter blocks enables modular learning that effectively leverages the already learned representations, thereby reducing the dependency on large amounts of data. We demonstrate the effectiveness of our method in task arithmetic, few-shot recognition and test-time adaptation, with supervised or unsupervised objectives. In particular, we show that (1) learned anisotropic scaling allows task vectors to be more disentangled, causing less interference in composition; (2) task vector composition excels with scarce or no labelled data and is less prone to domain shift, thus leading to better generalisability; (3) mixing the most informative parameter blocks across different task vectors prior to training can reduce the memory footprint and improve the flexibility of knowledge transfer. Moreover, we show the potential of aTLAS as a parameter-efficient fine-tuning method, particularly with less data, and demonstrate that it can be easily scaled up for higher performance.

Main paper: 
 introduction: 
 One practical advantage of neural networks is the fact that knowledge learned from a previous problem, in the form of network weights, can be transferred to solve other related problems. Commonly referred to as transfer learning [6, 73] , this technique is often applied when a model trained on a general-purpose dataset-ImageNet [52] for many years-is fine-tuned on other datasets to improve (1) , τ (2) has two parameter blocks, learning anisotropic scaling grants more flexibility when combining task vectors.performance on downstream problems. In the past, classification models [18, 53] have been used as the medium for such knowledge transfer, which played a crucial part in the success of detection and segmentation [7, 19, 51, [66] [67] [68] . In recent years, foundation models [4] trained on broad data, CLIP [47] particularly, have demonstrated strong performance on a multitude of tasks, even when applied in a zero-shot manner. Besides the conventional way of exploiting the knowledge in these models via fine-tuning, recent works [28, 44, 62] have presented more direct measures to manipulate the network weights. In particular, Ilharco et al. [28] showed that, a task vector, defined as the weight difference between a pre-trained and a fine-tuned model, can be used as a carrier of the task-specific knowledge learned via fine-tuning. As such, multiple task vectors, when combined with simple arithmetic, can form a multi-task model that largely retains its performance across all fine-tuning tasks. Linearisation techniques [44] , in addition, have been shown to further enhance this compositionality.Intrigued by this phenomenon, we investigate the potential of task vectors being knowledge carriers in this paper, by learning linear combinations of them (Figure 1a ) for various problems. In particular, parameter blocks, e.g., weights and biases, tend to encode different learned representations in different layers. We thus learn an independent scaling coefficient per block for more precise adjustments tailored to the unique roles of each parameter block. This results in anisotropic scaling of task vectors (Figure 1b ), and allows us to exploit their modularity in knowledge composition, granting higher controllability when steering the behaviours of a model for task arithmetic [28] .The potential applications of task vector composition extend beyond model editing. With the coefficients being the only learnable parameters, our method exploits the rich knowledge encapsulated in the task vectors by searching in a low-dimensional coefficient space. As a result, it is a competitive parameter-efficient fine-tuning (PEFT) method, and is particularly effective in cases where labelled data is scarce. This offers new opportunities for few-shot learning [34, 69] and test-time adaptation [35, 57] . Furthermore, for multi-purpose models such as CLIP [47] , variants of the model trained with different data sources or fine-tuned on different downstream tasks are often available [26] . These resources constitute a significant knowledge bank, with task vectors being the knowledge carrier. Many learning problems may be simplified to learning a combination of task vectors.Our primary contribution is a learning algorithm named aTLAS, wherein otherwise complex learning problems can be framed as learning linear combinations of task vectors. The algorithm is broadly applicable to optimising supervised and unsupervised objectives. Its effectiveness is demonstrated in task arithmetic, few-shot recognition, test-time adaptation and parameter-efficient fine-tuning, where we show that (1) learning linear combinations of task vectors directly exploits the low intrinsic dimensionality of pre-trained models [1, 33] , resulting in a small number of learnable parameters; (2) standard task vectors, otherwise inferior to linearised variants [44] in task arithmetic, can produce stronger multi-task models with learned anisotropic scaling; (3) aTLAS is effective in low-data regimes, and improves the accuracy of CLIP by 6.5 absolute points averaged over 22 datasets with unlabelled data; (4) aTLAS is complementary to previous few-shot adaptation methods, in that one third of the examples it improves upon are unique; (5) aTLAS as a few-shot learning method is less prone to domain shift, and achieves better generalisation on out-of-domain datasets; (6) the most informative parameter blocks from different task vectors can be mixed prior to training, allowing for flexible and efficient knowledge transfer under memory constraints; (7) aTLAS is a strong PEFT method when data is limited, and existing PEFT methods such as low-rank adaptations (LoRA) [23] can be seamlessly integrated into aTLAS to improve memory efficiency. 
models and task vectors: 
 As Ilharco et al. [28] demonstrated, task vectors exhibit many intriguing properties across a wide range of models, such as CLIP [47] , GPT-2 [46] and T5-based models [48] . To facilitate more in-depth experimentation and analysis, we focus on the CLIP model in this paper, due to its wide availability and manageable size. In particular, we follow previous practice [28, 44] and acquire task vectors by fine-tuning the image encoder, with the text representations frozen. This ensures that image encoders fine-tuned on different datasets produce features residing in the same representation space, through a common text encoder. The task vectors obtained from these fine-tuned encoders can thus be combined more effectively to form a unified multi-task model.Formally, denote the CLIP image encoder by f : X × Θ → Z, such that for input image x ∈ X and parameters θ ∈ Θ, z = f (x; θ) is the learned latent representation for the input image. Denote the weights of a pre-trained model by θ 0 , and the weights of its fine-tuned variant by θ i , i ∈ N + , where i indexes a dataset D i . We follow Ilharco et al. [28] and define a task vector as τ i = θ i -θ 0 . In addition, we investigate task vectors produced by linearised variants of the image encoder using the first-order Taylor expansion, g(x; θ) := f (x; θ 0 ) + (θ -θ 0 ) T ∇ θ f (x; θ 0 ).(1)Ortiz-Jiménez et al. [44] showed that, task vectors obtained from fine-tuning the linearised variants have low disentanglement errors, and exhibit strong compositional properties. 
learning task vector compositions: 
 Parameters in a neural network, depending on the depth of the layer, often have different significance.For instance, early layers in convolutional neural networks [18, 53] are known for extracting generic, low-level features, such as edges, corners, etc., while deeper layers produce features more specific to the task. We recognise the non-uniform impacts parameters at different layers can have, and do not perform isotropic scaling on task vectors. Instead, weights, biases and any other forms of parameterisation, which we collectively refer to as parameter blocks, will be scaled independently.Formally, denote a task vector with m parameter blocks by τ = τ (1) , . . . , τ (m) , where each parameter block τ (j) is vectorised, and round brackets denote column vector concatenation. We learn a block diagonal matrix Λ, parameterised asΛ =   λ (1) I (1) . . . 0 . . . . . . . . .EQUATIONwhere λ (j) ∈ R is a learnable coefficient; I (j) denotes an identity matrix with its number of columns matching the dimension of τ (j) ; and the superscript j ∈ N + indexes a parameter block. This results in anisotropic scaling of a task vector, that is,EQUATIONwhere the subscript i ∈ N + indexes a task vector. As such, assuming a supervised objective, finding the optimal composition of task vectors can be defined as the following optimisation problem arg min Λ1,...,ΛnEQUATIONwhere L is the loss function for a target task; n is the number of task vectors; y is the labels corresponding to inputs x; D t denotes a target dataset. The number of learnable parameters, as a result, is precisely mn, Let us denote the solution to the aforementioned optimisation problem by{Λ ⋆ i } n i=1 . In inference, model f (x, θ 0 + n i=1 Λ ⋆ i τ i )will be deployed, which incurs no additional computational cost compared to models trained in the conventional way. The accuracy is normalised by that of the fully fine-tuned model. Using task vectors to construct the projection matrix performs consistently better than using random bases on (a) MNIST [32] , (b) CIFAR100 [31] .In addition, we investigate the task vectors obtained from fine-tuning linearised variants of the model, i.e., g(x) in Eq. 1. Denote such task vectors by τ . The learning objective with linearised task vectors can be derived as follows arg min Λ1,...,ΛnEQUATIONA notable characteristic of aTLAS is its parameter efficiency. To offer more intuitions, we refer to previous findings [1, 33] that deep neural networks often produce solutions residing in a subspace with much lower intrinsic dimensionality. This is measured by finding a minimum number of d parameters, such that learning these parameters ( θ ∈ R d ) leads to approximately the same performance as optimising in the full parameter space (θ ∈ R D ). This can be expressed as followsEQUATIONwhere θ 0 ∈ R D denotes the pre-trained weights and P ∈ R D×d is a random projection matrix. We demonstrate that learning task vector compositions leads to the same formulation. For brevity of exposition, let us consider compositions at the block level. For the j-th parameter block, we haveEQUATIONWe draw a parallel between Eqs. 6 and 8 and note that aTLAS explicitly exploits the low intrinsic dimensionality by learning a small set of coefficients. The number of task vectors, i.e., n, is much smaller than the dimension of weight vector θ(j)i , and is analogous to the intrinsic dimensionality d. However, as opposed to using a random projection matrix P , aTLAS constructs the projection matrix from task vectors, making use of the learned representations. To demonstrate its advantage, we use the same number of bases for task vectors 2 and random bases 3 , and show that task vectors consistently achieve higher performance in Figure 2 . These results solidify our understanding of task vectors being knowledge carriers. We thus set out to apply aTLAS to various applications. 
task arithmetic: 
 Task arithmetic [28] is comprised of a few tasks aimed at editing pre-trained models using task vectors. Following previous practice [28, 44] , we conduct experiments under the settings of task negation and task addition on eight image classification datasets (details included in Appendix A). 2 A fixed number of task vectors are selected based on the blockwise gradient. Details can be found in Section 5.2 and Appendix D.6.3 Each random basis of the projection is drawn from a Gaussian distribution with the mean and standard deviation to match those of the pre-trained weights in the corresponding parameter block, i.e., θTable 1 : Performance of task negation averaged across eight datasets. Selected results must maintain at least 95% of the pre-trained accuracy on the control dataset, following previous practice [44] . Best performance in each section is highlighted in bold. Task vector is abbreviated as t.v. Results for each dataset are available in γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks (a) Coef. of param. blocks (negation) Figure 3 : Box-and-whisker plots for the learned coefficients. As each transformer layer consists of a fixed set of parameter blocks, we visualise the distribution of coefficients for these parameter blocks across all layers, for (a) task negation and (b) task addition, as well as (c) distribution of coefficients by layer. We denote the learnable LayerNorm parameters by γ and β. Weights and biases are denoted by W and b, respective, with attention layer parameters indexed by superscripts and the MLP parameters indexed by subscripts.EQUATIONPrevious works acquire the optimal isotropic scaling factor on task vectors via a hyper-parameter search on validation sets. As such, we learn anisotropic scaling matrices on the same validation sets, and visualise the learned coefficients to shed light on this mechanism.Task negation aims to reduce undesired biases, characterised by the performance, on a target task, while maintaining performance on a control dataset, ImageNet [52] in this case. Denote the validation sets for the target and control tasks by D t and D c , respectively. We perform a simultaneous gradient ascent on the target task and gradient descent on the control task, described as follows, arg minEQUATIONwhere τ t is the task vector for the target dataset, and cross-entropy loss is used. The learning objectives with linearised task vectors can be derived easily based on Eq. 5, and so are omitted.We summarise the task negation results in Table 1 , and show that our method significantly improves upon standard task vectors, while the improvement upon linear task vectors is less prominent. In particular, we observe that weights matrices tend to have much larger negative coefficients, as shown in Figure 3a . To investigate this, we instead only learn coefficients for the weight matrices, with zero coefficients on other parameter blocks, effectively reducing the number of learnable parameters by two thirds. With ViT-B/32 as the backbone, we observe an average accuracy of 20.14 (vs. 18.76) on target tasks and 61.23 (vs. 61.21) on the control task, which shows that weight matrices carry majority of the knowledge required for task negation.Task addition aims at producing a multi-task model using task vectors acquired from a range of datasets. We utilise task vectors from the eight image classification datasets, and learn the anisotropic 2 , where our method yields substantial improvements. Interestingly, we note that with previous methods [28, 44] , linear task vectors outperform the standard ones in terms of absolute accuracy, while the converse is true with our method. To investigate this, we compute the pairwise disentanglement error ξ [44] , which measures the percentage of data with inconsistent predictions when two task vectors are combined (more details in Appendix C.2). Results in Figure 4 show that standard task vectors with learned anisotropic scaling achieve the lowest average error, indicating less interference in task vector composition. Along with higher fine-tuning accuracy, previously referred to as the non-linear advantage [44] , standard task vectors demonstrate stronger performance in task addition. Furthermore, we again observe that weight matrices have consistently larger coefficients in Figure 3b , and learning coefficients on weight matrices alone results in an accuracy of 84.17 (vs. 84.98) using ViT-B/32. This suggests that weight matrices in transformers are the primary knowledge carrier, which enabled knowledge composition and negation. Note that for better clarity in visualisation, we add L 1 regularisation on the learned coefficients during learning, which causes marginal performance drop (84.23 vs. 84.98) but significantly improves interpretability. In addition, we observe substantially higher coefficients on deeper layers (Figure 3c ). This aligns with our understanding that early layers extract generic features that do not vary significantly across datasets [29] , while the deeper layers produce task-specific features and require more careful adaptations. 
knowledge transfer in low-data regimes: 
 Beyond model editing for task arithmetic, we explore the idea of transferring existing knowledge in task vectors to previously unseen tasks. To this end, we use the CLIP [47] model and a total of 22 image classification datasets, each of which produces a task vector. We defer the details of datasets and the process to acquire task vectors to Appendix A. Denote the set of available task vectors by T = {τ i } n i=1 , and the dataset corresponding to task vector τ i by D i . For each target dataset D t , we learn task vector compositions using the subset T \ {τ t }, excluding the task vector for the target dataset to avoid information leakage. We test our method in few-shot and test-time adaptation, to demonstrate its effectiveness in low-data regimes. Notably, we observe that task vectors complement existing few-shot methods. Combining aTLAS with them thus leads to significant improvements.Few-shot recognition requires learning new objects or concepts using a limited amount labelled data-k per class for k-shot. Following previous practice [69] , we approach this problem by adapting a pre-trained CLIP model [47] to each target dataset D t . We use the subset of task vectors T \ {τ t } and k ∈ {1, 2, 4, 8, 16} images from dataset D t . During training, we adopt the cross-entropy loss and minimise objectives described in Eqs. 4 and 5 for standard and linear task vectors, respectively.We compare against Tip-Adapter [69] and LP++ [25] using CLIP with ViT-B/32 backbone, across 22 datasets over three random seeds, and summarise the results in Figure 5a . We show that with k = 1, our approach, aTLAS, significantly outperforms previous methods, demonstrating the effectiveness of knowledge transfer with scarce labelled data. More importantly, we note that the idea of task vector composition is highly complementary to those presented in previous methods. As such, combining aTLAS with them results in significant improvements. This is also illustrated in Figure 5b as a Venn diagram, where we show the percentage of examples in the validation set that are incorrectly classified by the pre-trained model but correctly classified with few-shot methods. Out of the examples aTLAS improves upon, around half are unique compared against either Tip-Adapter or LP++, demonstrating its complementarity. We also found that standard task vectors generally perform better than their linearised counterparts, and so defer the results of linear task vectors to Appendix D.2.In addition, due to the low number of learnable parameters, aTLAS exhibits strong generalisability. To demonstrate this, we learn task vector composition on ImageNet [52] , and test it on out-of-domain (OOD) datasets: ImageNet-A [22] , ImageNet-R [21] , ImageNet-sketch [60] and ImageNetV2 [50] . We summarise the results in Figure 5c , which shows the performance difference against the pre-trained model. Notably, aTLAS is the only method that consistently improves upon the pre-trained model on OOD datasets, and combining aTLAS with other methods can improve their generalisability.We also test our method and variants integrated with Tip-Adapter and LP++ using other backbones, including ViT-{B/16, L/14} and ResNet-{50, 101}, and find that the results are consistent with those for ViT-B/32. More details can be found in Appendix D.3.In practical applications, there may only be a limited number of task vectors available, or the number of task vectors used in training may be restricted due to memory constraints. To this end, we study the influence of task vector budget b on few-shot recognition performance. We experiment with four selection strategies: For a task vector budget b ∈ {1, 2, 5, 10, 15, 21}, we summarise the few-shot recognition performance in Figure 6 . First, we note that the accuracy of aTLAS does not plateau with the maximum number of task vectors available (21) , indicating that more task vectors could be beneficial. Second, we find that selecting task vectors based on feature similarity is a simple yet effective approach with sufficient budgets (b > 5). Selecting whole task vectors with gradient is less effective, generally on par with random selection. Nevertheless, the blockwise variant achieves the best accuracy, particularly for very low budgets (b ∈ {1, 2}), as it is able to exploit knowledge from more task vectors than the budget dictates. We thus deduce that parameter blocks can function as knowledge carriers in isolation, independent of the task vectors to which they belong. In fact, a parameter block τ (1) as part of the task vector τ = τ (1) , . . . , τ (m) can be considered as a task vector by itself, i.e., τ (1) , 0, . . . , 0 . This modular nature underscores the potential of task vectors for flexible and efficient knowledge transfer.Test-time adaptation (TTA) [35, 57, 59] assumes no labelled data is available for the target task, requiring the model to adapt in an unsupervised fashion. We conduct experiments under the offline adaptation setting, which allows access to the target dataset. We consider three categories of selfsupervised techniques for TTA: constrastive objectives, entropy objectives and pseudo labelling.Contrastive objectives align representations of the same image under different data augmentations.For this category, we adopt SimCLR [9] , a simple yet effective method. Entropy objectives encourage the pre-trained model to produce confident predictions on unseen datasets by minimising the entropy over the predictions. This technique was previously explored by Yang et al. [65] in model merging.While effective in simpler cases, it can lead to catastrophic collapse in TTA. Therefore, we utilise a state-of-the-art sharpness-aware entropy minimisation algorithm named SAR [43] . Last, we experiment with an unsupervised pseudo-labelling algorithm inspired by FixMatch [54] , which we refer as unsupervised FixMatch (UFM). UFM selects an equal number of highly confident examples per class as the labelled set, and then employs FixMatch to produce pseudo-labels from rest of the unlabelled examples. Details are available in Appendix E.We summarise the results in Table 3 and compare our method, i.e., learning task vector compositions, against the conventional approach of tuning the layer normalisation parameters [43, 57, 59 ]. We show that under all self-supervised objectives, aTLAS achieves higher accuracy than tuning the LayerNorm.In particular, LayerNorm has 30k learnable parameters with ViT-B/32 while our method only has 3.5k learnable parameters. We note that with the UFM objective, aTLAS performs the best and improves the accuracy by an average of 6.5 absolute points over the zero-shot baseline.  
relation to parameter-efficient fine-tuning: 
 One of the key advantages of aTLAS is its ability to adapt pre-trained models with few learnable parameters, making it suitable for parameter-efficient fine-tuning (PEFT). Similar to popular PEFT methods such as low-rank adaptation (LoRA) [23] , our approach does not introduce additional modules, thereby avoiding an increase in inference complexity. In addition, since only the encoded weight matrices in LoRAs have non-zero weight difference, LoRAs are in fact sparse task vectors. They can thus be seamlessly integrated into our method, significantly reducing the memory cost.Due to the sparsity and rank deficiency, LoRAs as task vectors may have limited representation capacity and carry less knowledge. Therefore, they may be inferior to standard task vectors for knowledge transfer. We investigate this by learning linear combinations of LoRAs4 using our method, under the settings of few-shot recognition. Results are summarised in Table 4 . We first shed light on the impact of sparsity, and compare two variants of our method that either learns linear combinations of all parameter blocks or just the weight matrices. Results show that sparsity results in an accuracy decrease of around 0.5% on average, except for the one-shot setting. The rank deficiency, on the other hand, causes more substantial accuracy drop. Nevertheless, this can be largely mitigated by increasing the rank. Using a rank of 64 leads to similar performance compared to learning compositions of only weight matrices in standard task vectors. In conclusion, while the sparsity and rank deficiency introduce some performance drops, especially in low-shot settings, LoRAs are competitive alternatives to standard task vectors due to their low memory cost. Despite the parameter efficiency of aTLAS, its performance is not as competitive when sufficient training data is available. To address this, we devise a strategy to flexibly scale up the number of learnable parameters as needed. Specifically, we randomly divide each parameter block into K partitions, and assign a learnable coefficient to each partition, naturally increasing the number of learnable parameters by K-fold. We denote these variants by aTLAS ×K. We conduct experiments with these variants using {1, 5, 10, 25, 35, 50, 100}% of the total available training data across the 22 datasets used in Section 5. The results are summarised in Figure 7 , showing that our method consistently improves as K increases. Compared to LoRAs, particularly with limited training data, our method achieves higher performance with fewer learnable parameters. With sufficient training data, the variant aTLAS ×1200 leads to higher performance with a similar number of learnable parameters, as it is able to exploit the knowledge contained in the task vectors that may otherwise be unobtainable from the target dataset. 
related work: 
 Task vectors and model compositions. Recent studies have demonstrated the possibility of manipulating the behaviours of neural networks directly in the weight space [27, 62, 64] . In particular, task vectors [28] , as a carrier of the domain-specific knowledge learned through fine-tuning, exhibit strong compositional properties. Such compositionality can be enhanced via linearisation using first-order Taylor expansion [44] , and improves model editing with simple arithmetic, e.g., addition, negation, etc. Yang et al. [65] also investigated the idea of learning layer-wise coefficients to improve task arithmetic. In addition, low-rank adaptations [23] , as special forms of task vectors, were shown to also support such arithmetic operations. A recent study [3] also investigated the idea of learning combinations of LoRAs for few-shot recognition.Model-based transfer learning. One interpretation of transfer learning [73] is to exploit the knowledge encapsulated in a pre-trained model for a target domain. Amongst various sub-modules of a pre-trained model, transferring the feature extractor is the most extensively studied. This ranges from early convolutional neural networks [18, 53] to modern transformers [58] , from vision backbones [14, 37] to language models [13, 46] . For vision applications, classification models trained on ImageNet [52] have been used as the medium for knowledge transfer. In recent years, contrastively pre-trained multi-modal models such as CLIP [47] have emerged as a prevelant choice. Such models are trained on large volumes of data by aligning image and language representations, leading to strong baselines well suited for transfer learning. CLIP representations have since been use for medical imaging [70] , semantic segmentation [72] , satellite imaging [40] , etc.Model adaptation in low-data regimes. The performance of pre-trained models is often constrained when applied to specific tasks with limited labelled data. To address this limitation, extensive research has been conducted on few-shot adaptation of CLIP [47] . These studies focus on various techniques, including prompt engineering [71] , feature adaptation [16] , and more recently classifier adaptation [25, 69] . In addition to few-shot adaptation, test-time adaptation represents an even more challenging scenario where no annotated data is available. This typically requires leveraging self-supervised objectives to adapt the model, employing methods such as entropy minimisation [35, 43, 59 ], contrastive learning [8] , pseudo labelling [35] and image rotation prediction [57] . 
conclusion: 
 In this paper, we introduced aTLAS, a learning algorithm that leverages the rich knowledge encapsulated in task vectors through learned linear combinations with anisotropic scaling. Unlike conventional methods that learn network parameters, our approach focuses on learning coefficients on task vectors, significantly reducing the number of learnable parameters. We conducted experiments across task arithmetic, few-shot recognition, test-time adaptation and parameter-efficient fine-tuning, demonstrating the effectiveness of our method with supervised and unsupervised objectives. In particular, we highlighted several properties of aTLAS, including low disentanglement error, robustness against domain shift, effectiveness in low-data regimes, complementarity with existing few-shot methods, etc. These properties paved the way for efficient knowledge composition and transfer.Limitations. As a task vector is defined with respect to a specific pre-trained model, knowledge composition and transfer are not yet feasible across different architectures. This may become possible with suitable projections and remains part of the future work. In addition, combining large numbers of task vectors can consume a substantial amount of GPU memory when training larger models. This can be mitigated by selecting a subset of task vectors, using LoRAs as task vectors or by offloading the computation of task vector composition to CPU, at the cost of training speed decrease. It is also possible to perform task vector composition at bit-width lower than floating point precision, e.g., 4-bit. Similar features are being tested with popular deep learning frameworks such as PyTorch, and we expect the memory requirement of larger models to be less of a constraint in the future.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: Although we do not report compute requirements in the paper, on a single A100 or 4090GPU except for ViT-L/14 experiments that were performed on 2 A100. The typical run time for an experiment is 2 hours with ViT-L/14 experiments going up to 8h. We estimate the total compute needed to 1000 GPU-hours for repeated results over 3 seeds. The research project includes failed experiements and iterations on the method that are not reported.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognise that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: All datasets we use are credited in Section A. 
 
 Appendix: We acquire task vectors by fine-tuning CLIP [47] on a variety of 22 image recognition datasets:(1) Stanford Cars [30] , (2) DTD [11] , (3) EuroSAT [20] , (4) GTSRB [56] , (5) MNIST [32] , (6) RESISC45 [10] , (7) SUN397 [63] , (8) SVHN [41] , (9) CIFAR10 [31] , (10) CIFAR100 [31] , (11) ImageNet [52] , (12) STL10 [12] , (13) Food101 [5] , (14) Caltech101 [34] , (15) Caltech256 [17] , ( 16) FGVCAircraft [39] , (17) Flowers102 [42] , (18) Oxford Pets [45] , (19) CUB200 [61] , (20) PascalVOC [15] , (21) Country211 [47] , and ( 22) UCF101 [55] . Fine-tuning was conducted using AdamW optimiser [38] , with a learning rate of 10 -5 , batch size of 128 and weight decay of 0.1. Details of the datasets, additional dataset-specific hyper-parameters, and the accuracy after fine-tuning for an assortment of backbones are shown in Table 5 . We use the same hyper-parameters for the linearised variants of the model. To shed light on the semantic relationships amongst datasets, we extract the features of all images for each dataset, and visualise the distributions as ellipses (Figure 8 ). Specifically, for each dataset, the mean µ t ∈ R d and covariance Σ t ∈ R d×d of image features are computed. Principal component analysis (PCA) is used produce a projection matrix P ∈ R d×2 from the mean features µ t . Subsequently, the mean and covariance with reduced dimensionality can be expressed as P T µ t and P T Σ t P , respectively. The mean image features for all datasets are visualised as the ellipse center, with the dimensionality reduced to 2 using Principal Component Analysis (PCA). The dimensionality of covariance matrices are also reduced using the same principal components. We show visualisations with (a) ×1 and (b) ×3 standard deviations. Pre-trained CLIP [47] with ViT-B/32 is used to extract image features. The evaluation of task negation is conducted on eight classification datasets (1-8 in Table 5 ), following previous practice [28, 44] . In particular, we learn anisotropic scaling using the validation set of each dataset. We also adjust the learning rates and training epochs on the same validation set. The details are shown in Table 6 . We report detailed task negation results for each dataset in Table 7 . In addition, for more evidence that weight matrices learn large negative coefficients, we show a detailed visualisation of the learned coefficients in Figure 9 and distribution of the coefficients in Figure 10 . γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(c) EuroSAT (standard) -0.4 -0.2 0.0 0.2 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(d) EuroSAT (linear) -1.0 -0.5 0.0 0.5 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks (e) GTSRB (standard)-0.75 -0.50 -0.25 0.00 0.25 0.50 learned coef. (lambda)γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(f) GTSRB (linear) -0.8 -0.6 -0.4 -0.2 0.0 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks (g) MNIST (standard)-1.0 -0.5 0.0 0.5 learned coef. (lambda)γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(i) RESISC45 (standard) -0.6 -0.4 -0.2 0.0 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(j) RESISC45 (linear) -0.6 -0.4 -0.2 0.0 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(k) SUN397 (standard) -0.4 -0.2 0.0 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(l) SUN397 (linear) -0.8 -0.6 -0.4 -0.2 0.0 0.2 learned coef. (lambda) γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks (m) SVHN (standard) -0.5 0.0 0.5 learned coef. (lambda)γ1 β1 W I b I W O b O γ2 β2 W1 b1 W2 b2parameter blocks(n) SVHN (linear)Figure 10 : Additional box-and-whisker plots for the learned coefficients in task negation, beside previous visualisation on Cars (Figure 3a ), including results on (a, b) DTD, (c, d) EuroSAT, (e, f) GTSRB, (g, h) MNIST, (i, j) RESISC45, (k, l) SUN397 and (m, n) SVHN.Table 8 : Detailed performance for task addition across all eight datasets. We show additional performance with learned isotropic scaling, which have comparable accuracy to the simple hyper-parameter search used in previous methods [28, 44] . Highest performance in each section is highlighted in bold. The method search corresponds to model f (x; θ0 + α τ i), where α is determined via a hyper-parameter search. Methods iso. and aniso. use learned coefficients and correspond to models f (x; θ0 + αiτ i) and f (x; θ0 + Λiτ i), respectively. Task addition is also evaluated on datasets 1-8 shown in Table 5 . The hyper-parameters are identical to fine-tuning, except the learning rate is modified to 10 -3 . We show detailed performance on each dataset in Table 8 , where we compare our method against hyper-parameter search used in previous works [28, 44] , and another variant with learned isotropic scaling. We also visualise the learned coefficients with L 1 regularisation in Figure 12 . It can be easily observed that weight matrices, particularly those in the deeper layers, have significantly higher learned coefficients, which conforms to our observations in Figures 3b and 3c . Since our method involves learning the coefficients, unlike previous methods [28, 44] that only require a hyper-parameter search, we also compare against the direct fine-tuning approach. We fine-tune the pre-trained model on the union of eight datasets, assuming only the validation sets are available. The results are shown in Figure 11. Unsurprisingly, task vector compositions, whether the coefficients are searched or learned, are less susceptible to the lack of data, as the accuracy only starts to drop with less than 35% of the data. The performance of fullparameter fine-tuning, however, drops substantially as the amount of data available decreases.In addition, we provide more technical details and intuitions on the pairwise disentanglement error [44] , which was visualised in Figure 4 . Specifically, we make a few changes to the formulation proposed by Ortiz-Jiménez et al. [44] , and evaluate the disentanglement error only with the optimal Two variants of Tip-Adapter [69] were proposed for few-shot recognition where the weights of the adaptor are either fixed based on features of the few-shot examples or further fine-tuned. We only study the fine-tuned variant due to its higher performance. Tip-Adapter has two hyper-parameters, which in the original paper are optimised through hyper-parameter search on a separate validation set. This practice may not align with the principles of few-shot learning, where access to extensive validation data is typically limited. In addition, Huang et al. [25] note that the performance of Tip-Adapter is very sensitive to these hyper-parameters. We thus opt to learn these two hyper-parameters together with the feature adaptor through gradient descent. The learning rates for the feature adaptor and the hyper-parameters are set to 10 -3 and 10 -1 , respectively.For both Tip-Adapter and LP++ [25] , we conduct experiments using the publicly available codebase 5 . We train both LP++ and Tip-Adapter for 300 epochs on frozen zero-shot features. We apply a cosine annealing decay for Tip-Adapter and maintain fixed learning rates for LP++ as per the official implementation.We report the average few-shot accuracy over the 22 datasets in Table 9 , which corresponds to results in Figure 5a . In particular, we show results with linearised task vectors, as proposed by Ortiz-Jiménez et al. [44] . As highlighted in Section 4, learned anisotropic scaling allows standard task vectors to achieve stronger performance than the linear variants in task addition. For few-shot recognition, we again observe that standard task vectors result in superior performance in most cases. We, however, note the exception that linear task vectors when combined with LP++ achieve higher performance in the 1-shot setting. Nevertheless, the margin over standard task vectors is not very significant, and aTLAS using standard task vectors when integrated with Tip-Adapter is generally a stronger few-shot model.We use the AdamW [38] optimiser with a learning rate of 10 -1 and a weight decay of 10 -1 . Our method by itself is trained for 10 epochs with ViT backbones and 30 epochs with ResNet backbones.We show that state-of-the-art few-shot methods can be seamlessly integrated into our method, since both Tip-Adapter and LP++ focus on the classifier, while aTLAS improves the feature representations.We experiment with two strategies to combine aTLAS with previous methods, where we either (1) train our method first and use the frozen representations to train a previous method, or (2) train parameters in both methods jointly. Results in Table 10 shows that the joint training strategy results in higher performance, particularly in low-shot settings. We therefore adopt the joint training strategy when combing our method with Tip-Adapter. During training, we adopt different learning rates for different parameter groups, that is, 10 -1 for learnable coefficients in aTLAS and the hyper-parameters in Tip-Adapter, and 10 -3 for the adaptor. The joint training takes 20 epochs for ViT backbones and 60 epochs on ResNet backbones, twice the number of epochs when training aTLAS alone. On the other hand, The joint training strategy with LP++ is non-trivial, due to LP++'s superconvergence strategy being designed around frozen feature representations, which would have been updated every iteration by aTLAS. We thus use the sequential strategy to combine aTLAS and LP++. We include detailed results for each dataset with ViT-B/32 in Table 11 and additional results with different backbones in Table 12 , where we show our method scales well across different datasets and backbones.Table 12 : Detailed few-shot accuracy for each dataset across RN50, RN101, ViT-B/16, and ViT-L/14 backbones. We report results for the same random seed. In the case where the results are worse than the zero-shot accuracy, we report zero-shot accuracy. Highest performance and those within a range of 0.1 in each section are highlighted in bold. Tip-Adapter is abbreviated as Tip. We show detailed results for out-of-domain generalisation over k ∈ {4, 16} shots in Table 13 . These results correspond to those presented in Figure 5c . aTLAS is the only method that consistently improves test accuracy over the zero-shot model on out-of-domain images. When combined with LP++ or Tip-Adapter, aTLAS can be observed to improve the out-of-domain generalisation of these methods.In this section, we examine the informativeness of a task vector across different target datasets. To this end, we apply aTLAS to each of the 22 datasets using only one task vector. For each dataset, we compute the relative accuracy improvement, that is, the accuracy improvement of aTLAS normalised by that of fine-tuning in the full parameter space. Note that aTLAS is applied under the 16-shot setting, while standard fine-tuning uses all training data available. Results are shown in Figure 13 . We first note that certain datasets are more prone to accuracy improvement, such as EuroSAT, MNIST, etc., as indicated by the high percentage across entire rows. This is most likely due to the low intrinsic dimensionality of the task. In addition, we highlight the average improvement in the last row. Notably, certain task vectors, e.g., ImageNet task vector, are particularly informative while others, such as those from Flowers102 and OxfordPets are much less so. These results illustrate the varying contributions different task vectors can have depending on the target dataset, which also motivated subsequent efforts on careful task vector selection. An alternative to save memory during training is to utilise gradient-free methods to learn the coefficients. We follow previous work on the combination of LoRAs [24] and use the nevergrad [49] library. We observe a memory usage reduction of 60% from 10GB to 4GB calculated using a dedicated pytorch function 6 . Results for few-shot recognition are summarised in Table 16 . We show that although gradient-free optimisation improves upon the zero-shot model, the performance quickly plateaus as the amount of data increases. In addition, learning anisotropic scaling results in worse performance, most likely due to the relatively high number of parameters. We provide more details on the Unsupervised FixMatch (UFM) approach in this section. Fix-Match [54] utilises a labelled set to guide training, which is given as part of the semi-supervised learning protocol, while we produce a class-balanced "labelled" set from unlabelled images. Given a target dataset D t consisting of N unlabelled images, we first rank the examples by the prediction scores from the zero-shot model across C classes. We then select the top min(N/C, 100) examples, that is, at most 100 examples per class, as a trusted set in absence of a labelled set. The standard cross-entropy loss is applied to the trusted set. For the rest of the unlabelled images, we use a weakly augmented (Open-CLIP [26] validation augmentations) view of an image to produce pseudo-labels, and incur a loss on the strongly augmented view (Tip-Adapter [69] augmentations). Denote an image with weak augmentation by x, its strongly augmented view by x ′ , and the predictions made by network by ŷ and ŷ′ , respectively, the unsupervised loss can be expressed asEQUATIONσ(ŷ) = ŷ0.5EQUATIONwhere 1(•) denotes the indicator function, σ(•) performs re-normalisation with adjusted temperature scaling, and ω is a confidence threshold that is linearly adjusted from 0.9 to 1 during training. The trusted set is re-estimated at the beginning of each epoch to account for the improving accuracy of the model. In training, images in the trusted set are over-sampled to constitute one fourth of each batch, as this practice prevents the model from diverging due to confirmation bias [2, 54] .Dividing a parameter block into K random partitions allows us to introduce more learnable coefficients to each block, thus scaling up our method flexibly. One draw back of this approach, however, is that masks for the partitions have to be stored in memory, resulting in a linear memory increase with respect to the size of the parameter block and the value K. To reduce the memory consumption the of aTLAS ×K variants, we only apply it to LoRAs task vectors. Nevertheless, these memory requirements could most likely be reduced by exploiting sparse matrices or memory efficient matrix indexing techniques, which we plan to investigate in the future.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: We find no direct societal impact for this research paper as research is conducted on controlled open-sourced datasets. This paper did not study applicability to specialized datasets that can be used to impact society. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimising neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Details about the process to acquire LoRAs are included in Appendix D.7.github.com/fereshteshakeri/fewshot-clip-strong-baselinehttps://pytorch.org/docs/stable/generated/torch.cuda.memory_allocated.html 


Decision: The paper presents an improvement to task arithmetic using different trained coefficients for each block. Overall the reviewers thought that the idea was original, the paper well presented, and the experiments thorough. In the rebuttal, a more thorough discussion of prior work and a multitask finetuning baseline were provided. Otherwise I believe no major weaknesses remain, and I recommend the paper for acceptance. Congratulations to the authors!

Reviews: review: paper_summary: paper_summary: The paper enhances the performance of task arithmetic, a recent model editing technique based on weight interpolation, in vision-language models. Instead of the original task- and parameter-independent scaling coefficients of the task vectors, it proposes to learn anisotropic scaling coefficients from validation data, resulting in significant performance improvements, particularly in task addition. The method also proves effective in few-shot learning and test-time adaptation scenarios and as a parameter-efficient fine-tuning (PEFT) method in low-data regimes.
strength_weakness: strength_weakness: **Originality and references to previous work**

1. The idea of learning task-wise and layer-wise scaling coefficients for the task vectors is one of the main features of Adaptive Model Merging (*AdaMerging*, Yang et al., 2024). This work, which is not currently referenced, must be duly cited, along with a detailed discussion of the similarities and differences between aTLAS and AdaMerging. 
2. Similarly, the idea of using test-time adaptation techniques such as entropy optimization was also present in Yang et al. (2024).

**Clarity and missing details** 

3. The paper does not report the essential methodological details of aTLAS. Specifically, the loss, optimizer, learning rate, and hyperparameters for learning the scaling coefficients are missing.
4. The disentanglement error (line 143) must be defined and briefly explained. Moreover, as it seems to slightly differ from the original metric defined in Ortiz et al. (2024) – as likely it is computed only for the best set of scaling coefficients – this difference should also be mentioned.
5. In lines 153-155, some citations could be added, e.g., to the fact that the representation built by neural networks is often hierarchical and increases in complexity with the layer depth.
6. The experiments with ResNet backbones are never mentioned in the main text. 
7. A comparison of the computational costs compared to standard task arithmetic is missing and should be provided.

Yang, E., Wang, Z., Shen, L., Liu, S., Guo, G., Wang, X. and Tao, D., 2024. AdaMerging: Adaptive model merging for multi-task learning. The Twelfth International Conference on Learning Representations.

Ortiz-Jimenez, G., Favero, A. and Frossard, P., 2024. Task arithmetic in the tangent space: Improved editing of pre-trained models. Advances in Neural Information Processing Systems, 36.
questions: questions: 8. **Multi-task learning vs. learning the coefficients.** Differently from the original task arithmetic technique, aTLAS involves learning the optimal merging coefficients for each task and block from data. How would the performance of fine-tuning the pre-trained model on the same number of data and for the same number of steps compare with your task addition results? What about the accuracy of the model on a control dataset that was not used during fine-tuning, e.g., ImageNet?
limitations: limitations: The paper adequately discusses its limitations. I do not foresee any potential negative societal impacts arising from this study.
paper_summary: paper_summary: The paper introduces a method named aTLAS, which leverages task vectors and anisotropic scaling to enhance knowledge composition and transfer in pre-trained models. The authors investigate whether components of task vectors, particularly parameter blocks, exhibit similar characteristics and how these can be used to improve knowledge composition and transfer. The effectiveness of the proposed method is demonstrated in various tasks such as task arithmetic, few-shot recognition, and test-time adaptation.
strength_weakness: strength_weakness: - Knowledge composition and transfer are limited to the specific pre-trained model architecture, which may restrict its applicability across diverse model architectures.
questions: questions: - How does the method handle potential conflicts when combining task vectors from very dissimilar domains? Will it negatively affect performance?
limitations: limitations: - Knowledge composition and transfer are limited to the specific pre-trained model architecture, which may restrict its applicability across diverse model architectures.
paper_summary: paper_summary: This paper presents a method called aTLAS, which leverages task vectors to enhance transfer learning in neural networks. Task vectors represent the difference in weights between a pre-trained model and its fine-tuned variant. aTLAS introduces anisotropic scaling to these task vectors by learning different coefficients for each parameter block, which allows for more effective composition of knowledge from different tasks. The method is tested in various scenarios, including task arithmetic, few-shot recognition, and test-time adaptation, demonstrating improvements in performance and efficiency.
strength_weakness: strength_weakness: - **Only CLIP?** the method is primarily tested on the CLIP model and might not directly generalize to other architectures without significant modifications. Future work should explore the applicability of aTLAS across a broader range of model architectures.

- **Computational complexity (is this scalable)?**  While the method is parameter-efficient, the computation of task vector compositions during training can still be resource-intensive, especially for large models. Strategies to optimize this process or reduce its computational footprint would be beneficial.  The method's scalability with an increasing number of task vectors is not fully explored. While the paper shows that performance improves with more task vectors, it is unclear how this scales with very large sets of task vectors or in more complex multi-task environments. Also, earning different coefficients for each parameter block introduces additional complexity. The benefits of this complexity should be weighed against the potential for simpler approaches that might achieve similar results with less computational overhead.

- **How exactly should one select the task vectors?** The selection of task vectors for composition can significantly impact performance. The paper discusses various selection strategies but does not provide a definitive approach. Further research into more sophisticated selection mechanisms could enhance the method's robustness and effectiveness.


- **How to leverage this in real world?** The experiments are conducted in controlled settings with well-defined datasets. Evaluating the method's performance in more diverse and real-world scenarios would provide a better understanding of its practical applicability and limitations.
questions: questions: See weaknesses section
limitations: limitations: See weaknesses section

Comments: comment: Thank you for your detailed response, which successfully addressed all of my concerns. I encourage you to include the additional discussions in the revised version of your submission. Overall, I believe this paper merits acceptance, and I have adjusted my initial score accordingly. Best wishes.