Knowledge Composition using Task Vectors with Learned Anisotropic Scaling (2407.02880v2)

Abstract: Pre-trained models produce strong generic representations that can be adapted via fine-tuning. The learned weight difference relative to the pre-trained model, known as a task vector, characterises the direction and stride of fine-tuning. 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 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 labeled 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 PEFT method, particularly with less data, and demonstrate its scalibility.

Analysis of Knowledge Composition using Task Vectors with Learned Anisotropic Scaling

The research paper analyzes the utility and enhancement of task vectors derived from fine-tuning pre-trained models. Task vectors, defined as the difference between weights of pre-trained and fine-tuned models, have been acknowledged for their capability to carry task-specific information. The paper introduces an innovative method termed aTLAS, which builds upon this foundation by utilizing anisotropic scaling to improve knowledge transfer and composition across tasks.

Methodological Advancements

The core innovation of the paper lies in the introduction of aTLAS, an algorithm that applies learned anisotropic scaling to task vectors. Unlike previous methods that use isotropic scaling, aTLAS allows independent scaling of parameter blocks, such as weights and biases within the task vector, thereby increasing flexibility and optimization precision. This approach exploits the low intrinsic dimensionality of pre-trained model landscapes, where only a limited number of scaling coefficients are learned.

aTLAS facilitates modular learning by allowing task vectors from diverse domains to be linearly combined, maintaining the pre-trained model's general representation while adapting to specific tasks. This characteristic is particularly beneficial when data is scarce, paving the way for effective few-shot learning and test-time adaptation, both unsupervised and supervised.

Empirical Validation

The paper presents empirical evidence supporting the strength of the proposed method in multiple domains, including task arithmetic, few-shot recognition, and test-time adaptation. The anisotropic scaling allows task vectors to be more disentangled during composition, minimizing interference and improving generalization.

For task addition, a significant improvement was noted with a higher accuracy using aTLAS compared to prior methods. Specifically, the standard task vectors with learned anisotropic scaling surpassed linearized task vector compositions from earlier research by Ortiz-Jiménez et al., showcasing reduced disentanglement errors and superior performance, with learned scaling reflecting the significance of deeper neural network layers.

Practical Implications

The implications of this research are broad-ranging in the field of AI. The ability to seamlessly combine task vectors can significantly optimize the model adaptation process for varied tasks without extensive retraining. aTLAS offers parameter-efficient fine-tuning, making it advantageous in resource-constrained settings.

Moreover, by demonstrating that parameter blocks can be independently scaled and composed, the paper lays the groundwork for new modular approaches in neural network design, potentially leading to more versatile and adaptive models that are less dependent on large datasets for fine-tuning.

Future Directions

One avenue for further exploration is the application of aTLAS across different architectures or foundation models, where suitable projections could enable cross-architecture knowledge transfer. The paper also hints at potential advancements in memory efficiency and computational speed by leveraging LoRAs and gradient-free optimization strategies, which could see aTLAS being applied to even larger models.

In conclusion, the research provides a comprehensive and effective strategy to harness and enhance task vectors for dynamic knowledge composition, addressing the challenge of task-specific model adaptation in an innovative manner. By reducing reliance on extensive data and computational resources, aTLAS could play a pivotal role in advancing the scalability and adaptability of AI systems.