Task Arithmetic Through The Lens Of One-Shot Federated Learning
This work provides incremental improvements for researchers and practitioners in multi-task learning and model merging by bridging Task Arithmetic with Federated Learning.
The paper tackles the unclear factors behind Task Arithmetic's success in model merging by framing it as a one-shot Federated Learning problem, identifying data and training heterogeneity as key issues and showing that adapted Federated Learning algorithms can significantly boost merged model performance.
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.