Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic

TL;DR

This paper introduces a method to improve weight disentanglement in task arithmetic by fine-tuning only the attention modules in transformers, reducing interference among tasks and enhancing model efficiency.

Contribution

The study demonstrates that fine-tuning only attention modules significantly enhances weight disentanglement in task arithmetic, offering an efficient alternative to NTK linearization.

Findings

Attention modules exhibit kernel behavior.

Fine-tuning attention modules improves weight disentanglement.

Representation modules are crucial for disentanglement.

Abstract

In recent years, task arithmetic has garnered increasing attention. This approach edits pre-trained models directly in weight space by combining the fine-tuned weights of various tasks into a unified model. Its efficiency and cost-effectiveness stem from its training-free combination, contrasting with traditional methods that require model training on large datasets for multiple tasks. However, applying such a unified model to individual tasks can lead to interference from other tasks (lack of weight disentanglement). To address this issue, Neural Tangent Kernel (NTK) linearization has been employed to leverage a "kernel behavior", facilitating weight disentanglement and mitigating adverse effects from unrelated tasks. Despite its benefits, NTK linearization presents drawbacks, including doubled training costs, as well as reduced performance of individual models. To tackle this problem, we propose a simple yet effective and efficient method that is to finetune the attention modules only in the Transformer. Our study reveals that the attention modules exhibit kernel behavior, and fine-tuning the attention modules only significantly improves weight disentanglement. To further understand how our method improves the weight disentanglement of task arithmetic, we present a comprehensive study of task arithmetic by differentiating the role of the representation module and task-specific module. In particular, we find that the representation module plays an important role in improving weight disentanglement whereas the task-specific modules such as the classification heads can degenerate the weight disentanglement performance. (The code is available at https://github.com/kyrie-23/task_arithmetic_tangent)