Decomposing Task Vectors for Refined Model Editing

TL;DR

This work tackles interference in task-vector arithmetic by decomposing each task vector into shared and unique subspaces via column-space projections, enabling more precise model editing without unintended cross-task amplification. The DeVec framework identifies a shared subspace using SVD-based projections and then extracts task-specific unique components, recombining them with learned coefficients for editing across domains. Across image classification, image generation, and toxicity removal, DeVec yields notable gains: about a 4.6% average improvement in multi-task merging, clean style mixing in diffusion models by using only unique components, and up to 47% toxicity reduction while preserving general knowledge performance. The results highlight that shared subspaces promote transfer while isolating unique components prevents interference, offering a scalable framework for refined, safer, and more controllable model editing in large pretrained networks.

Abstract

Large pre-trained models have transformed machine learning, yet adapting these models effectively to exhibit precise, concept-specific behaviors remains a significant challenge. Task vectors, defined as the difference between fine-tuned and pre-trained model parameters, provide a mechanism for steering neural networks toward desired behaviors. This has given rise to large repositories dedicated to task vectors tailored for specific behaviors. The arithmetic operation of these task vectors allows for the seamless combination of desired behaviors without the need for large datasets. However, these vectors often contain overlapping concepts that can interfere with each other during arithmetic operations, leading to unpredictable outcomes. We propose a principled decomposition method that separates each task vector into two components: one capturing shared knowledge across multiple task vectors, and another isolating information unique to each specific task. By identifying invariant subspaces across projections, our approach enables more precise control over concept manipulation without unintended amplification or diminution of other behaviors. We demonstrate the effectiveness of our decomposition method across three domains: improving multi-task merging in image classification by 5% using shared components as additional task vectors, enabling clean style mixing in diffusion models without generation degradation by mixing only the unique components, and achieving 47% toxicity reduction in language models while preserving performance on general knowledge tasks by negating the toxic information isolated to the unique component. Our approach provides a new framework for understanding and controlling task vector arithmetic, addressing fundamental limitations in model editing operations.

Figures (17)

• Figure 1: Unlike conventional task arithmetic approaches, we start by (1) Finding the subspaces shared between available task vectors. (2)Project each task vector onto the shared subspaces to derive the shared component. (3) Subtract the shared component from the original task vector to derive the unique component (dashed vectors). (4) Recompose the decomposed components with component-specific coefficients $\lambda_i$, allowing for more controlled model editing.
• Figure 2: Arithmetic on ImageNet shows the shared component between task vectors from other large datasets consists of generalizable information allowing for accuracy improvement. However, unique components have negligible effects since the information embedded in unique vectors is related to the vector's own dataset.
• Figure 3: Arithmetic on GTSRB shows how the shared component between MNIST and SVHN improves accuracy and unique components fail to have meaningful effects.
• Figure 4: Text-to-image model outputs with prompt "pixelart of athletic adult woman dressed in a sleek black and dark grey jumpsuit, with a utility belt around her waist, armed with an arsenal of vampire-hunting gadgets".
• Figure 5: Generation toxicity is reduced with a larger negation coefficient but past a threshold, performance on control tasks degrades. The unique vector resulting from a larger pool of task vectors strikes a balance between toxicity and performance. See Appendix \ref{['sec:llamacheckpoints']} for the naming of task vector pools.

Theorems & Definitions (3)

• Definition 3.1: Projection Matrix golub_matrix_2013
• Proposition 3.1: Task Vector Decomposition
• Theorem 6.1: Subspace Intersection golub_matrix_2013