Task Arithmetic Through The Lens Of One-Shot Federated Learning
Authors: Zhixu Tao
Ian Mason
Sanjeev Kulkarni
Xavier Boix
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (05/2025)
Task Arithmetic Through The Lens Of One-Shot Federated Learning
Zhixu Silvia Tao'
skimu.aoo@@pinneton.ddd
Operations Reseorch and Financial Engineering
Princeton University
Ian Mason
imason &fujjitsu com
Pujilsu Reseorch of America
Sanjeev Kulkarni
kulkarni@@piinceton.cco
Operations Reseorch and Financial Engineering Electrical and Computer Engineering Princeton University
Xavier Boix
zboir com
Fujilsu Rescorch of America
Reviewed on OpenReview: https: // openreview. net/ formff id= Cgyo 7S7Dy0
Abstract
Task Arithmetic is model merging technique that enables the combination of multiple models' capabilities into single model through simple arithmetic in the parameter space, without the need for additional fine-tuning or access to the original training data. How- ever, the factors that determine the success of Task Arithmetic remain unclear. In this paper, we examine Task Arithmetic for multi-task learning by framing Federated Learning problem. We demonstrate that. Task Arithmetic is mathematically equivalent to the commonly used algorithm in Federated Learning, called Federated Aver- aging (FedAvg). By leveraging well-ostablished theoretical results from FedAvg, we iden- lify two key factors that impact the performance of Task Arithmetic: data heterogene- ity 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 boast performance of the merged model compared to the original Task Arithmetic approach. This work bridges Task Arithmetic and Federated Learning, offering new theoretical perspectives Task Arith- metic and improved practical methodologies for model merging. Our code is available at: wwm//iiviaaaa/////ooooooooaaaaal one shot
one-shot
arithmetic
1 Introduction
With the proliferation of fine-tuned models diverse domains, efficiently combining these models to achieve excellence across multiple tasks has emerged critical research challenge. Task Arithmetic (Ilharco et al. 2023). a recent technique in model merging, offers simple yet effective solution. Consider a set of T tasks, each defined by a training dataset and a loss function for fine-tuning. Let Opro E RRd be the parameters of a pre-trained model and 6, E Rd be the fine-tuned model parameters on task t. A task vector T, is generated by subtracting the pre-trained model parameters from the fine-tuned model parameters: 7. 8pre A.. The sum of these task vectors, Σi Tt, produces a direction that enhances performance of
Work done while Zhixu was
Fujitsu Research of America.

---

Published in Transactions on Machine Learning Research (05/2025)
the pre-trained model across multiple tasks for which the fine-tined models have been trained. In other words, updating the pre-trained model parameters by Opre + ΣE= results in a multi-task model. A key advantage of this approach is that it only involves dlement-wise operations in the weight space, eliminating the need for additional fine-tuning.
Despite its strong empirical performance, Task Arithmetic lacks substantial theoretical understanding. Only a small number of works have investigated this empirical success theoretically (Ortiz-Jimenez al.) 2023 Porrello et al., 2024). In this paper, we take step towards bridging the gap between theory and practice by framing Task Arithmetic a form of one-shot Federated Learning.
Federated Learning (McMahan et al. 2017], a distributed machine learning paradigm, enables devices to collaboratively train one shared model without exchanging the raw data. Federated Learning's goal is to retain data privacy and reduce computational costs, as all raw data remains stored locally on edge devices. In a typical Federated Learning training process, a server coordinates the training process by iterating through the following steps (Kairouz et al., 2021). First, the server broadcasts the current global model parameters and a training program to all the devices. Then each device locally computes an update to the model by using its own data. Finally, the server aggregates all the local updates from devices and updates the current global model by using the aggregated local updates. A commonly used algorithm for this training process is Federated Averaging (FedAvg) (McMahan et al. 2017). In one-shot Federated Learning, the server learns a global model in only a single round of communication between itself and all the devices (Guha et al., 2019
In this work, we show that using one-shot FedAvg is equivalent to Task Arithmetic, thus offering a new perspective on Task Arithmetic through the lens of one-shot Federated Learning. Using the connection between Federated Learning and Task Arithmetic, we can leverage the extensive thcorctical and algorithmic advancements in Federated Learning to better undersland when Task Arithmetic is effective and how it can be improved. To the best of our knowledge, this is the first study to bridge Federated Learning and Task Arithmetic. Our main contributions and the structure of the paper are summarized as follows.
we establish the connection bet wetween Task Arithmetic and one-shot Federated Averaging, formalizing Task Arithmetic by using notions from Federated Learning.
Bridge Task Arithmetic and Federated Learning: In Section 3
Analyze the Impact of Data and Training Heterogeneity in Task Arithmetic: In Section we use existing theoretical results from Federated Learning to deepen our understanding of Task Arithmetic. Specif- ically, we analyze how data heterogeneity, which slows convergence in FedAvg, also slows the convergence of Task Arithmetic (Section 4.11, and how training heterogencity, which causes objective inconsistencies in Federated Learning, similarly impacts Task Arithmetic (Section 4.2
Identify and Adapt Federated Learning Algorithms for Task Arithmetic: In Section we adapt several Federated Learning algorithms to mitigate the challenges posed by data and training heterogeneity, To achieve effective adaptions, we first discuss challenges in adapting Federated Learning algorithms (Section 5.1 and then we recommend four Federated Learning algorithms to enhance Task Arithmetic (Section 5
Experimentally Show That Federated Learning Algorithms Often Improve Task Arithmetic: We conduct experiments on vision-language model CLIP (Radford et al.) 2021) for image classification tasks in Scction and on large language model Llama2 (Touvron et al.] 2023) for instruction following, mathematical reasoning and code generation in Section 7 Our experiments confirm that adapting Federated Learning algorithms often improves the merged model's performance compared to Task Arithmetic.
2 Related Work
Our work bridges Federated Learning and Task Arithmetic, prominent approach within the growing domain of model merging. Thus, this section reviews related work on both model merging and Federated Learning.
Model Merging Task Arithmetic is one of many recent works on model merging (Wortsman et al. 2022a
Frankle et al., 2020b] Wortsman et al.) 20226) Matena & Raffel [022] Ainsworth et al. 2022 Don- Ychiya et al.) 2022). Though the term model merging relatively new, firstly formalized by Matena & Raffel (2022).
the concept has received significant investigation (Wortsman et al. 2022b Izmailov et. al.] 2018 Yang et al.
2

---

Published in Transactions on Machine Learning Research (05/2025)
2019
Tarvainen & Valpola 2017 Gupta et al. 2020 For example, Wortsman et al.
(202bb) average the pre-
trained model parameters and fine-tuned parameters to enhance the robustness of fine-tuned model against distribution shifts. Wortsman et al. (2022a) show that averaging parameters of multiple fine-tuned models with different hyperparamcter configurations can improve robustness and accuracy. [zmailov et al. (2018)
demonstrate that averaging parameters along a single SGD trajectory leads to better generalization
Task Arithmetic, introduced by Ilharco et al. (2023). refines model merging by introducing task vectors and hyperparameter À to control how much task vectors modify pre-trained model parameters. This method
has inspired various follow-up work on using simple arithmetic operations for model merging (Yu et al.
2024 Zhang et al. 2023; Ortiz-Jimenez ct al.) 2023 Stoica et al.
2023 Yadav et al.) 2024a; Yang et al.
2023 | such as sparsifying lask vectors (Yu et al. 2024), merging parameter--fficient modules (Zhang et al.
2023), fine-tuning in linearized model spaces
et al. 2023 and resolving sign interference of
task vectors (Yadav et al. 2024a).
A concurrent work of Task Arithmetic is Jin el al. (2022), ; who propose the RegMean. When merging two linear regression models, the model merging problem can be formulated as optimization problem which has closed-form solution. RegMean uses this solution to merge every linear layer of fine-tuned models. Unlike Task Arithmetic, RegMean is limited to linear layers and requires access to the input features of all linear layers. While Jin et al. (2022) also note that model merging is an extreme case of Federated Learning with single communication round, they do not explore this further. Moreover, RegMean focuses on merging linear layers, distinguishing it from our framework which addresses more general model merging problems.
There is substantial body of research dedicated to understanding the effectiveness of model merging Draxler et al. 2018 Entezari et al. 2021 Frankle et al. 2020a Garipov et al. 2018 Benton et al. 2021
Li et al. 2018 Foret et al.] 2020 Simsek et al.) 2021).
Some studies focus on the theory of linear model
connectivity (Draxler et al., 2018 Entezari et al. 2021 Frankle et al., 2020a] Garipov et al.] 2018)) while
others emphasize the flatness of the loss landscape (Benton et al. 2021 Li et al., (001)) Foret et al.. 2020
Simsek et al.] 2021). However, there has been relatively little work addressing the effectiveness of Task
Arithmetic except for Ortiz-Jimenes et al. (2023 Porrello et al. (2024). Our work aims to bridge this gap.
Federated Learning In Federated Learning, FedAug MMMMaaa et al., 2017 is widely used to solve the following distributed optimization problem across M devices: minger I.(0) = TV Em= T.m (0) where (0) == -.,...mm [e(O, xmm]] is the objective function on each device m, defined by some loss function é and data distribution Dm- The core idea behind FedAvy is to perform local Stochastic Gradient Descent (local SGD) each device, followed by model averaging on the server. A substantial body of research has analyzed
1
the performance of FedAvg and local SGD (Li et al. 2019 Karimireddy et al. 2020 Khaled et al. 2020:
Koloskova et al.) 2020 Woodworth et al.) 2020b a Glasgow et al. 2022 Patel et al. 2024 Wang et al.
2022
Woodworth et al.
2021
A key challenge in the theoretical analysis of FodAvg arises from data heterogeneity, where each device m has a different data distribution Dm In the homogencous setting where Dm = D Vm, Woodworth et al. (2021) have established the min-max complexity of Fed with smooth and convex loss functions. In the more complex heterogeneous setting, various works have derived the convergence rate of FedAvg under different assumptions about data heterogeneity (Li ct al.
2019
Khaled et al., 2020 Koloskova et al.) 2020 Woodworth
eet al. 2020b a: Glasgow et al. 2022 Patel et. al.
Wang el al. 2022 In this work, rather Ihan focusing on extending existing theoretical results, we focus on using these results to analyze Task Arithmetic.
2024
To address the challenges posed by data heterogencity, extensive rescarch has focusod on designing algorithms
to improve the performance of FedAvg (Karimireddy et al. 2020: Reddi et al., 2020 Li et al., 2020b Ye
et al. 2023b Li et al.
2021b Tenison et al.
2022 Wang et al. 2020).. Some work has enhanced optimization
algorithms by regulating the local training process (Karimireddy et al.) 2020 Reddi et al. 2020 Li et al.
2020b Wang et al.]
2020: Li et al.] 2021b). while others have proposed alternative aggregation methods
beyond simple averaging (Tenison et al. 2022 Ye et al. (202b)). Another line of work focuses on personalized
Fedcrated Learning (Smith et al. 2017 Tan et al.) 2022 Mansour et al. 2020 Li et al.
2021a Finn et al.
2017), addressing data heterogeneity by adapting the global model locally for each device.
3

---

Published in Transactions on Machine Learning Research (05/2025)
Aside from data heterogeneity, Wang et al. (2000) notice different local training process (which we refer to training hederogeneity) exacerbates the convergence of federated optimization algorithms, leading them to converge to a stationary point of an objective function inconsistent with the original objective function.
3 Task Arithmetic is One-Shot FedAvg
To deepen our understanding of the mechanism behind Task Arithmetic in multi-task learning, in this section, we establish a connection between one-shot FedAvg and Task Arithmetic.
Given T' tasks, the objective in multi-task learning to train model parameterized by E2 that performs well across all T tasks. This can be formulated minimizing the following multi-task objective function:
T Σ (0). L( = T 1=1
(1)
Here Le(0) B(,..,))>,,,(( (0(; 9t)] represents the objective function for lask I, where (12,3B) pair of input and output drawn from the data distribution D,, and f(.) denotes the loss function associated with the data and model. This formulation aligns with that used in Federated Learning, where each device t has local objective function L. L(0) is referred to as the global objective function.
In Federated Learning, the global objective function ( is often optimized using FedAvg. In FedAvg, cach local objective function is optimized through several iterations of Stochastic Gradient Descent (SGD), after which the server averages all the local updates. This process, also known as local Stochastic Grodient Descent (local SGD), is repeated over multiple communication rounds. Formally, given R communication rounds and initial global model parameters Ou, FedAvg follows the following update process Vr E [R]:
0 0,-1 o
AI E
(k++)) 0 0rr
Mi
(k) ((()
Ak: [0, 11, Vt [T]
T
K
0, = 0p-1 + T -
Σ k=0
(k) Thr 9e (
(2)
T
Here, 0, denotes the global model parameters obtained by optimizing the global objective at the end of the r-th communication round, and e denotes the local model parameters obtained by optimizing the t-th local (k) objective at the k-th local step during round r. The learning rate used for this step mh The stochastic gradient of the t-th local objective function L, is 9o()), and 8 is the outer step size used to aggregate all local updates. In the one-shot setting where R = 1, the update simplifies to the following:
T
T K
los = 00 + T
( - 00) 0o - 7
Σ
Σ
th
(k)
ge((((t)
(3)
k=
where los denotes the parameters generated by one-shot FedAvg,
In Task Arithmetic, the procedure mirrors the process in FedAvg. Each task / independently minimizes its own objective function L, by performing K, iterations of SGD with learning rates {np (0) (K+ 11)
..... :?):
starting from the same pre-trained model parameters 0g and converging to minimizer 02 E arg min L,, (0) This yiclds 07 9(0) - Ek
Kz-1 me (e(*)) where o(0) Bu Vt. The task vector T is defined by Tı 0; (0)
Using Task Arithmetic, a new set of parameters can be constructed
T 0u
T
K
0rA = 8
+
Σ 1=1
me
91 (
(4)
1=1
k=0
where À is a hyperparameter called the scaling coefficient [Taarco et al.
2023). which controls the extent
which the sum of task vectors is added back the pre-trained parameters. By comparing equations (3) and
4 we see that performing Task Arithmetic equivalent to one-shot FedAvg with outer step size / XT.
4

---

Published in Transactions on Machine Learning Research (05/2025)
4 Adapting Federated Learning Theory for Task Arithmetic
In this section, we extend theoretical insights from Federated Learning to Task Arithmetic, identifying two main factors that impact its performance: data heterogeneity and training heterogeneity. We analyze how these factors impact the convergence Task Arithmetic to the global minimum of convex objective function and the local minimum of a non-convex objective function.
4.1 Data Heterogeneity
This subsection is dedicated to understanding how data heterogeneity influences the performance of Task Arithmetic. Data heterogencity is common in Federated Learning and refers to the situation when data on each device is non-independent and identically distribuled (non-i.i.d.) (Li et al. 2020b; Ye el. al. 2023a
Wen et al. 2023p. In the context of Task Arithmetic, we define data heterogeneity as follows:
Definition 4.1. (Dala Helerogeneity) In Task Arithmelic, data helerogeneily refers to the non-i..d nature of training data across different tasks.
In the convergence analysis of FedAvg, data heterogeneity has been a longstanding issue. Given the con- nection between FedAvg and Task Arithmetic, in order to understand how data heterogeneity impacts Task Arithmetic, it is helpful to first review existing findings on data heterogeneity in FedAvg. We begin by in-
troducing several standard assumptions commonly used in the literature (Li et al. 2019 Karimireddy et al.
2020
Khaled et al., 2020 Koloskova et al.) 2020 Woodworth et al. 2020b
Glasgow ct al., 2022 Patel
et al. 2024 Wang el al. 2022 Woodworth et. al.) 2021).
Assumption 4.2. (Conuerity and Smoothness) Assume all the task objective functions L are conver and H -moooh.. That is, VI E [T] and v0, E R2, (o) S L((p) + (VL,())) - up) + =10 =0|²²
Assumption 4.3. (Bounded Stochastic Noise) The stochastic gradient computed by each task is unbiased with bounded voriance. That is, V0 .d,
E (( [V((0;
u:)] (V (0) and (( [IV u1) ([.()))
Assumption 4.4. (Bounded Initialization Ervor) Assume 40" E minger! L(0), 3B such that the inilial- ization 8n satisfies 110 - 11 < B.
To facilitate the analysis, we assume that all task objective functions are optimized with the same number of iterations, denoted K₁ = K Vt [2'], and that they use constant learning rates ny = Vt [7'] and Yk [K]. Additionally, set the outer step size to 1, reducing Task Arithmetic to model averaging.
Although there is no universal definition of data heterogeneity, several notions are commonly referenced in
the literature (Li et al. 2019 Karimireddy et al. 2020 Khaled et al. 2020 Koloskova et al. 2020
Woodworth
et al. 2020b a Glasgow et al. 2022 Patel et al.) 2024 Woodworth et al. 2021]. One widely adopted first-
order notion data heterogeneity is given by the following assumption (Koloskova al. 2020
Patel et al.
2024 Woodworth et al., 2020b Glasgow et al. 2022).
Assumption 4.5. (Bounded First-Order Data Heterogeneity at Optima) A set of objective functions {Le )7 satisfies the bounded first-order heterogeneity at optima if v6" E arg mingeR L(0), 3G. such that
{Lo)
T <L
The quantity ? measures the diversity among the set of functions (LiTT at the optima of the averaged multi-task objective function L()). Here, the notion of data heterogeneity is defined through objective func- tions, while whileeOriiz-iii Orrtiz-Jimenez et al. (2023) defines the Task Arithmetic property from the perspective of network functions. In Appendix A we further explore the connection between this notion of data heterogencity and the Task Arithmetic property proposed in Ortiz-Jimenez et al. (202)).
Using the notation from Patel et al. (024)), define learning problem satisfying Assumptions 4.2 4.3 4.4
and 4.5 belonging to class p B.B. Based on the upper bound from Koloskova et al. (2020) and the lower
bound Trom Patel et al. (2024) the following theorem characterizes the convergence of one-shot FedAvg.
5

---

Published in Transactions on Machine Learning Research (05/2025)
Theorem 4.6. Assume there is only one communication round R 1.
Then for any K 2,
min [4.1_, M.n.
- = HB HHOBB*
+
oB
+ (H
(5)
K1/3
VTK
and
max - 3 HB² + (Ho K1/3 VTK (1
oB
(6)
(1.)7 EP M.
Here, - and < denote inequalities that hold up to absolute constants. Based on the above theorem, we make several observations about Task Arithmetic. First, data heterogeneity C degrades the performance of Task Arithmetic. The term BH is non-vanishing error term introduced by (² highlighting the negative impact of data. heterogeneity.
Second, the one-shot learning nature of Task Arithmetic presents challenges that limit its performance. Notably, another non-vanishing term in Theorem 4.6 HB² arises due to the one-shot learning setup. In contrast, for FedAvg with It communication rounds, this term becomes HB² and diminishes as the number of Rt
K1/3 and aB decrease the number of <TK local steps K grows, they decay much more slowly the one-shot setting compared to rounds FedAvg, With multiple communication rounds, these terms given by (Ho"B*)1/ K1/5//255 and aB vTKЛ respectively (Patel 2024). This underscores the additional challenges introduced by the onc-shot learning paradigm. ct al.
communication rounds R increases. Morcover, although both (Ho2B4)1/8
Third, starting with a good pre-trained model is important. The influence of pre-training is captured by the term B introduced in Assumption 4.4 This quantity is a critical factor as appears in every error term, particularly in the non-vanishing term (HC3 Starting with well-suited pro-trained model that has a smaller B siguificantly mitigales the adverse effects of high data helerogeneily S as smaller B counteracts the beterogeneity. In fact, the significance of pre-trained models has been observed in experiments of both Task Arithmetic (Ortiz-Jimenez et al. 2023) and Federated Learning (Nguyen et al. 2022 Chen et al. 2022). Remark 4.7. The scaling coefficient À is important. As mentioned before, À 3. meaning that the scaling coefficient depends directly on the outer step size. Although in this section we assume 3 = 1 which yields == T and reduces Task Arithmetic to model averaging for simplicity, proper tuning of the scaling coefficient A is essential. Research indicates that the choice of /3 has a significant impact on FedAvg performance (Patel et a... 2024 Karimireddy et al. 2020 Charles & Konečný 2020 Jhunjhunwala et al., 2023 Malinovsky et al., 2023 et al. 2023b)). A similar sensitivity to À has been observed in Task Arithmetic: the performance of the final model depends heavily on A. For instance, Figure 15 in Ilharco et al. (2023) illustrates how Task Arithmetic's performance can vary dramatically with changes to X.
4.2 Training Heterogeneity
This subsection examines the effect of training heterogeneity on Task Arithmetic, which we define as follows: Definition 4.8. (Training Heterogeneity) In Task Arithmetic, training heterogeneity refers to the phc- nomenon where different task objective functions Lt are optimized with varying local learning rates and numbers of iterations during local training.
In the previous section, we assumed that all task objective functions are optimized with a fixed number of iterations K and constant learning rate n, However, in practice, each task objective function often optimized with different hyperparameter settings, which introduces training heterogeneity and can lead to objective inconsistency (Wang al. 2020). Now, we extend the setting in Section 4.1 to consider each task objective function L, being optimized with distinct hyperparameters (k) and K₁- We adopt notation and apply theoretical insights from Wang et al. (2020) to illustrate the impact of training heterogeneity.
First, let gr be the stochastic gradient of L,, and we define the following matrix of stochastic gradients for each task
G. (0))) (0(11))
a(((ki-1))) E Rddk..
6

---

Published in Transactions on Machine Learning Research (05/2025)
Next define the vector of normalized learning rates for each task t as
(0) T
1/ (K,
"
"
where 17 is a constant used to normalize the learning rates, whose purpose will be specified later. Using this for A as follows:
notation, we can rewrite equation
3 Σ n|aaا|| AA 80 A Σ G =00 -00-7
G [oll
(7)
where the second equality follows from À 4 mentioned in Section Next, we denote Taff 4 Ei= ||
the effective number of steps which measures the average number of updates accumulated using the the aggregation weight for task Then we can constant learning rate 7,, and denote wy [a,ll
||rlll
further rewrite equation 7
T
G4 |
OTA 00 - Toff
1
Notice the weight coefficients vector : differs from the original uniform coefficients 17: : #1 in the objective function from equation 1 This discrepancy is caused by training het- erogeneity, and leads Avd with multiple communication rounds converge to the stationary point of a different objective function L(0)= = E (0) which is inconsistent with the original objective function L. While Task Arithmetic involves only round of FedAvg, the inconsistency still remains due to training heterogeneity. Formally, present the following assumptions and adapt Theorems and 2 from Wang et al. 2020)) to illustrate this inconsistency in our setting.
Assumption 4.9. (Smoothness) Assume all the task objective functions I, are H-smooth. That is, VI E (T) and VO, & E Rdd (() - VL,(,)] < H|00 -11-- Assumption 4.10. (Bounded Gradient Heterogeneity) For any set of weights {wi)T- such that E we 1, there exist conslants a and such that v0 E ERdd
w. (V. < a3ll
Σ wi (L(0)))
Remark 4.11. Notice that Assumption 4.10 imposes more restrictive condition data heterogeneity com- pared to Assumption 4.5 Currently, no unified notion of data heterogeneity exists for Federated Learning. Since this section focuses on training heterogeneity, we adopt this more restrictive notion of data hetero- geneity, as done in Wang el al. 2020). to facilitate theoretical development.
Theorem 4.12. (Theorem and 2 from (Wang et al. 20201) Consider OT from update rule 7). Denote I(0) = Σ wrLe(0) and K a E K.. Let 17 V the following bound on the gradient norm [(Vi(OT
T/K. Under Assumption 4.3 4.9 and 4.10 we have
B[IVi(0TA)11I1 A(L((u) ( (K/TeD) VTK
4HO²A, VTK
6T
12TH2
+
(9)
K
K
Specifically, Linf inso I(0), A Tefi ΣT=
A2 -'T w:(llael1 - -a and As
max, {l Il ((aalll - am,11) where ar,-1 denotes the last coordinate of the vector az. Denote the RHS of Moreover, we have the following bound on gradient norm [|VL(OT)))
inequality g
E[]VL(0r)))1]1 2[x 1) + 1]e 2x
(10)
where X = ET (+ the chi-square divergence between the weight coefficient vectors = + 7 +] and [w1 Wa ... wr].

---

Published in Transactions on Machine Learning Research (05/2025)
The theorem above illustrates how heterogeneous local training affects the convergence of Task Arithmetic to stationary point of L. The convergence rate in 10 is influenced by two key factors: the term x Xpllu and the convergence rate to a stationary point of in (9.. For x xyites? when different training processes are used for each objective function, Xplte is non-zero, resulting in Xpllus a persistent error term. With training heterogeneity, Xe only vanishes if <² = O, indicating minimal data heterogeneity. This highlights the interaction between data heterogencity and training heterogencity: significant data heterogencity exacerbates the negative effects of training helerogeneity, intensilying the overall performance degradation.
When all task objective functions are optimized with the same number of iterations K and a consistent learning rate schedule
we have wt 7-- This yields Xpl|u 0 and L L, aligning the
actual objective function L being optimized with the original objective function L. In this sccnario, objective inconsisteney is effectively eliminated. The convergence rate in (10) is reduced to (
Remark 4.13. In Theorem 4.12 unlike in Theorem 1.6 we make no assumptions about the convexily of the objective functions, which naturally results in a looser convergence rate. Since the primary focus of this paper not on deriving a tighter convergence bound for non-convex scttings, we limit our analysis applying existing theoretical results understand the behavior of Task Arithmetie.
5 Adapting Federated Learning Algorithms for Task Arithmetic
In the previous section, we used insights from FedArg to analyze how data and training heterogeneity negatively impact Task Arithmetic. In order to address these challenges, numerous algorithms have been developed to improve FedAvg for more efficient Federated Learning (see Section 2 for related work). Thus, we adapt some Federated Learning algorithms to solve heterogeneity challenges in Task Arithmetic for better model merging performance. Selecting the right Federated Learning algorithms to implement requires a clear understanding of key challenges that complicate the adaptation. In this section, we begin with analyzing several key challenges.
5.1 Challenges in Adapting Federated Learning Algorithms for Task Arithmetic
First, the number of communication rounds is limited. Since Task Arithmetic is only one-shot Federated Learning, algorithms relying on multiple communication rounds are unsuitable. For instance, some Federated Learning algorithms add regularization terms to local objective functions (Li et al. 2020b] Durmus et al., 2021) to encourage local updates to remain close to the global model parameters transmitted from the previous communicatio round. However, in our onc-shot setting, only the pro-trained model parameters 6g are communicated. Applying such regularization would force each task's fine-tuned parameters to remain close to fo. potentially hindering both convergence and task-specific performance. Similarly, algorithms that address heterogeneity through iterative techniques such as variance reduction (Karimireddy et al.) 22000) Reddi et al. 2020) adaptive step size tuning (Jhunjhunwala et al. 2023 Li et al.) 2023b) require the accumulation and updating of certain metrics over multiple communication rounds. Since Task Arithmetic operates within a one-shot setting, implementing such iterative updates is impossible.
Second, no additional training is allowed. To counteract the effects of heterogeneity or limited communica- tion and data budget, some Federated Learning algorithms require additional training, imposing idditional computational cost. For instance. Ye et al. (2023b) ask each device to learn metrics that compare local and global data through training, and these metries are then used as additional scores for aggregation. [Zhag et al. (2022) train an additional generator to generate data which is subsequently used to train global model in one-shot Federated Learning setting.
Third, no additional datasets are available. Many Federated Learning algorithms rely on supplementary datasets, which is not feasible for modifying Task Arithmetic. In one-shot Federated Learning, a common approach to address data heterogeneity knowlodge distillation (Zhou et al., 2020 Guha et al. 2019 Li et al. 2020a). These methods often require access to extra datasets from which either local devices or the central server distills knowledge to improve model performance.
8

---

Published in Transactions on Machine Learning Research (05/2025)
Given the constraints and unique needs for adapting Federated Learning algorithms, we identify the following three criteria for selecting Federated Learning algorithms: adaptability to one-shot setting, no additional training required and no access to additional datasets required.
5.2 Algorithms
With criteria established above, we now explore four Federated Learning algorithms FedNova WWang et al. 2020), FedGMA (Tenison et al. 2022), Median (Yin et al. 2018) and CCLIP (Karimireddy et al. 2021] that align with these guidelines, explaining their motivations and how they modify Task Arithmetic. For more detailed explanation and further understanding of these algorithms, please refer to the original papers.
Fed (Wang et al., 2020) FedNova addresses objective inconsistency caused by training heterogene- ity by replacing the heterogeneous weight. vector [w] w2 ... wT] with the uniform weight vector 1 ... 7). ensuring consistent weighting across tasks. This approach adapts easily to the one-shot setting. Using the 4.2 FedNova modifies the Task Arithmetic update as:
notation from Section
OTA = 00 Toff
T
G.a,
0o + 7
T Σ (alll))
Σ t=
Tt
(11)
where Tell- = 2/ E | the effective number of steps defined in Section
4.2
To
G/Gra is the task
vector, and À = f is the scaling coefficient. In other words, FedNova normalizes each task vector by ||a]] and rescales the scaling coefficient by the average + Ei=nllaell ||
FedGMA (Tenison et al., 2022) FedGMA addresses data heterogencity by mitigating sign conflicts among local updates. Such sign conflicts can cause information loss and slower convergence. FedGMA uses gradient mask to reduce the impact of conflicting directions and preserves meaningful information. Specifically, it computes an agreement score A to measure alignment across task vectors (Te)²=)- Based on a threshold p. FedGMA constructs a mask M that emphasizes coordinates with strong agreement while if 2 where
1,
reducing the influence of others. Formally, define A D sign(T.) and My
Aj+ otherwise
denotes the j-th coordinate and sign(-) is applied in coordinate-wise manner. This yields
OT = 0o AM O
(12)
Median (Yin et al., 2018) Coordinate-Wand Wise Median Yin el al. 2018), originally designed to handle adversarial updates in Federated Learning, adapted here to address data and training heterogeneity in Task Arithmetic. Due to diverse data distributions or differing hyperparameter settings, some task vectors may have extreme values. By selecting the median value for each coordinate, this method reduces the influence of outliers while maintaining overall performance across tasks. It modifies Task Arithmetic
OTA 0 A medd(T.....,
(13)
where med(-) computes the coordinate-wise median of {Ta)i'
CCLIP (Karimireddy et al., 2021) CCLIP, short for centered clipping, is another widely applied robust aggregation method towards adversarial devices in Federated Learning. With the same motivation to using the Coordinate- Desei Median, we use CCLIP to reduce the impact of extreme task vectors. CCLIP is implemented with a predefined threshold and modifies Task Arithmetic as follows:
0 = 0o + Σ Ta (in
P [|T]|]
(14)
When the norm of task vector |T|||| exceeds the threshold p, this method identifies it as an outlier and shrinks its magnitude by factor of 2 [
9

---

Published in Transactions on Machine Learning Research (05/2025)
6 Experiments on Merging CLIP
In this section, we present and discuss our experimental results on CLIP-ViT-B-32 Radford ct al. 2021 for image classifications. We follow the same experimental paradigm as (Ilaacco el al. 2023). Specifically, we
use CLIP-ViT-B-33 [Radford et al. 2021 as the pre-trained model and eight datasets: Cars Krause et al.
2013). DTD (Cimpoi al., 2014). EuroSAT (Helber et al., 2019). GTSRB (Stallkamp et al. 2011), MNIST
(LeCun 1998). RESISC45 (Cheng et al. 2017), SUN397 (Xiao al. 2016 and SVHN (Netzer et al. 2011 to constrruct eight task vectors.
In total, there are 247 ways to select T different task vectors from these eight task vectors where T E [2, 8] 247 Er= (7). For each algorithm, therefore conduct 247 experiments. In each experiment, merge T selected task vectors and evaluate on the T' datasets corresponding to the task vectors used. Our evaluation metric is normalized accuracy (Ilharco et al. 2023 defined as the test accuracy normalized by the fine-tuned model's accuracy. That is,
accuracy on task
normalized accuracy on task t accuracy of the fine-tuned model t task t
6.1 Experimental Results
In the first part of the experiments (Section 6.1.1 to simulate practical conditions of training heterogeneity, we fine-tune CIIP-ViT-B-32 on each dataset using three learning rates (1e-4, le-5, 1e-6} and different numbers of iterations. Then we select the best fine-tuned checkpoints via cross-validation on validation datasets. Refer Appendix B.1 for further details on fine-tuning and cross-validation.
In the second part of the experiments (Section 6.1.2 to better understand the impact of training heterogene- ity, we the task vectors provided by [lharco et al. (2023), which were fine-tuned with the same number of iterations and the same learning rates, thereby eliminating training heterogeneity.
6.1.1 Merging with training heterogeneity
We first report experimental results using task vectors fine-tuned with training heterogencity. Table 1 summariizes the performance of various methods in specific experimental setup: merging all eight task vectors, corresponding to the scenario where T = 8. We report the average normalized accuracy as well as the normalized accuracy for each dataset. Task Arithmetic is used as the baseline method for comparison. As shown in the table, all four adapted Federated Learning methods outperform the bascline by a substantial margin. Moreover, we observe that. Median and CCLIP yield the most improvement.
Average
Methods
Normalized Accuracy
DTD
EuroSAT
GTSRB
SUN397
SVHN
MNIST
Cars
RESISC45
Task Arithmetic
67.33
57.43
53.86
41.00
82.40
78.58
87.76
71.94
65.72
Median
74.55 (t 7.22)
67.51
78.05
67.12
84.02
56.69
91.32
77.51
74.17
FedNova
69.57 (t 2.24)
57.08
50.37
61.47
86.62
77.68
74.22
63.96
FedGMA
68.55 (1 1.22)
60.01
58.69
45.02
84.13
71.19
86.65
74.33
68.13
CCLIP
74.82 (f 7.49)
66.76
75.42
73.87
83.18
58.51
92.03
76.40
72.39
Table 1: Merging all eight task vectors using five different methods. Each method is evaluated eight datasets, with normalized accuracy reported for each dataset. The highest and sccond-highest normalized accuracy values for each dataset are highlighted in bold and underlined, respectively.
Table 2 and Figure summarize the performance comparison between Task Arithmetic and other Federated Learning algorithms across 247 experiments. In Table 2 we report the percentage of 247 experiments in which the average normalized accuracy improves, remains unchanged, or degrades when using four Federated Learning methods compared the bascline, Task Arithmetic. The average normalized accuracy is calculated by averaging over the number of task vectors being used. In order better visualize the performance differences for each method, in Figure use histograms to show the frequencies of experiments within
10

---

Published in Transactions on Machine Learning Research (05/2025)
Percentage of Improved Experiments Compared to Task Arithmetic
Percentage of Unchanged Experiments Compared to Task Arithmetic
Percentage of Degraded Experiments Compared to Task Arithmetle
Meshods
Median
67. 07.61%
0%
32.39%
FedNova
63. 56%
0%
36. 11%
FedGMA
40.49%
18.22%
41 41.29%
CCLIP
91.50%
8.5%
Table Percentage of improved, unchanged, and degraded experiments using different methods compared to Task Arithmetic.
Median
FedNova
80
50
40
Count
30
30
20
-30%
-20%
-10%
0%
10%
20%
30%
-30%
-20%
-10%
0%
10%
20%
30%
Change Average Normalized Accurary
Change Average Normalized Accuracy
FedGMA
CCLIP
80
Count
60
Count
40
10
30%
20%
10%
0%
10%
20%
30%
30%%
-00%
-10%
10%
20%
30%
Change in Average Normalized Accuracy
Change in Average Normalized Accuracy
Figure 1: Histograms showing the change in average normalized accuracy for four different methods compared to Task Arithmetic. The x-axis shows the difference in average normalized accuracy relative to Task Arithmetic, with positive values indicating improvement and negative values indicating degradation. The y-axis represents the number of experiments within each range of change values.
each range of change in average normalized accuracy. Median, FedNova, and CCLIP consistently show the ability to improve upon Task Arithmetic in most cases, while FedGMA typically demonstrates either no change or slight improvements.
6.1.2 Merging without training heterogeneity
In Section 4.2 we analyzed how training heterogeneity causes objective inconsistency, degrading Task Arith- metic's performance. While in the experiments conducted by [[lhacco et al. (2023), all task vectors were homogeneously fnn-ttned using a consistent learning rate 1e-5, our approach in Section 6.1.1 employs het- erogeneous fine-tuning. We compare the performance of Task Arithmetic on these two sets of task vectors to validate our theoretical findings in Section 4.2
Table compares Task Arithmetic's performance heterogeneously and homogeneously fine-tuned task vectors when merging all eight task vectors. Again we report the average normalized accuracy and the normalized accuracy for each dataset. As evident from the table, Task Arithmetic with homogeneous fine- tuning consistently outperforms its heterogeneous counterpart across all datasets, except for SUN397.
Table and Figure compare Task Arithmetic's performance on homogencously and heterogencously fine- tuned task vectors across 247 experiments. Homogeneous fine-tuning outperforms heterogeneous fine-tuning in 92.31% 1 of cases, shown in Table Moreover, Figure shows that homogeneous fine-tuning improve
11

---

Published in Transactions on Machine Learning Research (05/2025)
the average normalized accuracy by up to more than 30%. These results highlight the significant negative impact of training heterogeneity on the performance of Task Arithmetic.
Average Normalizod Accuracy
Methods
DTD
EuroSAT
GTSRB
SUN397
SVHN
MNIST
RESISC15
Task Arithmetic with Heterogenoous Fine- Bun-Tunng
67.33
37.43
53.86
41.00
82.40
78.58
87.76
71.94
65.72
Task Arithmetic with Homogeneous Fine Tuning
77.34 (1 10.01) 64.90
77.93
69.47
80.64
80.26
96.42
75.98
73.01
Table 3: Using Task Arithmetic to merge eight task vectors from heterogeneous and homo- gencous fine-tuning processes. Each method is evaluated on eight datasets, with normalized accuracy reported for each dataset.
Percentage of Improved Experiments
Percentage of Unchanged Experiments
Percentage of Degraded Experiments
Task Arithmetic with Homogeneous Fine-Tuming
92.31%
0%
7.09%
Table 4: Percentage of improved, unchanged, and degraded experiments using task vectors with homogeneous fine-tuning process compared to those with heterogeneous fine-tuning process.
Task Arithmetic with Homogeneous Fine-Tuning
40
Count
30
-40% -30%
-20%
-10%
0%
10%
20%
30%
40%
Change in Average Normalized Accuracy
Figure Histogram showing the change in average normalized accuracy when using task vectors from homogeneous fine-tuning compared to heterogeneous fine-tuning. The x-axis shows the difference in average normalized accuracy relative to Task Arithmetic, with positive values indicating improvement indicating degradation. The y-axis represents the number of experiments within each range of change values.
negative values
6.2 Discussion on Experimental Results
We now discuss key observation from our experimental results: practice, training heterogeneity poses a greater challenge than data heterogeneity for Task Arithmetic.
highlights how data heterogeneity degrades Task Arithmetic, our experimental results in Section 6.1 show it is less problematic compared to training heterogeneity. First, FedNova, designed to address training heterogeneity, consistently outperforms Task Arithmetic more frequently and significantly than FedGMA which targets data heterogeneity. Table 2 and Figure I demonstrate that FedNova both improves performance more frequently and achieves greater overall gains than FedGMA.
While Section 4.1
Second, among Federated Learning algorithms, CCLIP and Median demonstrate the best performance. As discussed in Section 5. these methods are designed for robust aggregation in the presence of outliers. In setting, they effectively address training heterogeneity which causes certain task vectors to have
12

---

Published in Transactions on Machine Learning Research (05/2025)
disproportionately large norms and behave like outliers. For example, the cross-validation process selects much larger learning rate of le-4 for SVHN, compared to le-5 used for other datasets. This hyperparameter setup results in the SVHN task vector having a significantly larger norm (reported in Appendix [.1.5]] making an outlier that nogatively impacts the merged model's performance on other tasks when using Task Arithmetic. By employing robust aggregation methods like Median and CCLIP, we reduce the influence of the SVHN task vector, which improves the merged model's performance on other tasks.
Third, when comparing Table 2 and Table
that homogencous fine-tuning leads to more frequent improvements over Task Arithmetic compared to the other four algorithms Median, FedNova, FedGMA and CCLIP. Similarly, Figure 2
demonstrates that homogencous fine-tuning results in the most frequent and substantial positive changes in average normalized accuracy.
Further evidence is presented in Appendix B.2 where we evaluate the performance of Median, FedGMA and CCLIP on task vectors generatod via homogencous fino-tuning. Using the performance of Task Arithmetic these homogeneously fine-tuned task vectors as the baseline, we find that Federated Learning algorithms rarely improve upon the baseline. In fact, Task Arithmetic consistently emerges as the best-performing approach when merging those task vectors generated without training heterogeneity. This reinforces our observation that training heterogeneity is a more significant issue than data heterogeneity in practice.
7 Experiments on Merging LLMs
We now present and discuss experimental results merging LLMs for three tasks: instruction following, mathematical reasoning, and code generation. We follow the experimental paradigm of Yu et al. (2024). We merge task vectors from three models WizardLM-133B (Xu et al. 2023), WizardMath-11e (Luo et al. 2023 and 1ainaa---33-B- Alpaca (Chaudhary, 2023) for instruction following, mathematical reasoning, and code generation, respectively. All three models are fine-tuned from Alama2-1313 (Touvron et al. 2023 For instruction following, we evaluate the models on AlpacaEval (Li et al.
2023a For mathematical reasoning, 2021 For code generation, we use HumanEval (Chen et al.] [022]] and MBPP (Austin et al.) 2021p. Performance metrics include win rate for AlpacaEval, zero-shot accuracy for GSMSK and MATH, and pass@1 for HumanEval and MBPP.
we use GSM8K (Cobbe al. 2021) and MATH (Hendrycks et al.
Since all models used in this experiment are downloaded from HuggingFace, we do not have access to their fine-tuning hyperparameter settings. As a result, FedNova cannot be applied in this experiment because it requires knowledge of learning rates and the number of iterations, which are unavailable. Furthermore, when implementing Median, taking the median of two vectors reduces to averaging, which equivalent Task Arithmetic. Thus, implement Median only for merging three task vectors, with the corresponding results deferred to Appendix C.2 For additional details on experiments, please refer to Appendix C
7.1 Experimental Results
In Table 5 we compare the performance of thrce methods: Task Arithmetic, FedGMA, and CCLIP. The results show that. when merging two out of three task vectors, FedGMA and CCLIP often outperform Task Arithmetic. However, when merging all three task vectors, Task Arithmetic shows superior performance code generation and instruction-follovii tasks. Notably, Task Arithmetic consistently excels in instruction- following tasks, achieving either the highest accuracy or accuracy comparable to the other methods.
7.2 Discussion on Experimental Results
In this section, discuss key observation from our experimental results: only from differences in hyperparameters but also from variations in tuning methods.
training heterogeneity arises not
In Section 4.2 we theoretically analyzed how using different learning rates and number of ilerations creates training heterogeneity and thus leads to objective inconsistency. However, our experimental results in Section revcal that employing different fine-tuning methods further exacerbates training heterogeneity. Specifically, in our experiments, the 11am--2-13-- 2-133BCode--B Code-Alpaca Alpaca model fino-tuned using QLoRA (Dettmers et al. 2024 parameter-efficien fine-tuning (PEFT) approach.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Mathematical Reasoning
Code Generation
Instruction
Tasks
Methods
Following
GSM8K
MATH
HumanEval
MBPP
AlpacaEval
Task Arithmetic
64.22
14.1
1.22
8.66
Math Code
FedGMA
65.5
12.66
15.85
21.8
CCLIP
65.81
13.48
4.27
Task Arithmetic
65.88
13.32
69.96
Instruction Math
FedGMA
66.72
14.48
62.04
CCLIP
64.75
13.18
69.99
Task Arithmotic
32.32
32.2
79.76
Instruction
FodGMA
20.12
49.55
Code
CCLIP
32.32
34.2
76.02
Instruction
Task Arithmetic
58.45
12.06
25.16
31
70.89
Math
FedGMA
57.16
11.96
20.12
27.4
64.13
Code
CCLIP
62.93
12.96
20.12
27.6
66.91
Table 5: Performance of merging LLMs. The best performance for each dataset. is highlighted in bold.
PEFT adjusts only small subset of parameters while leaving the rest unchanged (Han et al. 2024 Conse- quently, task vectors generated by PEFT typically have smaller norms compared to those generated through standard fine-tuning. In our experiments, Ala-ma--133--oode-AA Alpaca, which is fine-tuned for code generalion using PEFT, produces a task vector with a notably small norm of 5.05. In contrast, WizardL M-1B and WizardMath-133. fine-tuned for instruction following and mathematical reasoning via standard fine-tuning, generale task vectors with much larger norms of 142.61 and 52.62, respectively. This discrepancy can pose challenges when merging task vectors. Simply regulating the behavior of task vectors with larger norms can lead to unintended negative effects. As shown in Table 5 when merging tasks include both instruction following (large norm) and code generation (small norm), FedGMA and CCLIP either fail to outperform Task Arithmetic or achieve comparable performance on these two tasks. This highlights that addressing training heterogeneity by focusing solely on differences in hyperparameters is insufficient in practice.
While some studies have explored the challenges of merging large models fine-tuned via PEFT (Zhang et al.
2023 Wu al. 2024), merging PEFF-genneaate task vectors with those produced by standard fine- tuning remains an open research question. Further investigation is required to design effective strategies for combining such task vectors in Task Arithmetic. Additionally, more research is needed to develop robust aggregation methods in Federated Learning to address this type of practical training heterogeneity.
8 Limitations and Conclusion
Limitations First, Federated Learning, there remains a gap theoretical understanding and algorith- mic development for addressing data hetcrogeneity in large-scale experiments. This might arise from the lack of universal notion of data heterogeneity. While existing theoretical work primarily focuses on first-order heterogeneity such as Assumption 4.5 recent studies (Patel et al. 2024 Wang et al. 2022) suggest that this perspective may be insufficient to fully capture data heterogeneity in practice. This underscores the need for more refined definitions of data heterogeneity and the development of corresponding algorithms.
Sccond, a recent study (Yadav et al. 2024b) shows that when merging very large models such as PaLM-2-64B (Anil et al. 2023), the differences in performance between various merging methods tend to diminish. While Section 4.1 discusses how high-quality pre-trained model mitigates the adverse effects of data heterogeneity, our framework does not yet provide theoretical explanation for why the large size of the pre-trained model can also alleviate data heterogeneity, and we leave this as an open question for future work.
Conclusion In this paper, we established a connection between Task Arithmetic and one-shot Feder- ated Learning. By leveraging thcoretical insights from Federated Learning, we identified and analyzed two key sources of heterogeneity- data heterogeneity and training heterogeneity and their impact on Task Arithmetic. Also, adapted Federated Learning algorithms for model merging, demonstrating their great
14

---

Published in Transactions on Machine Learning Research (05/2025)
potential to significantly improve over the performance of Task Arithmetic. We hope this work serves foundation for advancing the understanding, enhancing the algorithms, and expanding the applications of Task Arithmetic through the lens of Federated Learning.
Author Contributions
X. Boix ideated the research; Z. Tao conceptualized the theoretical part of the research with contributions from I. Mason and X. Boix; Z. Tao and I. Mason conceptualized the experimental part of the research with contributions from X. Boix; Z. Tao wrote the code and ran the experiments with contributions from I. Mason; Z. Tao analyzed the experimental results with contributions from Mason, S. Kulkarni and X. Boix; Z. Tao wrote the paper with contributions from I. Mason, S. Kulkarni and X. Boix; K.Kularr and X. Boix supervised the project.
Acknowledgments
We thank the reviewers and action editor for their insightful comments and constructive feedback, which have significantly improved the quality of this work. We are also grateful to Kasper Vinken and Mehdi Bahrami for valuable discussions and feedback throughout the project.
References
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXie preprint 2XXx:2209.04836 2022.
Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, al. Palm technical report. arXiv preprinl arXi::2305 10403, 2023.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.077 07732, 2021.
Gregory Benton, Wesley Maddox, Sanae Lotfi, and Andrew Gordon Gordon Wilson. Loss surface simplexes for mode connecting volumes and fast ensembling. In International Conference Machine Learning, pp. 769 779. PMLR, 2021.
Zachary Charles and Jakub Konečný, On the outsized importance of learning rates in local update methods. arXiv preprint trXtc:2007000877 2020.
Sahil Chaudhary. Code alpaca: An instruction-follovini Ilama model for code generation. https://github can/shh112801144cdeeeapece. 2023.
Hong-You You Chen, Cheng-Hao Tu, Ziwei Li, Han-Wei Shen, and Wei-Lun Chao. On the importance and applicability of pre-training for federated learning. arXie preprind ar arXiv:2206. 111.2006..1488. 2022.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint ar arXio:2107.08371, 2021.
Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Benchmark and state of the art. Proceedings of the IEEE, 105(10):1865 1883, 2017.
Mircea Cimpoi, Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. Describing textures in the wild. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3606 3613, 2014.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Retcchrr Nakano, al. Training verifiers solve math word problems. arXiv preprint ar Xin221 10. 14168 2021.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Newral Information Processing Systems, 36, 2024.
Shachar Don- Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Cold fusion: Collaborative descent for distributed multitask finctuning. arXiv preprint arXiv:211201197 2022.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309-1318. PMLR, 2018.
Alp Emre Durmus, Zhao Yue, Matas Ramon, Mattina Matthew, Whatmough Paul, and Saligrama Venkatesh. Federated learning based on dynamic regularization. In International conference on learn- ing representations, 2021.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in lincar mode connectivity of neural networks. arXiv preprint Xiv:2110.062900 2021.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXi::2010.014: 2020.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. PMLR, 2020a.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2259-3269 PMLR, 13-18 Jul 2020b. URL https://proceedinng mlr wttps////111tt/ttele/e// html
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew Wilson. Loss surfaces, mode connectivity, and fast, ensembling of duns. Advances in neural information processing systems, 31, 2018.
Margalit R Glasgow, Honglin Yuan, and Tengyu Ma. Sharp bounds for federated averaging (local sgd) and continuous perspective. In International Conference on Artificial Intelligence and Statistics, pp. 9050 9090. PMLR, 2022.
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local syd generalize better than sgd? arXiv preprint 2xXt:22333.0125, 2023.
Neel Guha, Ameet Talwalkar, and Virginia Smith.
One-shot federated learning. arXio preprint
arXiv:1902. 11175, 2019.
Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Large- batch training that generalizes well. arXiv preprint ar XrXv:2000.00311. 2020.
Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. cerameter-offccionn fino-tuning for large models: comprehensive survey. ar preprint Xiv::22403. 14608, 2024.
Patrick Helber, Benjamin Bischke, Andreas Dengel, and Damian Borth. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 12(7):2217-2226 2019.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103003874 2021.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidtt, Hannanch Hajishirzi, and Ali Parhadi. Editing models with task arithmetic. In The Fleventh International Conference on Learning Representations, 2023. URL
jrj
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint XXiu:1803.054400 2018.
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp: Speeding up federated averaging via extrapolation. arXiv preprint 000,000000000 2023.
Xisen Jin, Xiang Ren, Daniel Preotiuc-Piete and Pengxiang Cheng, Dataless knowledge fusion by merging weights of language models. arXiv preprint arXiv:222 09849, 2022.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and trends@ in machine learning. 14(1-2))1-210, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference machine learning, pp. 5132-5143. PMLR, 2020.
Sai Pranceth Karimireddy, Lie He, and Martin Jaggi. Learning from history for byzantine robust optimiza- tion. In International Conference on Machine Learning, pp. 5311 5319. PMLR, 2021.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519-4529.
PMLR, 2020.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory decentralized sgd with changing topology and local updates. In International Conference on Machine Learning, pp. 5381-5393. PMLR, 2020.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained catego- rization. In Proceedings of the IEEE international conference computer vision workshops, pp. 554 561, 2013.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/erdb/mnistt 1998.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
Qinbin Li, Bingsheng He, and Dawn Song. Practical one-shot federated learning for cross-silo setting. arXiv preprint arXiv:2010.010 01011,,, 2020a.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning, In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10713-10722, 2021a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Amect Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-455 2020b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data, orXiv preprint orXio:1907.02188. 2019.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning non-iid features via local batch normalization. arXiv preprint 2rXii:2102.076223 2021b.
Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpacaeval: An automatic evaluator of cndtruction-oolowiiin models. https: ///gituub.co//taaatuutu wwtpe/llllaaaaaaaaaaaaaaaaaaaaaa alpaca_evaaa eval. 2023a.
17

---

Published in Transactions on Machine Learning Research (05/2025)
Zexi Li, Tao Lin, Xinyi Shang, and Chao Wu. Revisiting weighted aggregation in federated learning with neural networks. In International Conference on Machine Learning, pp. 19767-19788. PMLR, 2023b.
Tao Lin, Sebastian Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini hani-batches, use local sgd. arXiv preprint Xiv: 100::18080000 008.07217.00 2018.
Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qingwei Lin, Shifeng Chen, and Dongeei Zhang. Wizardmath: Empowcring mathematical reasoning for large language models via reinforced evol-instruct. arXiv preprint arXiv:2308.008 09583, 2023.
Grigory Malinovsky, Konstantin Mishchenko, and Peter Richtárik. Server-side stepsizes and sampling without replacement provably help in federated optimization. In Proceedings of the 4th International Workshop on Distributed Machine Learning, pp. 85-104, 2023.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personaliza- tion with applications to federated learning. arXiv preprind 2ix::2002.1661, 2020.
Michael S Matena and Colin A Raffel. Merging models with isher-weighted averaging. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Ad- vances in Newral Information Processing Systems, volume 35, pp. 17703 17716. Curran As- sociates, Inc., 2022. URL https://proceeddn// nourigs...ur...... nouripp.ccc/ppper htaappscccppppe_flleeeeee fil.es//appor//222///1 7D22687700E44000661111110666
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. cettuumiiiatii--ffiiiieee Learning of Deep Networks from Decentralized Data. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference Artificial Intelligence and Statistics, vol- PMLR, 20 22 Apr 2017. URL
une 54 of Proceedings of Machine Learning Research, pp. https::/proceeddae mmw.e...........mmm.m.m.....mm.. press//54/mcmmhanl// html
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Reading digits natural images with unsupervisod fcature lcarning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. 4. Granada, 2011.
John Nguyen, Jianyu Wang, Kshitiz Malik, Maziar Sanjabi, and Michael Rabbat. Where to begin? on the impact of pre-training and initialization in federated learning. arXiv preprinl arxi 15387, 2022.
Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances én Neural Information Processing Systems, volume 36, pp. 66727-66754. Cur- ran Associates, Inc., 2023. URL httpp::/ppooceeddiiiiit wwweseseeriiippeee cc/ accppaprr__ files//papprr/223/fiiel/ 22607688160004433544111126ee0 Paper Pwwer-aoofeeeeeeeeeeeeee pdf
weurips.cc/papee
Kumar Kshitij Patel, Margalit Glasgow, Ali Zindari, Lingxiao Wang, Sebastian U Stich, Ziheng Cheng, Nirmit Joshi, and Nathan Srebro. The limits and potentials of local sgd for distributed heterogencous learning with intermitient communication. arXie preprinl arXin:2405.1166; 2024.
Angelo Porrello, Lorenzo Bonicelli, Pietro Buzzega, Monica Millunzi, Simone Calderara, and Rita Cuc- chiara. A second-order perspective on compositionality and incremental learning. arXiv preprint arXiv:2405.1635) 2024.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8700 PMT.R, 2021.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. aix preprint 2rXiv:2008.00295, 2020.
18

---

Published in Transactions on Machine Learning Research (05/2025)
Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape in overparameterriee neural networks: Symmetries and invariances. In International Conference on Machine Learning, pp. 9722 9732. PMLR, 2021.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Aduances in neural information processing systems, 30, 2017.
Johannes Stallkamp, Mare Schlipsing, Jan Salmen, and Christian Igel. The german traffic sign recognition benchmark: a multi-class classification competition. In The 2011 international joint conference on neural networks, pp. 1453-1460. TREE, 2011.
George Stoica, Daniel Bolya, Jakob Bjorner, Pratik Ramesh, Taylor Hearn, and Judy Hoffman. Zipit! merging models from different tasks without training. arXiv preprint arxiv:2305.00 03053,, 2023.
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE tronsactions on neural networks and learning systems, 34(12))9587 9603, 2022.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. Advances in newral information processing systems, 30, 2017.
Trene Tenison, Sai Aravind Sreeramadas, Vaikkunth Mugunthan, Edouard Oyallon, Trina Rish, and Eugene Belilovsky. Gradient masked averaging for federated learning. arXiv preprint arXic:2201.0 11000..19910 11986, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bash- lykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama Open foundation and fine-tuned chat models. Xiv preprint arXiv:2307.09288 2023.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33: 7611-7623, 2020.
Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. arXiv preprint Xrx: 0473,, 2022.
Jie Wen, Zhixia Zhang, Yang Lan, Zhihua Cui, Jianghui Cai, and Wensheng Zhang. A survey on federated learning: challenges and applications. International Journal of Machine Learning and Cybernelics, 14(2): 513 533,535,22 2023.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Memahan, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? In International Conference on Machine Learning, pp. 10334 10343. PMLR, 2020a.
Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogencous distributed learning. Advances in Neural Information Processing Systems, 33:6281 6292, 2020b.
Blake E Woodworth, Brian Bullins, Ohad Shamir, and Nathan Srebro. The min-max complexity of dis- tributed stochastic convex optimization with intermittent communication. In Conference on Learning Theory, pp. 1386-4437. PMLR, 2021.
Mitchell Wortsman, Gabriel Ioharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Mor- cos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Korbblith, al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Interna- tional conference machine learnóng, pp. 11965-23999 PMLR, 2022a.
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca. Roelofs, Raphael Gontijo Lopes, Hannanch Hajishirzi, Ali Farhadi, Hongscok Namkoong, et al. Robust fine- tuning of zero-shot models. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 2022b.
19

---

Published in Transactions on Machine Learning Research (05/2025)
Xun Wu, Shaohan Huang, and Furu Wei. Mixture of lora experts. arXiv preprint arXiv:2404.1363. 2024.
Jianxiong Xiao, Krista Ehinger, James Hays, Antonio Torralba, and Aude Oliva. Sun database: Exploring large collection of scene categories. International Journal of Computer Vision, 119:3-22 12:3-22,222 2016.
Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin Jiang. Wizardlm: Empowering large language models to follow complex instructions. arXiv preprint (al:2334,122444 2023.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin A Raffel, and Mohit Bansal. Ties-merging: Resolving interference when merging models. Advances in Neural Information Processing Systems, 36, 2024a.
Prateek Yadav, Tu Vu, Jonathan Lai, Alexandra Chronopoulou, Manaal Faruqui, Mohit Bansal, and Tsend- suren Munkhdalai. What matters for model merging at scale? arXiv preprint arXiv:24 10.1107,........ 2024b.
Enneng Yang, Zhenyi Wang, Li Shen, Shiwei Liu, Guibing Guo, Xingwei Wang, and Dacheng Tao. Adamerg- ing: Adaplive model merging for multi-Lask learning. arXiv preprint Kiw:2310.02575, 2023.
Guandao Yang, Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning, pp. 7015 7024. PMLR, 2019.
Mang Ye, Xiuwen Fang, Bo Du, Pong C Yuen, and Dacheng Tao. Heterogeneous federated learning: State- of-thc-art and research challenges. ACM Computing Surveys, 56(3):1-44, 2023a.
Rui Ye, Mingkai Xu, Jianyu Wang, Chenxin Xu, Siheng Chen, and Yanfeng Wang. Feddisco: Federated learning with discrepaney-aware collaboration. In International Conference Machine Learning, pp. 19879-39902 PMLR, 2023b.
Dong Yin, Yudong Chen, Ramchandran Kannan, and Peter Bartlett. Byzantine-robust distributed learning: Towards optimal statistical rates. In International conference on machine learning, pp. (650-5659, Pmlr, 2018.
Le Yu, Bowen Yu, Haiyang Yu, Fei Huang, and Yongbin Li. Language models are super mario: Absorbing abilities from homologous models as free lunch. In Forty-forst International Conference on Machine Learning, 2024.
Jie Zhang, Chen Chen, Bo Li, Lingjuan Lyu, Shuang Wu, Shouhong Ding, Chunhua Shen, and Chao Wu. Dense: Data-free one-shot federated learning. Advances in Neural Information Processing Systems, 35: 21414 21428, 2022.
Jinghan Zhang, Junteng Liu, Junxian He, et al. Composing parameter-efficient modules with arithmetic operation. Advances in Neural Information Processing Systems, 36:12580-12610 2023.
Yanlin Zhou, George Pu, Xiyao Ma, Xiaolin Li, and Dapeng Wu. Distilled one-shot federated learning, anXiv preprint arXix:2009. 07999, 2020.
Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized sgd and average- direction sam are asymptotically equivalent. In International Conference on Machine Learning, pp. 43005 43036. PMLR, 2023.
20

---

Published in Transactions on Machine Learning Research (05/2025)
A Task Arithmetic Property
In this section, we review a paper that provides theoretical insights into Task Arithmetic (Ottiz-imennz et al. 2023). We examine the relationship between data heterogeneity and the Task Arithmetic property proposed in their work. To facilitate a stronger connection between their framework and our perspective, we adapt their definition of the Task Arithmetic property as follows.
Property A.1. (Task Arithmetic Property from Ortiz-Jimencz et al. 2023 Consider a set of task vectors (Tr)==1 with associaled task data distributions {DiJi=ı- Suppose all the distribulions {D:) have non-intersectinn supports. Let f be a network function. We say network function f satisfies the Task Arithmetic property around do with respect to {Te)E= and (D.) if
f(x,00 0 + Σ Tt) f(z, 0u + TiVe e supp(D,)
and
f(z, 8o + Σ n f(2, & supp())D).
Notice that the Task Arithmetic property is defined through the network function f. while Assumption 4.5 for data heterogeneity is defined through the objective functions L, Recall that L,(0) = (( [e(0: 24, vc)]. The objective function is related to the network function ʃ as follows:
(0) = (( ([ 9+)]
(15)
This connection highlights how the objective function depends on the underlying network function f. Notice that the Task Arithmetic property a property of the network function f, but it does not guarantee the effectiveness of Task Arithmetic. For example, if 0, + T is far from being optimal for objective function Lu then the performance of f(x, 0o + + Tτ) will also be significantly suboptimal. In the subsequent analysis, we will demonstrate that data heterogencity serves a necessary condition for Task Arithmetic achieving optimal performance, assurning 0o + Tr is optimal for Lt-
Proposition A.2. (Necessary Condition) Suppose the Task Arithmetic property holds true and all the objective functions La are H Hssmooth and conver. in 0. Further assume O BT+ is optimal for Ia, ie., 1oo Tς) = 0. Then 0o + X -t= To is optimal for L, and the data heterogeneily at Oo + A Etm Te is zero, i.e.
1
L. + Σ R)) = 0.
T
Proof.
T
T L. + +xΣT) Ti)
1
LO0 + A
Σ
Till
T 1=1
t=1
1
VB ((f 0. + + Σ ye)l by equation 15
V
T
V 10 +T), m.), by Task Arithmetic property
T
(.. + Til
0 by the optimality of 0g Te-
21

---

Published in Transactions on Machine Learning Research (05/2025)
Therefore, 0u + Et= To is also optimal for L. Next,
T
(L.
+
+
Ti)l
T
t=1
VL
E>
Tr) (L, +))l since VL. Ti) -0
1=1
(
2H
T
Tτ) L, + 71) (VI,(00 + Tz), T2 T.,00 + Σ
Tt)
T
2H
T
T
Σ
B(D [e(f(1,, ( 0p + > ut)1 - BI [2) 80 + Tr), vc)
T
1m1
tm1
(
2H
(f(f 9o + Te), s0) (e 8u + v1).
T
0.
Note that incquality (i) follows from the property that for any H-smooth and convex function L, 21
1
V < L(y) L(x) + (V L(r), x ---) y), and equality (ii) follows from Task Arithmetic property. Therefore, the data heterogeneity at Go >E is zero.
The above proposition shows that data heterogeneity being zero is actually necessary condition for the optimal performance of Task Arithmetic. In other words, the paramcters generated by Task Arithmetic have to be a shared oplimum for all L,-
B Merging CLIP
B.1 Additional Experiment Details on Merging CLIP
In this section, we provide details on the hyperparameter searching process for experiments using CLIP. All experiments on CLIP were conducted on a single NVIDIA V100 GPU.
B.1.1 Fine-tuning
For each dataset, we fine-tuned ViT-B-32 using three different learning rates combined with four different numbers of epochs, resulting total of 12 distinct hyperparameter configurations per dataset. The selected numbers of epochs were chosen roughly correspond training for 1000, 2000, 3000 and 4000 iterations, assuming batch size of 128 for each dataset. To determine the optimal hyperparameter configuration, we used the validation accuracy to select the best combination of learning rate and epochs. Table 6 summarizes the fine-tuning hyperparameters and cross-validation details.
Datasets
Learning Rates
Epochs
Best Hyperparameters Configuration
DTD
[1e-4, 1e-5, 1e-6]
{38, 76, 114, 152)
{le-5, 114)
GTSRB
[lc-, le-5, 1c-6]
{6, 11, 17, 22)
[le-5, 6)
SUN397
[1c-4, 1o-5, 1c-6)
{7, 14, 21, 28)
[1c-5, 7)
MNIST
[le-4, le-5, 1e-6]
{3, 5, 8, 10)
{le-5, 8)
SVHN
[1c-4, 1e-5, 1e-6]
[2,4, 6,
[1c-4, 6)
EuroSAT
[1e-4, 1e-5, 1e-6]
{6, 12, 18, 24)
[1e-5, 18)
Cars
[le-,, le-5, lc-6]
[18, 35, 53, 70)
[le-5, 35)
RESISC45
[lc-4, 1c-5, 1c-6]
{8, 15, 23, 30)
[lo-5, 23)
Table 6: Fine-tuning and cross-validation details for CLIP ViT-B-32
22

---

Published in Transactions on Machine Learning Research (05/2025)
B.1.2 Scaling coefficient
To determine the optimal scaling coefficient X, we searched the range [0.05, 0.1,0.15,... 1.95, 2.0], selecting the value that produces the highest average normalized accuracy the validation data sets.
B.1.3 Hyperparameter for FedGMA
To determine the optimal sign-agreement threshold for FedGMA, we scarched the range 00.1, selecting the value that yields the highest average normalized accuracy in validation datasets.
B.1.4 Hyperparameter for CCLIP
To determine the optimal threshold for CCLIP, for each experiment, we searched the range generated by a sequence of five evenly spacod numbers between the minimum task vector norm (including) and maximum task vector norm (excluding) used in the experiment, selecting the value that yields the highest. average normalized accuracy on validation datasstss
B.1.5 Task vector norm
In Table 7 we present the norms of the task vectors used in Section Ö Homogencous fine-tuning task vectors, as provided by [harco al. (2023), are used Section 6.1.2. while heterogeneous fine-tuning task vectors,
developed as part of our work (detailed in Appendix B.1.1)
used in Section 6.1.1
DTD
EuroSAT
GTSRB
SUN397
SVHN
MNIST
Cars
RESISC45
Homogcncous Fine-Tuning
2.47
2.27
2.35
2.91
2.70
2.45
2.80
2.54
Hetcrogencous Fine- TonnTunng
2.77
2.71
1.92
2.04
23.90
3.03
2.80
3.07
Table 7: Norms of task vectors
B.2 Additional Experiments Using Task Vectors with Homogeneous Fine- Tuning
In this section, we present additional experiments on task vectors generated through a homogencous fine- tuning process, provided by Thharco et al. (2023). We apply Median, FedGMA, and CCLIP to these task vectors. Due the homogeneous nature of the fine-tuning process, FedNOVA is not applicable. The hyperparameter scarch for each method follows the same procodure described in Appendix [B.]
Table 8 summarizes the percentage of experiments that show improvement, no change, degradation when compared to Task Arithmetic. Additionally, Figure 3 uses histograms to depict the frequency of experiments within each range of change in average normalized accuracy.
In most cases, these Federated Learning algorithms fail outperform Task Arithmetic. Instead, they tend degrade the performance of the merged models, albeit usually by small margin. These findings further reinforce the key observation discussed in Section 6.2 training heterogeneity is critical factor in practice. Simply regulating the fine- fne-tuning process to eliminate training heterogeneity significantly enhance the performance of Task Arithmetic.
Percentage of Improved Experiments
Percentage of Unchanged Experiments
Percentage of Degraded Experiments
Median FedGMA 9.72%
5.67%
0%
94.33%
34%
56.28%
CCLIP
49.39%
0%
50.61% %
Table Percentage of improved, unchanged, and degraded experiments using different methods compared Task Arithmetic.
23

---

Published in Transactions on Machine Learning Research (05/2025)
Modian
FedGMA
100
140
125
100
Count
75
10
-15.0%-7.5%
Change in Average Noraalized Accuracy
Change Average Nomalizes Accuracy
Change in Avdrage Normalzed Acouracy
Figure 3: Histograms showing the change in average normalized accuracy for three different methods compared to Task Arithmetic by using task vectors with homogeneous fine-tuning. The x-axis shows the difference in average uormalized accuracy relative to Task Arithmellic, with positive values indicaling improvement and negative values indicating degradation. The y-axis represents the number of experiments within each range of change values.
C Merging LLMs
C.1 Additional Experiment Details for Merging LLMs
In this section, we present additional experimental details for merging LLMs. All experiments in this part were conducted on four NVIDIA V100 GPUs. In Table we provide HuggingFace download links for fine-tuned models used in our experiments.
Model
Download Link
Mathematical Reasoning
WisardMath-13B
Htt///WWaaMMa--IH.. 13B-
Code Generation Instruction Following
https: htttps//huugee httpss////hgggggfce..oom co/layoric Snwooric/lune 1aamaa---1b-code--ee paca
Wi23
https: httss///ggggnggaac....../// wwwp://aawaa..e//oa..aaaa.... A /Wizadd M- 133---11..
Table 9: Fine-tuned model download information
In order conduct hyperparameter search, randomly split 5% of GSM8K, MATH, HumanEval and AlpacaEval into validation datascts. To determine the optimal scaling cocfficient X, we scarched the range [0.2, 0.4, 0.6, 0.8, 1.0] for Task Arithmetic, FedGMA and CCLIP.
To determine the optimal sign agreement threshold p for FedGMA, notice that when there are two task vectors merged together, the sign agrecment score for each coordinate is either (opposite sign) 1 (same sign). Therefore, we simply set p be 0.1.
To determine the optimal threshold p for CCLIP, for each experiment, searched the range generated by sequence of five even spaced numbers between the minimun task vector norm (including) and maximum task vector norm (excluding) used in the experiment.
C.2 Additional Experiment Results on Merging LLMs by Median
Table 10 presents the experimental results of applying Median to merge all three task vectors. Compared to the results in Table 5 Median enhances the merged model's mathematical reasoning capabilities by preserving most of its task vector, as its task vector has a norm of middle value. However, this improvement comes at the cost of compromising the modcl's ability for code generation and instruction following. These findings suggest that applying Median to task vectors generated through different fine-tuning methods may be suboptimal, highlighting the need for developing new model merging techniques.
24

---

Published in Transactions on Machine Learning Research (05/2025)
Mathematical Reasoning
Code Generation
Instruction
Tasks
Method
Following
GSM8K
MATH
HumanEval
MBPP
ApacaEval
Instruction
Math
Median
65.73
13.7
10.37
11.6
53.5
Code
Table 10: Performance of Median on merging LLMs
D Generalization Ability of Task Arithmetic
In this section, explore potential reason behind the strong empirical performance of Task Arithmetic observed in several experiments. We conjecture this success is linked to the strong generalization capabilities that Task Arithmetic may inherit, from FedAvg or local SGD.
Research on local SGD has shown that, compared to mini-batch SGD which is the core algorithm used in standard centralized training, local SGD can offer better generalization properties (Gu et al., 2023 Lin et al. 2018 Zhu al., 2023)- Lin et al. (2018) first observed that switching to local SGD after several epochs of mini-batch SGD training enhances model generalization, leading them to propose a post-local SGD approach. In this scheme, local SGD is only employed in the second training phase, after initial mini- batch SGD training. This two-phase strategy mirrors Task Arithmetic, where we start with a mini-batch pre-trained model, switch to local SGD, and ultimately aggregate the updates.
Gu el al. (2023) provided theoretical insights into why local SGD improves generalization. They derived Stochastic Differential Equation (SDE) that models the long-term behavior of local SGD, observing that it induces larger drift term compared to standard SGD, thereby adding regularizing effect. Later on, Zhu et al. (2023) proved that decentralized SGD is asymptotically equivalent to minimizing the loss function of an average-directio sharpness-aware minimization algorithm, which enhances generalization by seeking flatter regions in the loss landscape. This challenges the common belief that centralized training always outperforms decentralized approaches.
The similar phenomenon has been observed in the context of Task Arithmetic. For example, Yu al. (2024) report in Table 1 that Task Arithmetic occasionally surpasses task-specific fine-tuning. In our experiments, particularly when merging LLMs, Task Arithmetic exhibits strong performance. As shown in Table Task Arithmetic achieves the best results when combining three task vectors. This performance is likely attributed to the remarkable generalization ability of local SGD, even in the onc-shot setting of Task Arithmetic.
25


Reviews:
[{"review_id": "oQQbFVmDH6", "reviewer": "Reviewer_kFLT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1. The study effectively links Task Arithmetic to Federated Learning, offering a novel viewpoint on the integration of models.\n2. The examination of data and training diversity under the recognized assumptions of Federated Learning is thorough and perceptive.\n3. The modification of current algorithms for Task Arithmetic, along with thorough experiments, greatly enhances their practical application.\n4. The experiments explore homogeneous and heterogeneous environments alike, thoroughly confirming the theoretical assertions.\n5 The paper is organized logically, with clear descriptions of both theoretical and experimental elements.\n\nWeaknesses\n1. The proposed methods may face challenges when applied to extremely large-scale scenarios with highly diverse tasks.\n2. The paper could benefit from a broader comparison with alternative model merging techniques beyond Task Arithmetic.\n3. Lack an overview picture about the structure of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper establishes a novel connection between Task Arithmetic (a model merging technique) and Federated Learning, specifically one-shot Federated Averaging (FedAvg). The authors leverage theoretical insights from Federated Learning to analyze factors affecting Task Arithmetic, such as data and training heterogeneity, and propose adaptations of Federated Learning algorithms to enhance Task Arithmetic's performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Provide an overview picture about the architecture of the paper.\n2. Compare with more methods in the area of federated learning.\n3. Improve the scale of the experiments and show the results.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZyW1mOWTe4", "reviewer": "Reviewer_9fnG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1. Robust Theoretical Analysis: The paper provides a solid and well-structured theoretical foundation.  \n\n2. Clarity and Presentation: The content is well-presented and easy to follow.  \n\n3. Comprehensive Coverage: The paper includes extensive references and thorough experiments, incorporating large language models (LLMs).  \n\nWeaknesses: \n1. Unclear Problem Definition: The definition of the TA problem lacks clarity and could benefit from further elaboration.  \n\n2. Overclaims: The study focuses primarily on statistical heterogeneity across tasks and variations in training steps, yet it generalizes findings to broader aspects of data and training heterogeneity, which may be overstated.  \n\n3. Limited Novelty and Originality: The paper heavily relies on adapting existing concepts and theoretical results from FL to the TA problem, with a direct application of FL algorithms rather than proposing fundamentally new approaches.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors formulate the Task Arithmetic (TA) problem within the Federated Learning (FL) framework as a one-shot FL problem. They extend convergence analysis from the FL domain to study the convergence properties of Task Arithmetic. The analysis highlights the impact of data and training heterogeneity on Task Arithmetic performance. Additionally, the authors adapt several FL algorithms to the Task Arithmetic setting and conduct empirical evaluations on benchmark datasets to validate their theoretical findings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Weakness 1: The paper would benefit from a clearer introduction to the TA problem, including formal definitions to provide a more precise understanding.  \n\nWeakness 2: The description of heterogeneity should be more precise, distinguishing between data heterogeneity, specifically statistical heterogeneity across tasks, and training heterogeneity, which refers to variations in local training steps.  \n\nWeakness 3: While the paper provides a solid theoretical analysis of the TA problem, it would be more impactful to focus on TA-specific challenges and highlight the contributions to TA directly, rather than emphasizing comparisons with FL.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "NU3LbbHPnJ", "reviewer": "Reviewer_NweM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "After reading this submission, I have the following questions and suggestions: \n\n1. Reorganizing Related Work and giving proper credit\n\nI suggest the authors move the Related Work section to the second section of the paper. This would provide better context early in the discussion. Additionally, the authors should give more credit to [1], which highlights that \"model merging is an extreme case of Federated Learning, where only one round of communication occurs.\"\n\nCurrently, the Related Work section is placed at the end of the paper, which diminishes the acknowledgment of prior contributions. As a result, it is challenging to clearly see the similarities and differences between FL and Task Arithmetic. This distinction is not well-discussed.\n\nMy impression after reading the submission is that the paper primarily applies FL methods to Task Arithmetic. However, the specific challenges and difficulties faced during this process are not clearly articulated.\n\n2. Missing comprehensive baselines and comparisons\n\nThe paper lacks richer baselines and more detailed comparisons with [1]. This omission limits the understanding of the novelty and effectiveness of the proposed approach.\nMoreover, [2], which is a well-known one-shot FL algorithm, is not discussed or compared. Including these discussions would strengthen the paper's contribution and provide more context for readers.\n\n[1] Dataless Knowledge Fusion by Merging Weights of Language Models. ICLR 2023.\n\n[2] DENSE: Data-Free One-Shot Federated Learning. NeurIPS 2022.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Similar to another paper that also examines Task Arithmetic for multi-task learning by framing it as a one-shot Federated Learning problem, this paper takes the approach a step further ----- By leveraging well-established theoretical results from FedAvg, the authors identify two key factors that influence the performance of Task Arithmetic: data heterogeneity and training heterogeneity. Furthermore, they conducted several experiments using standard Federated Learning algorithms to enhance the effectiveness of Task Arithmetic.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Move the Related Work to section 2, add more discussions on some related papers, and give proper credit to these works.\n\nClearly discuss the challenges of applying methods from FL to Task Arithmetic.", "broader_impact_concerns": "no", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper applies task arithmetic, a model merging technique, to tackle heterogeneous challenge in federated setting. The heterogeneous is a major challenge of federated learning, thus a new solution will benefit to tackle this challenge. The proposed technique solution is technique sounds and well-supported by theoretical and experimental analysis. All three reviewers are learning to accept this paper, and there is no major concern raised.
Decision: Accept (poster)