Knowledge Composition using Task Vectors with Learned Anisotropic Scaling

We appreciate the reviewer for recognizing the novelty, parameter efficiency and thorough experimentation in our work, and we are thankful for the feedbacks. We now address the questions and concerns as follows.

W1. Knowledge composition and transfer are limited to specific pre-trained model architecture.

While the experiments primarily use ViT-B/32 as the visual encoder, we showed that the proposed method works consistently across other architectures including ViT-B/16, ViT-L/14 (Table 1, 2), ResNet50 and ResNet101 (Table 13). Nevertheless, we acknowledge that aTLAS cannot combine task vectors obtained from different architectures, or transfer the knowledge in task vectors to a different architecture. This may require finding appropriate projections, and remains part of the future work.

Q1. Task vectors from dissimilar domains.

We studied the potential conflicts between task vectors using disentanglement error (Ortiz-Jimenez et al., 2023), as shown in Figure 3 of the paper. Specifically, Each row reflects the percentage of data in the corresponding dataset that have altered predictions after combining two task vectors. Therefore, low disentanglement error indicates low interference between task vectors. More importantly, this interference also depends on the dataset. For two datasets A and B, combining their corresponding task vectors into one model may hurt the performance on one dataset more than the other.

Our method, particularly with standard task vectors (Figure 3c), generally decreases the disentanglement error across all pairs, which very effectively reduces the conflict between task vectors, whether they are obtained from similar or dissimilar domains.

References:

• Ortiz-Jimenez et al. Task arithmetic in the tangent space: Improved editing of pre-trained models. NeurIPS'23.

We are encouraged that the reviewer find our method interesting and the experiments comprehensive and insightful. We are thankful for the feedbacks. Below, we address the questions and concerns.

W1. Applicability across different model architectures

For task arithmetic, we included results with ViT-{B/32, B/16, L/14} backbones (Table 1, 2) following previous practice (Ilharco et al., 2023 and Ortiz-Jimenez et al., 2023). For transfer learning applications including few-shot adaptation and test-time adaptation, we additionally showed results with ResNet-{50, 101} backbones besides ViTs (Table 13). Beyond CLIP models, Ilharco et al. (2023) showed that the properties of task vectors also apply to GPT-2 and T5 models. We believe task vector composition can also be applied to these models and plan to investigate in future work.

W2. Scalability

Research on larger pools of task vectors remains part of the future work. In such scenarios, LoRA task vectors emerge as the current optimal solution for mitigating memory usage. As shown in Fig. 6a and Table 14 in Appendix F, LoRA task vectors substantially reduce memory requirements while maintaining an accuracy score comparable to full task vectors, demonstrating their feasibility and potential for large number of task vector. This finding underscores the potential of LoRA task vectors as a memory-efficient solution for exploiting extensive collections of task vectors.

W3. Definitive approach on task vector selection

Thank you for pointing this out. We find in Fig. 5 that block-wise gradient-based selection is the best strategy, especially with low budget on task vectors. Therefore, it is our recommended approach and have been explicitly stated at L213 in our revision.

W4. Applying aTLAS to the real world

We have demonstrated strong results of our method in five applications. In particular, few-shot adaptation, parameter-efficient fine-tuning, etc., have direct use cases in real-world scenarios. We also conducted experiments across 22 datasets that cover a wide range of domains, in order to test its general practicality. We believe the observations we made on these datasets can reasonably reflect the challenges in the real world.

References:

• Ilharco et al. Editing models with task arithmetic. ICLR'23
• Ortiz-Jimenez et al. Task arithmetic in the tangent space: Improved editing of pre-trained models. NeurIPS'23.

We appreciate the reviewer's recognition of the novelty, quality and impact of our work, and we are thankful for the feedbacks. In what follows, we now address the questions and concerns.

W1. Comparison to AdaMerging (Yang et al., 2024)

We thank the reviewer for suggesting this comparison, and have added this in our revision (Section 3, Tables 2 and 3) with appropriate citations. In short, the idea of anisotropic scaling is a more general formulation, while layer-wise scaling in AdaMerging is a specific variant. Besides, AdaMerging is designed for model merging, similar to task addition, which is one of the five applications we investigated. Below, we detail the key differences between our work and AdaMerging.

• Formulation. We formulated anisotropic scaling as $\Lambda \mathbf{\tau}$, where $\mathbf{\tau}$ denotes a task vector and $\Lambda$ is a block-diagonal scaling matrix. When parameters in the same layer share one scaling coefficient, this formulation specializes to AdaMerging. However, for complex use cases such as parameter-efficient fine-tuning (Sec. 6.2 and Appendix F), we show that scaling different rows, columns or random partitions of weight matrices differently is necessary to increase the representation power. This variant can also be captured by our formulation. In addition, AdaMerging constrains the learned coefficients to be in $[0,1]$ while aTLAS does not. We observe that larger or negative coefficients can be beneficial for task addition in Fig. 10. The formulation of anisotropic scaling therefore covers a wider range of use cases.
• Task vector variants. We additionally studied combining linearized task vectors (Ortiz-Jimenez et al., 2023) in Section 4 and Appendix D.3, and using LoRAs as sparse task vectors to reduce memory consumption (Section 6.1).
• Training objectives. AdaMerging uses entropy as an unsupervised objective, while we experimented with constrastive, regularized entropy and pseudo-labelling objectives in Section 5.3. Note that our unsupervised pseudo-labeling algorithm UFM outperforms the entropy objective for test-time adaptation (Table 3).
• Applications. AdaMerging focuses on model merging, whereas our work investigates additional applications such as task negation (Sec. 4.1), few-shot recognition (Sec 5.1), test-time adaptation (Sec 5.3) and parameter-efficient fine-tuning (Sec 6.2), which, as noted by the reviewer, are original contributions to the field.
• Novel insights for model merging. As discussed in L141-146, anisotropic scaling reduces interference when merging task vectors (Fig. 3), a novel insight uncovered in our paper, highlighting its advantage over isotropic scaling and negating the need for linearized task vectors.

Last, we kindly note that AdaMerging's publication in ICLR'24 coincides with the submission period of NeurIPS'24. Nevertheless, we have added comparison and acknowledgements where appropriate throughout the paper and thank the reviewer for the suggestion.

W2. Entropy optimisation in AdaMerging

We thank the reviewer for pointing this out and have added a discussion in L224 to acknowledge this. We highlight that UFM, the pseudo-labelling algorithm we designed, outperforms (regularized) entropy optimisation (SAR, Table 3). AdaMerging has also been included as a baseline in this table.

W3. Technical details

They can be found in Appendix A for task addition or negation; Appendix D.1 and D.2 for few-shot learning. In our revision, we have added hyperlinks at L160 for clarity.

W4. Disentanglement Error

The reviewer is correct that the disentanglement error is evaluated at the optimal coefficients. We have added detailed mathematical formulas in our revision for clarity, as follows

$\xi(\tau_1, \tau_2) = **E**_{\mathbf{x} \in \mathcal{D}_1} \big[ \delta \big( f(\mathbf{x}; \mathbf{\theta}_0 + \Lambda^\star_1 \mathbf{\tau}_1), f(\mathbf{x}; \mathbf{\theta}_0 + \Lambda^\star_1 \mathbf{\tau}_1 + \Lambda^\star_2 \mathbf{\tau}_2) \big) \big]$,

where $\Lambda^\star$ denotes the optimal coefficients and $\delta(x_1, x_2)$ is a distance function that returns 1 if $x_1 \neq x_2$ and 0 otherwise.

W5. Citations

We have added references to Yosinski et al. (2014) and Kornblith et al. (2019) in our revision.

W6. ResNet Experiments

We have added a hyperlink to the results with ResNets (Table 13) in our revision.

W7. Computational cost

Standard task addition performs hyper-parameter search on one scaling coefficient ranging from 0 to 1, with an interval of 0.05, and therefore runs inference 21 times. Our method is trained for 10 epochs. Using an RTX 4090, our method takes 12min to train while the hyper-parameter search takes 20min.

Q1a. Comparison against multi-task fine-tuning

As shown in figure in the attached PDF, our method consistently outperforms multi-task fine-tuning using different percentage of the validation data, since the training data is considered unavailable. However, with more data, fine-tuning is expected to achieve better performance.

Q1b. Performance on a control dataset

With our method, the merged model achieves 58.1% accuracy on ImageNet, which retains 91.6% of the zero-shot accuracy 63.4%. In comparison, the model from multi-task fine-tuning achieves 57.3% accuracy on ImageNet, which shows our method generalizes better.

References:

• Yang et al. AdaMerging: Adaptive model merging for multi-task learning. ICLR'24
• Ortiz-Jimenez et al. Task arithmetic in the tangent space: Improved editing of pre-trained models. NeurIPS'23.
• Yosinski et al. How transferable are features in deep neural networks? NeurIPS'14.
• Kornblith et al. Similarity of neural network representations revisited. ICML'19.