Task Arithmetic Through The Lens Of One-Shot Federated Learning

TL;DR

This paper explores Task Arithmetic in model merging, revealing its equivalence to Federated Averaging and identifying key factors affecting its success, while proposing adaptations from Federated Learning to enhance performance.

Contribution

It establishes the theoretical connection between Task Arithmetic and Federated Averaging, and introduces federated learning algorithms to improve model merging effectiveness.

Findings

Task Arithmetic is mathematically equivalent to Federated Averaging.

Data and training heterogeneity significantly impact Task Arithmetic performance.

Adapting federated learning algorithms can boost merged model performance.

Abstract

Task Arithmetic is a model merging technique that enables the combination of multiple models' capabilities into a single model through simple arithmetic in the weight space, without the need for additional fine-tuning or access to the original training data. However, the factors that determine the success of Task Arithmetic remain unclear. In this paper, we examine Task Arithmetic for multi-task learning by framing it as a one-shot Federated Learning problem. We demonstrate that Task Arithmetic is mathematically equivalent to the commonly used algorithm in Federated Learning, called Federated Averaging (FedAvg). By leveraging well-established theoretical results from FedAvg, we identify two key factors that impact the performance of Task Arithmetic: data heterogeneity and training heterogeneity. To mitigate these challenges, we adapt several algorithms from Federated Learning to improve the effectiveness of Task Arithmetic. Our experiments demonstrate that applying these algorithms can often significantly boost performance of the merged model compared to the original Task Arithmetic approach. This work bridges Task Arithmetic and Federated Learning, offering new theoretical perspectives on Task Arithmetic and improved practical methodologies for model merging.