Knowledge Composition using Task Vectors with Learned Anisotropic Scaling

TL;DR

ATLAS, an algorithm that linearly combines parameter blocks with different learned coefficients, resulting in anisotropic scaling at the task vector level, is introduced, showing the potential of aTLAS as a PEFT method, particularly with less data, and its scalibility.

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.

Figures (13)

• Figure 1: Illustration of (\ref{['fig:teaser-a']}) learning task vector compositions ($n=2$, $\boldsymbol{\theta}_0$ denotes the weights of a pre-trained model) and (\ref{['fig:teaser-b']}) the flexibility of anisotropic scaling. Assume a task vector $\boldsymbol{\tau} = \mathopen{}\mathclose{\left(\boldsymbol{\tau}^{(1)}, \boldsymbol{\tau}^{(2)}\right)$ has two parameter blocks, learning anisotropic scaling grants more flexibility when combining task vectors.
• Figure 2: Recognition accuracy versus the number of bases when optimising in a low-dimensional subspace. 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 (\ref{['fig:intrinsic-a']}) MNIST mnist_1998, (\ref{['fig:intrinsic-b']}) CIFAR100 2009_CIFAR.
• 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 (\ref{['fig:neg-vis-a']}) task negation and (\ref{['fig:add-vis-a']}) task addition, as well as (\ref{['fig:add-vis-b']}) distribution of coefficients by layer. We denote the learnable LayerNorm parameters by $\gamma$ and $\beta$. Weights and biases are denoted by $W$ and $\mathbf{b}$, respective, with attention layer parameters indexed by superscripts and the MLP parameters indexed by subscripts.
• Figure 4: Disentanglement errors between each pair of datasets. Each row reflects the percentage of data in the corresponding dataset that have altered predictions after combining two task vectors. Our method achieves stronger task addition performance as a result of less interference amongst task vectors.
• Figure 5: Few-shot experiment results averaged across $22$ datasets and three seeds, showing (\ref{['fig:few_shots']}) comparison against state-of-the-art few-shot methods with ViT-B/32 backbone and (\ref{['fig:separateimprov']}) percentage of images in the validation sets that become correctly classified after applying few-shot methods. We also show (\ref{['fig:oodrobust']}) performance difference compared to pre-trained CLIP model on OOD datasets. More detailed results are included in Appendix \ref{['apdx:few-shot']}.