Task Vector Bases: A Unified and Scalable Framework for Compressed Task Arithmetic

TL;DR

This paper introduces Task Vector Bases, a framework that compresses multiple task vectors into fewer basis vectors, enabling scalable task arithmetic with theoretical guarantees and improved empirical performance.

Contribution

It proposes a novel basis compression method for task vectors that maintains arithmetic capabilities and enhances scalability in transfer learning.

Findings

Outperforms heuristic basis methods in experiments.

Reduces storage and computation while preserving task arithmetic performance.

Sometimes surpasses full task vector collections in downstream tasks.

Abstract

Task arithmetic, representing downstream tasks through linear operations on task vectors, has emerged as a simple yet powerful paradigm for transferring knowledge across diverse settings. However, maintaining a large collection of task vectors introduces scalability challenges in both storage and computation. We propose Task Vector Bases, a framework compressing $T$ task vectors into $M < T$ basis vectors while preserving the functionality of task arithmetic. By representing each task vector as a structured linear combination of basis atoms, our approach supports standard operations such as addition, negation, as well as more advanced arithmetic ones. The framework is orthogonal to other efficiency-oriented improvements in task arithmetic and can be used in combination with them. We provide theoretical analysis showing that basis compression retains addition generalization guarantees and enables principled unlearning, with error bounds depending on reconstruction quality. Empirically, our proposed basis construction methods consistently outperform heuristic basis construction baselines and, in some cases, even surpass the performance of full task vector collections across diverse downstream applications while reducing storage and computational requirements. The code is available at https://github.com/uiuctml/TaskVectorBasis.