Knowledge Composition using Task Vectors with Learned Anisotropic Scaling

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.

Strengths

- An interesting way of using task vectors: the paper extends the concept of task vectors by introducing anisotropic scaling, which enhances the flexibility and effectiveness of knowledge transfer. This approach allows for fine-grained control over the composition of different task vectors, leading to better performance in multi-task settings. - Evaluation and analysis: the method is evaluated across multiple tasks, including task arithmetic, few-shot learning, and test-time adaptation. The comprehensive set of experiments provides strong evidence for the effectiveness of the proposed approach. It provides insights on the behavior and importance of various parts of the neural network during task vector composition. - Parameter efficiency: aTLAS is shown to be a parameter-efficient method for fine-tuning, which is particularly valuable in scenarios with limited data. The ability to achieve high performance with fewer learnable parameters is a significant advantage for practical applications.

Weaknesses

- **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

See weaknesses section

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses section

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.

Strengths

- The introduction of anisotropic scaling at the task vector level is novel and offers higher controllability in model behavior, particularly for task addition and negation. - The method is thoroughly validated across multiple tasks and datasets, showing significant improvements in performance. - aTLAS demonstrates strong parameter efficiency, making it suitable for scenarios with limited data. - The method complements existing few-shot adaptation techniques, leading to additional improvements in performance when combined.

Weaknesses

- Knowledge composition and transfer are limited to the specific pre-trained model architecture, which may restrict its applicability across diverse model architectures.

Questions

- How does the method handle potential conflicts when combining task vectors from very dissimilar domains? Will it negatively affect performance?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- Knowledge composition and transfer are limited to the specific pre-trained model architecture, which may restrict its applicability across diverse model architectures.

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.

Strengths

- **Originality:** While some parts of the method are not entirely novel (see Weaknesses section), overall, aTLAS goes beyond previous work. The few-shot adaptation application and its relation to PEFT (e.g., using LoRAs as task vectors) are also original contributions to the field of model editing/merging. - **Quality:** The work is generally of good quality. - **Significance:** Editing foundation models is an emerging field with promising real-world impact. The results obtained with the proposed method are very good. Specifically, the authors merge the parameters of 8 ViT-L-14 CLIP models while retaining 97.07% of the performance of the single models (see Tab. 2).

Weaknesses

**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

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?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper adequately discusses its limitations. I do not foresee any potential negative societal impacts arising from this study.