FedFiOT: Fisher-Weighted Optimal Transport for Model Merging in One-Shot and Few-Round Federated Learning

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging in One-Shot and Few-Round Federated Learning

Quan Minh Nguyen1 , Hoang M. Ngo1 , Bhakti Chowkwale2 Yasaman Keshtkarjahromi2 , My T. Thai1

arXiv:2611.06783v1 [cs.LG] 4 Nov 2026

1

University of Florida

Abstract

One-shot and few-round federated learning (FL) reduce communication burden by constructing a global model within one or a few communication rounds, but often achieve lower performance than conventional multiround FL. Without frequent synchronization, client models may drift apart, causing neurons that learn similar features to occupy different positions across clients. At the same time, local data heterogeneity can cause corresponding parameters to differ in their importance to each client’s predictions. Together, these challenges cause direct averaging to combine mismatched neurons and assign equal weight to parameters despite differences in their importance. We propose FedFiOT, a Fisher-weighted optimal transport framework for one-shot and few-round FL that leverages Fisher information to quantify parameter importance and guide neuron alignment and model aggregation. FedFiOT jointly optimizes neuron correspondences and global model parameters through alternating minimization of a shared Fisher-weighted objective. Theoretically, we show that FedFiOT’s optimization procedure monotonically decreases the layer-wise merging objective and establish a stationarity bound that quantifies the effects of client heterogeneity and nonuniform Fisher weighting. Experiments across multiple datasets demonstrate consistent improvements over competing baselines in both one-shot and few-round federated learning settings.

2

Seagate Technology

1

INTRODUCTION

Federated learning (FL) McMahan et al. (2017) enables collaborative model training without sharing private data, but often requires a large number of communication rounds to achieve good accuracy. This multi-round training paradigm has several drawbacks. First, clients must frequently be available to connect to the server to send and receive model updates. Second, clients must repeatedly expend computational resources to update the global model whenever they participate in training. Third, repeated communication rounds can increase vulnerability to backdoor, model poisoning or inference attacks by allowing adversaries to continuously refine their attacks in response to the evolving global model Gu et al. (2022); Xie et al. (2025); Zhu et al. (2025); Zhang et al. (2023). Furthermore, frequent communication rounds can amplify the communication and computational costs of privacy-preserving mechanisms such as secure aggregation Bonawitz et al. (2017) and homomorphic encryption Zhang et al. (2020), which impose additional overhead on clients in each round. These limitations have motivated FL approaches that seek to reduce communication frequency. One-shot federated learning constructs a global model using only a single round Guha et al. (2019); Zhang et al. (2022); Tang et al. (2024), while few-round federated learning allows only a small number of communication rounds Park et al. (2021); Jhunjhunwala et al. (2024). Fewround FL provides a middle ground between one-shot and conventional multi-round FL, where the goal is to approach the performance of fully iterative FL using as few communication rounds as possible. However, with fewer synchronization opportunities, the quality of each aggregation step becomes more important, especially when clients’ data are heterogeneous. A key challenge in one-shot and few-round FL is neuron misalignment. Due to permutation invariance in neural networks, two models can represent the same

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

function while differing only in the ordering of their hidden neurons or channels Ainsworth et al. (2022); Entezari et al. (2021). Consequently, similar features may be represented by different neurons across client models, causing naive parameter averaging to combine mismatched neurons. Neuron-matching methods Liu et al. (2022); Yurochkin et al. (2019); Wang et al. (2020); Singh and Jaggi (2020); Akash et al. (2022) address this issue by identifying correspondences between neurons across client models based on the similarity of their weights or learned representations, and then permuting neurons into a common ordering before parameter aggregation.

correspondences and the global model, using only the local models and Fisher estimates sent by the clients (Sections 3.1 and 3.2). • We formulate neuron alignment as a damped Fisher-weighted optimal transport problem, assigning larger penalties to discrepancies in parameters that are more important to each client’s predictions. We solve the joint alignment and aggregation problem by alternating between computing layer-wise transport plans and updating global model parameters through Fisher-weighted averaging of aligned client parameters (Section 3.3). • Our theoretical analysis shows that the alternating updates monotonically decrease the layer-wise merging objective, with damping making each aggregation step strongly convex so that the change in the merged weights vanishes at an O(S −1/2 ) rate. We further establish a stationarity bound across communication rounds that quantifies the effects of client heterogeneity, stochastic noise, and nonuniform Fisher weighting (Section 4).

A complementary problem is determining how parameters from different models should be combined during model fusion. From a Bayesian perspective, model merging can be viewed as finding parameters that maximize the joint likelihood of the models’ posterior distribution. Fisher merging Matena and Raffel (2022) adopts this perspective by approximating each model’s posterior as a Gaussian whose precision matrix is given by its Fisher information. The merged model is then obtained through a Fisher-weighted average, giving greater influence to parameter values with higher estimated precision. Intuitively, a large Fisher value indicates that changing a parameter can substantially affect the model’s predictions, whereas a small Fisher value suggests that the parameter has relatively little influence on the model’s behavior. These two challenges are closely coupled. Since Fisher information is a local property of a single parameter value, Fisher-based merging may become less reliable when the models being merged are far apart in parameter space Matena and Raffel (2022). Consequently, Fisher weighting is most effective when corresponding neurons are properly aligned across models. At the same time, neuron alignment should place greater emphasis on correctly matching neurons whose associated parameters have high Fisher information values, since incorrectly matching parameters that strongly influence the model’s predictions can be substantially more harmful than misaligning parameters with low Fisher values. Existing approaches typically consider Fisher-based weighting and neuron alignment as separate problems, potentially limiting their effectiveness in one-shot and few-round federated learning. To address this gap, we propose FedFiOT, a Fisherweighted optimal transport framework for one-shot and few-round FL. We summarize our contributions: • We propose FedFiOT and show how neuron alignment and Fisher-weighted aggregation, which prior work treats separately, can be posed as a single optimization over the client-specific neuron

• We evaluate FedFiOT on a range of one-shot and few-round FL tasks and show that it gives a consistent improvement in global model accuracy compared to competing baselines involving knowledge distillation, neuron matching or model fusion (Section 5).

2

RELATED WORKS

2.1

One-Shot and Few-Round FL

Existing one-shot federated learning (FL) methods broadly fall into four categories: (i) knowledge distillation, (ii) neuron matching, (iii) model fusion and (iv) ensemble-based methods. Knowledge distillation methods combine client models into an ensemble and transfer its knowledge to a single global model using either server-side auxiliary data Guha et al. (2019); Li et al. (2020a); Gong et al. (2021) or synthetic data generated by GANs or VAEs Zhang et al. (2022); Heinbaugh et al. (2023). However, auxiliary data may poorly represent client distributions, while sharing client data or generating synthetic samples can raise privacy concerns, e.g. data reconstruction attack Hitaj et al. (2017). Data generation and distillation also increase server computation and require careful hyperparameter tuning, complicating their use in one-shot FL Jhunjhunwala et al. (2024). Neuron matching methods Liu et al. (2022); Yurochkin et al. (2019); Wang et al. (2020); Singh and Jaggi (2020); Akash et al. (2022) exploit neural networks’ permutation invariance to align corresponding neurons across

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

client models before averaging their weights. Another line of work, known as model fusion, aims to combine the capabilities of multiple existing models into a single model Choshen et al. (2022); Ilharco et al. (2022); Jin et al. (2022); Matena and Raffel (2022); Yadav et al. (2023). Although not explicitly developed for one-shot federated learning, some of these methods offer improvement to simple parameter averaging. However, they primarily focus on merging pre-trained models and do not explicitly account for heterogeneity in the data used to train them. FedBEns Talpini et al. (2025), an ensemble-based approach, captures multimodal client posteriors through mixtures of Laplace approximations, but requires clients to train and communicate multiple models and retains an ensemble for prediction. Few-round FL relaxes the single-round communication constraint, allowing clients to refine their models over a small number of communication rounds Jhunjhunwala et al. (2024); Park et al. (2021). 2.2

Permutation Invariance & Neuron Alignment in Federated Learning

Due to permutation invariance, direct parameter averaging can combine mismatched neurons across clients. Neuron alignment methods address this issue by establishing correspondences between neurons before aggregation. PFNM Yurochkin et al. (2019) uses Bayesian nonparametric inference to determine whether a local neuron should match an existing global neuron or introduce a new one, while FedMA Wang et al. (2020) extends this approach to CNNs and LSTMs by alternating layer-wise matching with local retraining. Other methods formulate neural layers as distributions over neurons and use optimal transport (OT) to establish correspondences before merging Singh and Jaggi (2020) or graph matching that captures dependencies between connected neurons Liu et al. (2022). Although these methods mitigate permutation mismatch, alignment alone does not resolve differences in features learned from heterogeneous client data. 2.3

Model Fusion in Federated Learning

Most model fusion approaches target pre-trained or fine-tuned models and offer no theoretical guarantees in FL. A closely related work is Fisher merging Matena and Raffel (2022), which approximates each model’s posterior as a Gaussian whose precision is given by its diagonal Fisher information and casts merging as maximizing the models’ joint posterior likelihood. FedFisher Jhunjhunwala et al. (2024) extends this idea to one-shot FL and provides a theoretical analysis for two-layer over-parameterized ReLU networks. However, since the Fisher information is a local property of a single parameter value, Fisher-based merging may

perform poorly when the models being merged are misaligned and lie far apart in parameter space.

3

PROPOSED ALGORITHM: FEDFIOT

In this section, we first introduce the federated modelmerging setup and formulate neuron alignment and parameter aggregation as a joint optimization problem. We then present FedFiOT, and describe its application to one-shot and few-round FL. 3.1

Federated Model-Merging Setup

We consider an FL system with K clients, where client k holds a private dataset Dk . All clients share an Llayer network architecture. At each communication round, the server broadcasts a global model, which each participating client uses to initialize local training. Client k then returns its updated parameters ek ⪰ 0 Wk ∈ W := Rd and a diagonal Fisher estimate F computed from its local data. We describe a single aggregation round and omit the round index for clarity. The Fisher estimate provides a measure of parameter importance. Specifically, a larger entry indicates greater sensitivity of the client’s predictive distribution to changes in the corresponding parameter Martens (2020); Matena and Raffel (2022). Details of the Fisher estimation procedure are provided in Appendix B. Let mℓ denote the number of neurons in layer ℓ. ℓ For a fully connected layer, let wk,j ∈ Rmℓ−1 denote the incoming weights of neuron j in client k’s ℓ model, and let fek,j denote the corresponding diagonal Fisher block. We omit biases from the notation. ℓ ℓ Client k’s layer is wkℓ , with neuron weights {wk,j }m j=1 . The global parameters at this layer are denoted by ℓ wℓ ∈ W ℓ := Rmℓ ×mℓ−1 , with neuron weights {wiℓ }m i=1 . Neurons that learn similar features may appear at different positions within the same layer across clients, and their parameters may have different Fisher values. The aggregation problem therefore involves two coupled decisions: (1) identifying neuron correspondences and (2) determining the global parameter values. Our goal is to formulate these decisions jointly while accounting for the parameter importance indicated by the Fisher estimates. 3.2

Joint Alignment and Aggregation Objective

We formulate the merge at layer ℓ as a joint optimization over neuron correspondences and global parameters. For this formulation, the correspondence between

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Client models

1

Fisher-weighted optimal transport Cost:

` Weights wk,j ` Fisher fek,j

Client 1

Client 2

Client 3

A B C

B C A



`,[s]  Ck = ij

2

2 `,[s−1] ` wi − wk,j `,λ fe

(4)

k,j

Shared reference Global w`,[s−1]

Aligned clients

`,[s] Neuron Alignment π 1

=

A B C

`,[s]

=

A B C

`,[s]

=

A B C

C A B

Neuron Alignment π 2

A C B

Neuron Alignment π 3

Minimize shared objective (2) `,[s−1] : minπ` J` 1 Fix w 1:K `,[s] 2 Fix π 1:K : minw` J `

Repeat steps 1–2

Fisher-weighted model merging

`,[s] = wi

Matched neuron pair? 1: matched; 0: not matched

P K

−1 P K b`,λ b`,λ b` (7) k=1 fk,i k=1 fk,i w k,i

Alternating optimization Server iterations s = 1, . . . , S

Use w`,[s] as reference for iteration s + 1

Updated global layer w`,[s]

Global neuron i’s weights

m K X̀ X
` min J ` w` , π `1:K = min πk,ij

w` ,π `1:K

` wi` − wk,j

w` ,π `1:K k=1 i,j=1

2 `,λ fek,j

(3)

Parameter importance (damped Fisher)

Client k, neuron j’s weights

s = 1 → 2 → ··· → S

A B C

Fisher merging

Stop early if alignments stabilize

Figure 1: Overview of FedFiOT. At each layer, the server minimizes the joint objective J ℓ by alternating between Fisher-weighted neuron alignment and Fisher-weighted model merging. The merged layer serves as the alignment reference for the next iteration. This repeats for up to S iterations, or until no client’s alignment ℓ changes. Here, ℓ indexes the layer and [s] the iteration; wiℓ and wk,j denote the weights of global neuron i and of neuron j of client k. In each permutation matrix, rows are global neurons, columns are client neurons, and ℓ,λ ℓ orange cells mark matched pairs. w bk,i and fbk,i denote the weights and damped Fisher block of the neuron of client k matched to global neuron i. For illustration, we provide a simple numerical example in Appendix G. neurons in layer ℓ − 1 is held fixed and determines a common ordering for the incoming coordinates. For the first hidden layer, this ordering is given by the shared input features. Let P ℓ denote the set of mℓ ×mℓ permutation matrices. For client k, an alignment πkℓ ∈ P ℓ specifies the neuℓ ron correspondences. Specifically, πk,ij := [πkℓ ]ij = 1 means that client neuron j corresponds to global neuℓ ℓ ron i. We write π1:K = (π1ℓ , . . . , πK ) for the collection of client alignments. The output-layer alignment is restricted to the identity, since all clients share the same class ordering.

(Appendix B). This tempering preserves the ordering of parameter importance while compressing its range. e k and To keep the notation light, we absorb ρ into F ℓ,λ e refer to fk,j simply as the damped Fisher. Writing ∥x∥2A := x⊤ Ax, we define mℓ K X X
2 ℓ ℓ ℓ . (2) J ℓ wℓ , π1:K := πk,ij wiℓ − wk,j feℓ,λ k=1 i,j=1

The resulting layer-merging problem is min ℓ

To account for parameter importance, we use the damped Fisher blocks ℓ,λ ℓ fek,j := fek,j + λI,

λ > 0,

(1)

where λ > 0 ensures positive precision even for coordinates with zero estimated Fisher information. Since diagonal Fisher entries often span several orders of magnitude, a few coordinates can dominate a Fishere k with each weighted average. We therefore compute F empirical Fisher entry raised to a power ρ ∈ (0, 1]

k,j

w ∈W

ℓ

ℓ π1:K ∈(P ℓ )K

ℓ J ℓ wℓ , π1:K .

(3)

The objective penalizes discrepancies between matched neurons, with greater weight assigned to coordinates with larger Fisher values. Thus, the same importance estimates guide both the choice of neuron correspondences and the global parameter values. We propose FedFiOT, a solution for the layer-merging problem which constructs a merged model through alternating alignment and aggregation updates.

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

3.3

Alternating Alignment and Aggregation

We solve (3) by alternating between neuron alignment and parameter aggregation. We index these inner iterations by a bracketed superscript [s]. For a fixed layer ℓ, the global weights wℓ,[0] can be initialized from either the broadcast model or a locally trained client model. Each inner iteration then consists of the following two updates. Fisher-weighted alignment. With the global weights fixed at wℓ,[s−1] , we construct the matching cost 2  ℓ,[s]  ℓ,[s−1] ℓ − wk,j := wi , (4) Ck ℓ,λ ij fek,j

ℓ,[s]

where Ck ∈ Rmℓ ×mℓ collects these costs. The objective separates across clients, so each alignment is updated independently: E D ℓ,[s] ℓ,[s] . (5) πk ∈ arg min π, Ck π∈P ℓ

This assignment problem is equivalent, up to normalization, to optimal transport between two distributions over the client and global neurons. Specifically, ℓ,[s] πk /mℓ is an optimal transport plan. We compute the assignment using the Hungarian algorithm Kuhn (1955), which takes O(m3ℓ ) time. The Fisher-weighted cost makes mismatches in high-importance coordinates more expensive. Fisher-weighted aggregation. Once the alignℓ,[s] ments πk are fixed, we reorder each client’s neurons and their Fisher blocks into the global neuron order. For global neuron i, define ℓ w bk,i := ℓ,λ fbk,i :=

mℓ X

ℓ πk,ij wk,j ,

j=1

ℓ,[s] ℓ,λ πk,ij fek,j .

ℓ,[s]

j=1

mℓ X

(6)

Each row of a permutation matrix contains exactly one nonzero entry, so these expressions select the client neuron matched to global neuron i and its corresponding Fisher block. With these correspondences fixed, setting the gradient of (2) with respect to wiℓ to zero gives

matrix being inverted is positive definite. We repeat alignment and aggregation until the alignments stop changing or S iterations are completed. Merging across layers. The server processes layers in order. Before merging layer ℓ, it reorders working copies of the client weights’ incoming coordinates and the corresponding Fisher entries using the final alignment of layer ℓ − 1. This makes the neuron comparisons consistent with the preceding layer. The input alignment is the identity. At the output layer, the neuron alignment is also fixed to the identity, so only one aggregation update is needed. After all layers are merged, the server compares the full-model merge objective of the aligned model with that of the identityaligned merge, and returns the latter if alignment does not reduce the objective. One-shot and few-round operation. The returned model initializes the next communication round. Clients recompute their Fisher estimates after local training, and the server recomputes the alignments during each merge. Setting T = 1 gives oneshot FL. Larger T repeats the same procedure for fewround FL. For partial participation, all client sums above are restricted to the selected set K(t) . Algorithm 1 summarizes the procedure.

4

THEORETICAL ANALYSIS

We analyze FedFiOT at two levels. First, we establish stability of the alternating updates within each layer. Second, we bound the average squared gradient of the federated training objective in terms of local training, client heterogeneity, and nonuniform Fisher weighting. Proofs and lemmas are deferred to the appendix. 4.1

Stability of FedFiOT Aggregation

Fix a communication round and a layer ℓ, with the incoming-coordinate correspondence held fixed. The client weights and Fisher estimates remain unchanged throughout the inner iterations. We write J ℓ,[s] := ℓ,[s] J ℓ (wℓ,[s] , π1:K ) for the objective after iteration s and define the change in the merged layer by ∆ℓ,[s] := wℓ,[s] − wℓ,[s−1]

2

.

(8)

(7)

To define the initial objective J ℓ,[0] , we use the broadℓ,[0] cast layer wℓ,[0] and identity alignments πk = Imℓ for all clients.

Since the Fisher blocks are diagonal, this update averages each weight coordinate according to its corresponding Fisher values. Damping ensures that the

For fixed alignments, damping makes the layer objective 2Kλ-strongly convex in the global weights. Together with the exact assignment update, this yields the following guarantee.

ℓ,[s]

wi

=

!−1 K K X X ℓ,λ ℓ,λ fb fb w bℓ . k,i

k=1

k,i

k,i

k=1

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Theorem 1 (Stability of the Layer-Wise Merge). Supℓ pose λ > 0 and fek,j ⪰ 0 for every client k and neuron ℓ,⋆ j. Let J denote the minimum value of (3). Under the exact updates (5) and (7), the objective values satisfy J ℓ,[s] ≤ J ℓ,[s−1] for every s ≥ 1. Moreover, after S ≥ 1 inner updates, r J ℓ,[0] − J ℓ,⋆ ℓ,[s] . (9) min ∆ ≤ 1≤s≤S KλS Theorem 1 shows that alignment and aggregation consistently reduce the same layer objective. The available objective decrease limits how long the merged weights can continue to change substantially: among the first S updates, the smallest consecutive change is bounded at an O(S −1/2 ) rate. This guarantee concerns the inner solver for a fixed layer and does not imply global optimality of the complete network merge. 4.2

Stationarity of Federated Training

We now study FedFiOT across communication rounds and restore the round index, omitted in Section 3, as a superscript (t). The analysis follows the training procedure. Specifically, we first bound the disagreement created by local training, then control its effect on the server’s aggregation, and finally establish a stationarity guarantee. Local training and optimization assumptions. We consider full participation, where all K clients participate in every PK round. The global objective is L(W) := K −1 k=1 Lk (W). At round t, each client starts from the broadcast model W(t) and performs E local SGD steps with learning rate η. Client k returns (t) Wk = W(t) − η

E−1 X

(t)

gk,e ,

(10)

e=0

(t)

where gk,e is the stochastic gradient evaluated at its e-th local iterate. Assumption 1 (Federated Optimization). Each Lk is Lf -smooth, and the global objective L is bounded below by L⋆ . Stochastic gradients are conditionally unbiased given the preceding randomness, have conditional variance at most σ 2 , and have norm at most G almost surely. Clients use independent fresh mini-batches at each local step. Moreover, for every W, K

1 X 2 ∥∇Lk (W) − ∇L(W)∥2 ≤ ζ 2 . K

(11)

same model. Smoothness and the gradient bound control the additional drift caused by taking multiple local steps. Dispersion of the returned client models. Before analyzing alignment and aggregation, we quantify how far the client models separate during local training. Define their arithmetic mean and dispersion before alignment as W

(t+1)

K

:=

1 X (t) Wk , K k=1 K

V

(t)

1 X (t+1) 2 (t) Wk − W . := K 2 k=1

For compactness, let   L2f η 2 (E − 1)2 G2 1 σ2 Rloc := ζ 2 + 1 − + . (13) K E 4 Lemma 1 (Dispersion of Local Models). Under Assumption 1, the dispersion of the returned client models satisfies EV (t) ≤ 3η 2 E 2 Rloc . (14) The proof decomposes each local update into the gradient at the broadcast model, the drift along the local trajectory, and stochastic-gradient noise. After subtracting the client average, these three components give the three terms in Rloc . Lemma 1 therefore provides a bound on the disagreement that the server must handle during aggregation. Effect of FedFiOT aggregation. We next compare the model returned by FedFiOT with the arith(t+1) metic mean W of the original client uploads. Let P denote the full-parameter permutations induced by consistent reordering of hidden neurons, their incoming and outgoing weights, and biases when present, with input and output coordinates fixed. We consider architectures for which these permutations preserve each client loss: Lk (ΠW) = Lk (W) for every Π ∈ P. Proposition 1 in the appendix establishes this property for the fully connected setting. (t)

e denote client k’s uploaded diagonal Fisher esLet F k e λ,(t) := F e (t) + λI with λ > 0. timate, and write F k k For full-model alignments Π1:K = (Π1 , . . . , ΠK ) ∈ P K , where Πk is induced by the layer alignments πk1 , . . . , πkL−1 , the full-model counterpart of the layer objective in (2) is

k=1

Here, σ 2 controls stochastic-gradient noise, while ζ 2 measures disagreement between client gradients at the

(12)

J (t) (W, Π1:K ) :=

K X

k=1

(t)

W − Πk Wk

2 λ,(t)

e Πk F k

Π⊤ k

.

(15)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

This objective sums the layer objectives (i.e.,
PL ℓ ℓ ℓ J (t) (W, Π1:K ) = ℓ=1 J w , π1:K ) with the incoming coordinates and Fisher entries reordered consistently according to the preceding layer’s permutation. †,(t) Π1:K

Let denote the permutations associated with the candidate selected by the server, including identity permutations when the identity-aligned candidate is retained. For fixed permutations, the aggregation updates in (7) minimize the full-model objective: †,(t)

W(t+1) = arg min J (t) (W, Π1:K ). W

(16)

The server’s selection rule (Algorithm 1) guarantees (t+1) †,(t)
, I) J (t) W(t+1) , Π1:K ≤ min J (t) (W, I) ≤ J (t) (W W

(17) where I is identity alignment for all clients. To convert this objective comparison into a bound on model dispersion, we control the range of the Fisher entries. Assumption 2 (Fisher Bounds). The undamped e (t) ]cc ≤ Fmax for every Fisher entries satisfy 0 ≤ [F k client k, round t, and coordinate c ∈ {1, . . . , d} of the full parameter vector. After damping, every Fisher entry lies between λ and Fmax + λ. Define κ :=

Fmax + λ , λ

(18)

which controls the conversion from Fisher-weighted cost to Euclidean dispersion. We also quantify how the aggregation weights differ from uniform averaging. For the selected permutations, let
e λ,(t) Π†,(t) ⊤ b λ,(t) := Π†,(t) F F k k k k

be client k’s damped Fisher in the global coordinate ordering. Define b λ,(t) ]cc [F (t) ωk,c := PK k λ,(t) , b ′ ]cc [F ′

(19)

k =1

β (t) := max k,c

k (t) Kωk,c − 1

.

(t)

Here, ωk,c is the aggregation weight assigned to client k at coordinate c, and β (t) = 0 corresponds to uniform averaging of the aligned models. Assumption 2 gives β (t) ≤ Fmax /λ. Together, permutation invariance, the objective comparison, and the Fisher bounds yield the following loss guarantee. Lemma 2 (Loss Control after FedFiOT Aggregation). Under Assumptions 1 and 2, suppose the selected per-

mutations preserve every client loss, then we have: L(W(t+1) ) ≤ L(W

(t+1)

+ β (t)

p

)+

Lf (1 + κ)V (t) 2

κV (t) ∇L(W(t+1) )

(20) 2

.

The key to the proof is to combine (17) with the Fisher bounds to control the aligned models’ dispersion around the merge by κV (t) . Lemma 2 shows that the loss of the FedFiOT aggregate exceeds that of the (t+1) plain average W by at most terms controlled by the client dispersion V (t) and the non-uniformity β (t) of the aligned Fisher weights. Stationarity guarantee. The arithmetic mean (t+1) W satisfies a local-SGD descent inequality. Lemma 2 transfers this inequality to the actual FedFiOT output, while Lemma 1 bounds the aggregation error through Rloc . Together, these results give the following guarantee. Theorem 2 (Stationarity of FedFiOT). Under Assumptions 1 and 2, if β (t) ≤ βmax for every round t and 0 < ηE ≤ 1/(4Lf ), then
T 2 4 L(W(0) ) − L⋆ 1 X (t) E ∇L(W ) ≤ T + 1 t=0 ηET 2 (21) 2 3 8Lf ησ + L2f η 2 (E − 1)2 G2 + 2  K  2 + 6Lf (1 + κ)ηE + 12κβmax Rloc .

Theorem 2 separates the effects of local optimization and server aggregation. The aggregation contribution depends on the disagreement created by local training and the variation in Fisher weights. Its component proportional to ηE decreases with the learning rate, whereas persistent nonuniform weighting can leave a residual term. In particular, for fixed E, bounded κ, and a uniform bound βmax independent of T , by choosing η = Θ(T −1/2 ), we have: T 2 1 X E ∇L(W(t) ) T + 1 t=0 2      1 σ2 −1/2 2 2 . =O T + κβmax ζ + 1 − K E (22) Uniform Fisher weights in the aligned coordinates therefore recover a vanishing O(T −1/2 ) stationarity bound. With persistent nonuniformity, the guarantee permits a residual error governed by client heterogeneity and stochastic noise. Smaller βmax tightens this residual bound, but a small fixed value does not generally make it vanish.

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

5

EXPERIMENTAL RESULTS

5.1

Experimental Setup

We evaluate FedFiOT against state-of-the-art (SOTA) neuron matching, knowledge distillation, model fusion, and Bayesian ensemble baselines on CIFAR-10, SVHN Netzer et al. (2011), and CINIC-10 Darlow et al. (2018) for image recognition in a FL setting. Dataset details are provided in Appendix C.1. Baseline Methods We compare our method with nine baselines: (i) FedAvg McMahan et al. (2017), (ii) PFNM Yurochkin et al. (2019), (iii) OTFusion Singh and Jaggi (2020), (iv) WB Fusion Akash et al. (2022), (v) DENSE Zhang et al. (2022), (vi) FisherMerge Matena and Raffel (2022), (vii) RegMean Jin et al. (2022), (viii) FedFisher Jhunjhunwala et al. (2024), and (ix) FedBEns Talpini et al. (2025). FedAvg serves as the standard parameter-averaging baseline. PFNM and OTFusion are neuron-matching methods that align neurons across local models before aggregation. WB Fusion jointly optimizes neuron alignment and aggregation through a Wasserstein barycenter formulation, but does not incorporate Fisher-based parameter importance or provide explicit convergence analysis. DENSE uses data-free knowledge distillation with GAN-generated data at the server. FisherMerge, RegMean, and FedFisher represent model fusion approaches, with FisherMerge and FedFisher using Fisher information to guide aggregation. FedBEns captures multimodal client posteriors using mixtures of Laplace approximations and performs ensemble prediction. We exclude inherently multi-round methods, e.g. FedProx Li et al. (2020b) and SCAFFOLD Karimireddy et al. (2020) since their performance would be similar to FedAvg for one or few round. Models and Other Details. Following prior works Jhunjhunwala et al. (2024); Talpini et al. (2025), we use the CNN proposed in Wang et al. (2020). To simulate data heterogeneity, we partition each dataset among K clients using Dirichlet sampling with parameter α, where smaller values indicate greater heterogeneity. We use the same local optimization procedure across all algorithms. Further details, including hyperparameter tuning, are provided in Appendix C.2. Communication Efficiency With a diagonal Fisher approximation, each client uploads d Fisher entries alongside its d model parameters, doubling the number of scalars it sends compared to FedAvg. We find that this cost can be reduced to match that of FedAvg using standard quantization techniques without significantly affecting accuracy. For all experiments, we ensure that FedFiOT’s communica-

tion cost matches that of FedAvg. Further practical considerations are discussed in Appendix C.3. 5.2

One-Shot Performance

FedFiOT outperforms baselines: Table 1 reports test accuracy after a single communication round. FedFiOT achieves the highest mean accuracy across all six settings, outperforming the strongest baseline, FedFisher K-FAC, by 3.3–7.6 percentage points. Compared with FedAvg, FedFiOT improves accuracy by 13.0–42.5 percentage points without increasing communication costs. FedFiOT also outperforms neuron-matching methods (OTFusion, PFNM, WB Fusion) and Fisher-weighted aggregation methods (Fisher Merge and FedFisher K-FAC). These results support the effectiveness of jointly incorporating neuron alignment and Fisher-weighted aggregation. We further evaluate FedFiOT under extreme non-i.i.d. data distributions (Appendix D.1), where it continues to outperform all baselines. Comparison against ensemble methods. Figure 3 compares FedFiOT with FedBEns, which trains a mixture of M = 5 models on each client, increasing the computation and communication cost by approximately 5 times. Nevertheless, FedFiOT achieves performance comparable to FedBEns, even outperforming it in multiple settings (Appendix D.2). Different client initialization. We further consider one-shot FL in which clients do not share a common initialization, as in model merging of independently trained networks (Singh and Jaggi, 2020; Ainsworth et al., 2022). This setting arises when clients hold pre-existing models or the server cannot coordinate the initialization, and it removes the partial alignment that a shared initialization provides. In this setting, the gaps between FedFiOT and the baselines widen considerably (Appendix D.3). Ablation Study. We ablate the components of FedFiOT, including alignment, the Fisher-weighted matching cost, and Fisher-weighted aggregation and show that removing any of them reduces accuracy by up to 17–29 points (Appendix F). 5.3

Extension to Few-Round FL

A natural question is whether FedFiOT provides additional utility in the few-round setting. Figure 2 reports four rounds with α = 0.1 and K = 6. We omit PFNM, which modifies the global architecture and increases client computation and communication in every round, and FedBEns, which is designed for the one-shot setting. FedFiOT attains the highest accuracy, reaching 66.7% on CIFAR-10 and 50.8% on

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Table 1: One-shot test accuracy (%) with the number of clients K = 6 fixed and the heterogeneity parameter α varied. Mean ± std over three seeds. FedFiOT uploads quantized weights and diagonal Fishers. WB Fisher FedFisher FedFiOT Dataset α FedAvg OTFusion PFNM RegMean DENSE (K-FAC) (Ours) Fusion Merge 0.2 54.4±2.8 0.1 24.3±2.8

54.8±0.9 31.0±1.5

60.8±1.0 69.5±1.3 61.6±4.3 70.6±0.9 50.7±1.3 34.7±2.3 44.1±2.8 41.1±5.0 48.4±1.1 38.4±3.2

72.2±1.8 61.6±3.6

78.9±1.9 66.8±0.8

0.2 34.0±0.8 CIFAR10 0.1 30.3±1.0

37.0±1.2 32.3±2.8

40.3±1.2 46.0±2.3 39.2±1.7 42.5±1.7 31.9±6.7 29.0±1.1 32.8±1.3 30.9±2.5 33.3±1.8 25.2±2.8

48.7±1.1 38.6±2.1

52.0±0.9 43.3±1.8

0.2 31.9±1.1 0.1 26.6±1.2

34.4±1.3 29.2±0.7

36.4±1.1 39.6±1.4 31.6±1.5 33.2±1.2 26.4±4.6 25.4±0.8 29.8±1.2 25.6±1.3 26.0±0.4 25.6±0.1

40.5±2.9 32.0±0.5

45.6±1.4 39.6±0.3

SVHN

CINIC10

References Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836, 2022. Aditya Kumar Akash, Sixu Li, and Nicolás Garcı́a Trillos. Wasserstein barycenter-based model fusion and linear mode connectivity of neural networks. arXiv preprint arXiv:2210.06671, 2022.

Figure 2: Test accuracy (%) over four rounds of FL with K = 6 clients and α = 0.1.

CINIC-10 after four rounds. Interestingly, the gap between FedFiOT and FedFisher shrinks over the rounds. This suggests that after a few rounds of synchronization, clients start from a shared, well-trained model and their neurons become increasingly aligned, so the permutation mismatch that FedFiOT corrects diminishes (Appendix D.4). FedFiOT’s advantage is therefore most pronounced when communication is limited.

6

CONCLUSION

We propose FedFiOT, a Fisher-weighted optimal transport framework for one-shot and few-round FL that jointly optimizes neuron alignment and aggregation. Its alternating updates monotonically decrease the layer-wise merging objective, and we establish a stationarity bound across communication rounds that quantifies the effects of client heterogeneity, stochastic noise, and nonuniform Fisher weighting. Empirically, FedFiOT consistently outperforms one-shot and fewround FL baselines. Future work will extend FedFiOT to richer curvature approximations e.g. K-FAC and make it compatible with secure aggregation.

Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for privacy-preserving machine learning. In proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, pages 1175–1191, 2017. Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining. arXiv preprint arXiv:2204.03044, 2022. Luke N Darlow, Elliot J Crowley, Antreas Antoniou, and Amos J Storkey. Cinic-10 is not imagenet or cifar-10. arXiv preprint arXiv:1810.03505, 2018. Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:2110.06296, 2021. Xuan Gong, Abhishek Sharma, Srikrishna Karanam, Ziyan Wu, Terrence Chen, David Doermann, and Arun Innanje. Ensemble attention distillation for privacy-preserving federated learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 15076–15086, October 2021. Yuhao Gu, Yuebin Bai, and Shubin Xu. Cs-mia: Membership inference attack based on prediction confidence series in federated learning. Journal of Information Security and Applications, 67:103201, 2022. Neel Guha, Ameet Talwalkar, and Virginia Smith.

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

One-shot federated learning. arXiv:1902.11175, 2019.

arXiv preprint

Clare Elizabeth Heinbaugh, Emilio Luz-Ricca, and Huajie Shao. Data-free one-shot federated learning under very high statistical heterogeneity. In The Eleventh International Conference on Learning Representations, 2023. Briland Hitaj, Giuseppe Ateniese, and Fernando Perez-Cruz. Deep models under the gan: Information leakage from collaborative deep learning. In Proceedings of the 2017 ACM SIGSAC conference on computer and communications security, pages 603–618, 2017. Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. arXiv preprint arXiv:2212.04089, 2022. Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedfisher: Leveraging fisher information for one-shot federated learning. In International Conference on Artificial Intelligence and Statistics, pages 1612–1620. PMLR, 2024. Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietro, and Pengxiang Cheng. Dataless knowledge fusion by merging weights of language models. arXiv preprint arXiv:2212.09849, 2022.

Michael S Matena and Colin Raffel. Merging models with fisher-weighted averaging. Advances in Neural Information Processing Systems, 35:17703– 17716, 2022. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273–1282. Pmlr, 2017. Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. URL http://ufldl.stanford.edu/ housenumbers/nips2011_housenumbers.pdf. Younghyun Park, Dong-Jun Han, Do-Yeon Kim, Jun Seo, and Jaekyun Moon. Few-round learning for federated learning. Advances in Neural Information Processing Systems, 34:28612–28622, 2021. Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045–22055, 2020. Jacopo Talpini, Marco Savi, and Giovanni Neglia. Fedbens: One-shot federated learning based on bayesian ensemble. arXiv preprint arXiv:2503.15367, 2025.

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 on machine learning, pages 5132– 5143. PMLR, 2020.

Zhenheng Tang, Yonggang Zhang, Peijie Dong, Yiuming Cheung, Amelie C Zhou, Bo Han, and Xiaowen Chu. Fusefl: One-shot federated learning through the lens of causality with progressive model fusion. Advances in neural information processing systems, 37:28393–28429, 2024.

Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2 (1-2):83–97, 1955.

Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.06440, 2020.

Qinbin Li, Bingsheng He, and Dawn Song. Practical one-shot federated learning for cross-silo setting. arXiv preprint arXiv:2010.01017, 2020a. Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429– 450, 2020b. Chang Liu, Chenfei Lou, Runzhong Wang, Alan Yuhan Xi, Li Shen, and Junchi Yan. Deep neural network fusion via graph matching with applications to model ensemble and federated learning. In International Conference on Machine Learning, pages 13857–13869. PMLR, 2022. James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1–76, 2020.

Yueqi Xie, Minghong Fang, and Neil Zhenqiang Gong. Model poisoning attacks to federated learning via multi-round consistency. In 2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 15454–15463. IEEE, 2025. 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:7093– 7115, 2023. Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International conference on machine learning, pages 7252–7261. PMLR, 2019.

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Chengliang Zhang, Suyi Li, Junzhe Xia, Wei Wang, Feng Yan, and Yang Liu. {BatchCrypt}: Efficient homomorphic encryption for {Cross-Silo} federated learning. In 2020 USENIX annual technical conference (USENIX ATC 20), pages 493–506, 2020. Hangfan Zhang, Jinyuan Jia, Jinghui Chen, Lu Lin, and Dinghao Wu. A3fl: Adversarially adaptive backdoor attacks to federated learning. Advances in neural information processing systems, 36:61213– 61233, 2023. 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. Gongxi Zhu, Donghao Li, Hanlin Gu, Yuan Yao, Lixin Fan, and Yuxing Han. Fedmia: An effective membership inference attack exploiting” all for one” principle in federated learning. In 2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 20643–20653. IEEE, 2025.

Supplementary Materials

A

PROOFS

ℓ Theorem 1 (Stability of the Layer-Wise Merge). Suppose λ > 0 and fek,j ⪰ 0 for every client k and neuron ℓ,⋆ j. Let J denote the minimum value of (3). Under the exact updates (5) and (7), the objective values satisfy J ℓ,[s] ≤ J ℓ,[s−1] for every s ≥ 1. Moreover, after S ≥ 1 inner updates, r J ℓ,[0] − J ℓ,⋆ ℓ,[s] min ∆ ≤ . (23) 1≤s≤S KλS

Proof. For a fixed merged model wℓ,[s−1] , the OT alignment step minimizes the objective over the client permutations. Therefore,   ℓ,[s] J ℓ wℓ,[s−1] , π1:K ≤ J ℓ,[s−1] . (24) ℓ,[s]

Next, fix π1:K . The merging subproblem is 2Kλ-strongly convex because its Hessian with respect to each wiℓ is P ℓ,λ 2 k fbk,i ⪰ 2KλI, and wℓ,[s] is its minimizer. Hence,     ℓ,[s] ℓ,[s] J ℓ wℓ,[s−1] , π1:K − J ℓ wℓ,[s] , π1:K ≥ Kλ wℓ,[s] − wℓ,[s−1]

Combining (24) and (25) gives

2 2

.

 2 J ℓ,[s−1] − J ℓ,[s] ≥ Kλ ∆ℓ,[s] .

(25)

(26)

Thus, the objective is monotonically non-increasing. Summing (26) over s = 1, . . . , S gives Kλ

S  X s=1

∆ℓ,[s]

2

Since S min

1≤s≤S

≤ J ℓ,[0] − J ℓ,[S] ≤ J ℓ,[0] − J ℓ,⋆ .

(27)



(28)

∆ℓ,[s]

2

≤

Eq. (23) follows.

S  X

∆ℓ,[s]

s=1

2

,

Lemma 1 (Dispersion of Local Models). Under Assumption 1, the dispersion of the returned client models satisfies EV (t) ≤ 3η 2 E 2 Rloc . (29) (t)

Proof. Let Wk,e denote the local model of client k immediately before local step e. Thus, (t)

Wk,0 = W(t) , (t)

(t)

so that Wk = Wk,E .

(t)

(t)

(t)

Wk,e+1 = Wk,e − ηgk,e ,

(30)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

We decompose the local stochastic gradients into three components. Define   (t) (t) (t) ξk,e := gk,e − ∇Lk Wk,e ,

(31)

dk :=

(32)

(t)

E−1   i 1 Xh (t) ∇Lk Wk,e − ∇Lk (W(t) ) , E e=0 E−1

(t)

zk :=

1 X (t) ξ . E e=0 k,e

(33)

(t)

(t)

Here, dk measures the drift caused by evaluating gradients along the local trajectory, while zk is the average stochastic-gradient error. (t)

(t)

Using the local SGD update, the local change ∆k := Wk − W(t) satisfies (t)

∆k = −η

E−1 X

(t)

gk,e

e=0

h i (t) (t) = −ηE ∇Lk (W(t) ) + dk + zk .

(34)

(t)

We first control the local drift. Since ∥gk,e ∥2 ≤ G almost surely, (t)

Wk,e − W(t)

2

= −η ≤η

e−1 X

(t)

gk,h

h=0

e−1 X

2

(t)

gk,h

h=0

2

≤ ηeG.

(35)

By Lf -smoothness of Lk ,   1 X (t) ∇Lk Wk,e − ∇Lk (W(t) ) E e=0 2 E−1

(t)

∥dk ∥2 ≤

E−1

Lf ηG X e ≤ E e=0 =

Lf η(E − 1)G . 2

(36) (t)

We next control the stochastic-gradient error. Let F (t) contain all randomness preceding round t, and let Fe contain F (t) and all randomness preceding local step e. By conditional unbiasedness, h i (t) E ξk,e | Fe(t) = 0, (37) and

E



(t)

ξk,e

2 2

 | Fe(t) ≤ σ 2 .

(38)

Errors from different local steps have zero cross moments because the earlier error is measurable with respect to the filtration preceding the later step. Fresh independent mini-batches across clients also give zero cross moments between different clients. Therefore, (t)

E(t) zk = 0, where E(t) [·] := E[· | F (t) ].

(t)

E(t) ∥zk ∥22 ≤

σ2 , E

(39)

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Let (t)

d

K

K

1 X (t) dk , K

:=

z(t) :=

k=1

1 X (t) zk . K

(40)

k=1

Using independence across clients, the cross moments between different clients vanish, so 2

E(t) z(t)

2

K

=

1 X (t) (t) 2 σ2 E ∥zk ∥2 ≤ . 2 K KE

(41)

k=1

The centering identity K

K

k=1

k=1

1 X 1 X ∥ak − a∥22 = ∥ak ∥22 − ∥a∥22 K K

(42)

  1 σ2 1− . K E

(43)

then gives K

2 1 X (t) (t) zk − z(t) ≤ E K 2 k=1

Similarly, by (36), K

K

1 X (t) 2 1 X (t) (t) 2 dk − d ∥dk ∥2 ≤ K K 2 k=1

k=1

L2f η 2 (E − 1)2 G2 . ≤ 4

(44)

Averaging (34) over clients and subtracting the result from the update of client k gives (t)

Wk − W

(t+1)

h = −ηE ∇Lk (W(t) ) − ∇L(W(t) ) (t)

(t)

+ dk − d

i (t) + zk − z(t) .

Applying

∥a + b + c∥22 ≤ 3 ∥a∥22 + ∥b∥22 + ∥c∥22 ,

(45)

(46)

averaging over clients, and taking conditional expectation gives (t)

E V

(t)

2

≤ 3η E

2

"

K

2 1 X ∇Lk (W(t) ) − ∇L(W(t) ) K 2 k=1

K

+

1 X (t) (t) 2 dk − d K 2 k=1

# K 2 1 X (t) (t) (t) + E zk − z . K 2

(47)

k=1

Using Assumption 1, Eq. (43), and Eq. (44), " (t)

E V

(t)

2

≤ 3η E

Taking total expectation proves Eq. (29).

2



1 ζ + 1− K 2



# L2f η 2 (E − 1)2 G2 σ2 + . E 4

(48)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Proposition 1 (Permutation invariance). Let P denote the set of coordinate permutations of W induced by tuples of permutation matrices (π 1 , . . . , π L−1 ), π ℓ ∈ P ℓ , through wℓ 7→ π ℓ wℓ (π ℓ−1 )⊤ , bℓ 7→ π ℓ bℓ , with π 0 = π L = I. If every activation commutes with coordinate permutation, then every Π ∈ P is orthogonal and Lk (ΠW) = Lk (W) for every client k and every W. Proof. Each Π ∈ P only reorders the coordinates of W, so it is a permutation matrix and hence orthogonal. Let h0 = x, let hℓ = ϕ(wℓ hℓ−1 + bℓ ) denote the activations of the original network, and let h′ℓ denote those of the permuted network. If h′ℓ−1 = π ℓ−1 hℓ−1 , then

h′ℓ = ϕ π ℓ wℓ (π ℓ−1 )⊤ π ℓ−1 hℓ−1 + π ℓ bℓ = ϕ π ℓ (wℓ hℓ−1 + bℓ ) = π ℓ hℓ ,

where the last step uses that ϕ commutes with coordinate permutations. By induction from π 0 = I, and since π L = I, both networks produce the same output on every input. Hence Lk (ΠW) = Lk (W).

Proposition 1 ensures that FedFiOT can reorder neurons without changing the client objectives. This lets the analysis compare aligned client models with their original counterparts while preserving both their losses and Euclidean norms. Lemma 2 (Loss Control after FedFiOT Aggregation). Under Assumptions 1 and 2, suppose the selected permutations preserve every client loss, then we have: L(W(t+1) ) ≤ L(W

(t+1)

+ β (t)

)+

Lf (1 + κ)V (t) 2

p κV (t) ∇L(W(t+1) )

(49) 2

. †,(t)

Proof. Fix a realization of local training, Fisher estimation, and the selected permutations Π1:K . For compactness, define †,(t) (t) X := W(t+1) , Xk := Πk Wk , (50) and let

K

X :=

1 X Xk K

(51)

k=1

be the arithmetic mean of the aligned client models.

We first control how far the aligned client models lie from the FedFiOT aggregate. By (17),   †,(t) J (t) W(t+1) , Π1:K ≤ min J (t) (W, I) W  (t+1)  ≤ J (t) W ,I . Since

we have

e λ,(t) ⪯ (Fmax + λ)I, F k

(53)

K  (t+1)  X (t+1) 2 (t) J (t) W , I ≤ (Fmax + λ) Wk − W 2

k=1

= K(Fmax + λ)V

(52)

(t)

.

(54)

On the other hand, damping and orthogonality of the permutation matrices give b λ,(t) ⪰ λI. F k

(55)

Therefore, defining the aligned dispersion around the FedFiOT aggregate by K

(t)

Vali :=

1 X ∥Xk − X∥22 , K k=1

(56)

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

we obtain   †,(t) J (t) W(t+1) , Π1:K

(t)

Vali ≤

Kλ Fmax + λ (t) ≤ V = κV (t) . λ

(57)

We next control the difference between the Fisher-weighted aggregate X and the arithmetic aligned mean X. Let   (t) (t) (t) Ak := diag ωk,1 , . . . , ωk,d . (58) The closed-form FedFiOT update gives

X=

K X

K X

(t)

Ak Xk ,

k=1

(t)

Ak = I.

(59)

k=1

From the definition of β (t) , (t)

Ak − Since

K  X

k=1

we may write

β (t) 1 ≤ I . K op K

(60)



(61)

(t)

Ak −

= 0,

 1 I Xk K k=1   K X 1 (t) = Ak − I (Xk − X). K

X−X=

K  X

1 I K

(t)

Ak −

(62)

k=1

Hence, K

β (t) X ∥Xk − X∥2 ∥X − X∥2 ≤ K k=1

K

≤β

(t)

1 X ∥Xk − X∥22 K k=1

q (t) (t) =β Vali .

!1/2

We now compare the losses of the aligned and unaligned arithmetic means. By Lf -smoothness,    (t+1)  (t) L Wk ≤ L W D  (t+1)  E (t+1) (t) , Wk − W + ∇L W +

Lf (t+1) 2 (t) Wk − W . 2 2

(63)

(64)

Averaging over clients cancels the linear terms, giving  (t+1)  L 1 X  (t)  f L Wk ≤ L W + V (t) . K 2 K

k=1

(65)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Similarly, the lower smoothness inequality around X gives L(Xk ) ≥ L(X) + ∇L(X), Xk − X Lf − ∥Xk − X∥22 . 2

(66)

Again averaging cancels the linear terms, and therefore L(X) ≤ By Proposition 1,

K

K

k=1

k=1

Lf X 1 X L(Xk ) + ∥Xk − X∥22 . K 2K

    †,(t) (t) (t) L(Xk ) = L Πk Wk = L Wk .

(67)

(68)

Combining this with Eqs. (65) and (67) gives

 (t+1)  L f V (t) + L(X) ≤ L W 2 K Lf X ∥Xk − X∥22 . + 2K

(69)

k=1

Finally, apply the lower smoothness inequality at X with X and rearrange: L(X) ≤ L(X) + ∇L(X), X − X +

Lf ∥X − X∥22 . 2

(70)

The variance decomposition K

(t) Vali =

1 X ∥Xk − X∥22 + ∥X − X∥22 K

(71)

k=1

therefore yields

  (t+1)  L  f (t) + L(X) ≤ L W V (t) + Vali 2 + ∇L(X), X − X .

(72)

Using Cauchy–Schwarz and Eq. (63),

(t)

 (t+1)  L   f (t) L(X) ≤ L W + V (t) + Vali 2 q (t) (t) Vali ∥∇L(X)∥2 . +β

(73)

Finally, Vali ≤ κV (t) from (57) and X = W(t+1) , so

L(W(t+1) ) ≤ L(W

(t+1)

+ β (t) This proves Eq. (49).

)+

Lf (1 + κ)V (t) 2

p κV (t) ∇L(W(t+1) )

2

.

(74)

Theorem 2 (Stationarity of FedFiOT). Under Assumptions 1 and 2, if β (t) ≤ βmax for every round t and 0 < ηE ≤ 1/(4Lf ), then
T 2 4 L(W(0) ) − L⋆ 1 X (t) E ∇L(W ) ≤ T + 1 t=0 ηET 2 (75) 8Lf ησ 2 3 + + L2f η 2 (E − 1)2 G2 2  K  2 + 6Lf (1 + κ)ηE + 12κβmax Rloc .

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Proof. For compactness, define γ := ηE,

D :=

Lf η(E − 1)G . 2

(76)

Let F (t) contain all randomness preceding communication round t, and write E(t) [·] := E[· | F (t) ].

(77)

We first establish descent for the arithmetic average of the locally trained models. Averaging the decomposition in (34) over clients gives i h W

(t+1)

(t)

= W(t) − γ ∇L(W(t) ) + d

+ z(t) .

(78)

From the proof of Lemma 1,

(t)

∥d ∥2 ≤ D, By Lf -smoothness of L,

σ2 . KE

(79)

2 Lf γ 2 (t) ∇L(W(t) ) + d + z(t) . 2 2

(80)

E(t) z(t) = 0,

E(t) ∥z(t) ∥22 ≤

 (t+1)  ≤ L(W(t) ) − γ∥∇L(W(t) )∥22 L W D E (t) − γ ∇L(W(t) ), d + z(t) +

Since ∇L(W(t) ) is F (t) -measurable and E(t) z(t) = 0, D E E(t) ∇L(W(t) ), z(t) = 0.

(81)

For the local-drift term, Young’s inequality gives D E (t) − ∇L(W(t) ), d ≤ D∥∇L(W(t) )∥2 ≤

1 ∥∇L(W(t) )∥22 + D2 . 4

(82)

Moreover, (t)

∇L(W(t) ) + d

+ z(t)

2

2 (t) 2 ≤ 2∥∇L(W )∥2 + 2∥d + z(t) ∥22 ≤ 2∥∇L(W(t) )∥22 + 4D2 + 4∥z(t) ∥22 . (t)

Substituting Eqs. (82)–(83) into (80) and taking conditional expectation yields    (t+1)  3 (t) (t) E L W ≤ L(W ) − γ − Lf γ ∥∇L(W(t) )∥22 4 2Lf γ 2 σ 2 + γ(1 + 2Lf γ)D2 + . KE Since Lf γ = Lf ηE ≤ we have

3 1 − Lf γ ≥ , 4 2

1 , 4

1 + 2Lf γ ≤

(83)

(84)

(85) 3 . 2

(86)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Thus,

 (t+1)  γ ≤ L(W(t) ) − ∥∇L(W(t) )∥22 E(t) L W 2 3γ 2 2Lf γ 2 σ 2 + D + . 2 KE

We next transfer this descent inequality to the actual FedFiOT model. By Lemma 2,  (t+1)  L f L(W(t+1) ) ≤ L W + (1 + κ)V (t) 2 p + β (t) κV (t) ∥∇L(W(t+1) )∥2 . Since β (t) ≤ βmax , Young’s inequality gives p β (t) κV (t) ∥∇L(W(t+1) )∥2 ≤

Therefore,

κβ 2 γ ∥∇L(W(t+1) )∥22 + max V (t) . 4 γ

γ L(W(t+1) ) − ∥∇L(W(t+1) )∥22 4   (t+1)   L κβ 2 f ≤L W (1 + κ) + max V (t) . + 2 γ

(87)

(88)

(89)

(90)

To telescope this inequality, define the potential Φ(t) := L(W(t) ) −

γ ∥∇L(W(t) )∥22 . 4

(91)

Combining Eqs. (87) and (90) gives γ ∥∇L(W(t) )∥22 4 3γ 2 2Lf γ 2 σ 2 + D + KE 2  Lf κβ 2 + (1 + κ) + max E(t) V (t) . 2 γ

E(t) Φ(t+1) ≤ Φ(t) −

It remains to verify that Φ(t) is bounded below. By smoothness,   1 ⋆ L ≤L W− ∇L(W) Lf 1 ≤ L(W) − ∥∇L(W)∥22 . 2Lf

(92)

(93)

Hence, ∥∇L(W)∥22 ≤ 2Lf (L(W) − L⋆ ) .

Consequently,

γ Φ(t) − L⋆ = L(W(t) ) − L⋆ − ∥∇L(W(t) )∥22 4    Lf γ  ≥ 1− L(W(t) ) − L⋆ ≥ 0. 2

(94)

(95)

Thus Φ(t) ≥ L⋆ , while Φ(0) ≤ L(W(0) ). Moreover, by (94) and Lf γ ≤ 1/4, Φ(T ) − L⋆ ≥

 1 γ γ − ∥∇L(W(T ) )∥22 ≥ ∥∇L(W(T ) )∥22 . 2Lf 4 4

(96)

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Taking full expectation in (92), summing over t = 0, . . . , T − 1, telescoping, and lower-bounding EΦ(T ) by (96), we have: T

γX E∥∇L(W(t) )∥22 ≤ L(W(0) ) − L⋆ 4 t=0   3γ 2 2Lf γ 2 σ 2 +T D + 2 KE   TX −1 Lf κβ 2 + (1 + κ) + max EV (t) . 2 γ t=0

(97)

Dividing by γ(T + 1)/4 and using T /(T + 1) ≤ 1 yields

T 4 L(W(0) ) − L⋆ 1 X 8Lf γσ 2 (t) 2 E∥∇L(W )∥2 ≤ + 6D2 + T + 1 t=0 γT KE  TX  −1 2 1 2Lf (1 + κ) 4κβmax + EV (t) . + γ γ2 T t=0

(98)

By Lemma 1, T −1

1 X EV (t) ≤ 3γ 2 Rloc . T t=0

(99)

Substituting this into (98) gives

Finally,

T 4 L(W(0) ) − L⋆ 1 X 8Lf γσ 2 (t) 2 E∥∇L(W )∥2 ≤ + 6D2 + T + 1 t=0 γT KE   2 + 6Lf (1 + κ)γ + 12κβmax Rloc . D2 =

γ = ηE, so 6D2 = and

L2f η 2 (E − 1)2 G2 , 4

3 2 2 L η (E − 1)2 G2 2 f

8Lf γσ 2 8Lf ησ 2 = . KE K

(100)

(101) (102)

(103)

Substituting these expressions into Eq. (100) gives exactly Eq. (75).

B

FISHER INFORMATION MATRIX AND FISHER MERGING

The Fisher information matrix measures the sensitivity of a model’s predictive distribution to changes in its parameters Jhunjhunwala et al. (2024); Matena and Raffel (2022). For client k, let pWk (y | x) denote the predictive distribution of its locally trained model. Its Fisher information matrix is Fk =

Nk   1 X Ey∼pWk (·|xk,n ) sk,n (y)sk,n (y)⊤ , Nk n=1

(104)

where Nk is the number of local examples used for estimation, xk,n is the n-th input, and sk,n (y) = ∇W log pW (y | xk,n )|W=Wk

(105)

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

is the gradient of the log predictive probability evaluated at Wk . Larger diagonal Fisher entries indicate greater sensitivity to changes in the corresponding parameters. To avoid storing the full d × d matrix, we use a diagonal empirical Fisher estimate computed from observed labels and tempered by an exponent ρ ∈ (0, 1]: !ρ Nk X 1 2 e k ]cc = [F [sk,n (yk,n )]c , c = 1, . . . , d, (106) Nk n=1 where yk,n is the observed label and c indexes a parameter. This estimate uses squared per-example gradients in place of the expectation over model-generated labels and requires only O(d) storage. Setting ρ = 1 recovers the standard diagonal empirical Fisher. Smaller ρ compresses the range of the entries (Section 3.2).

C

ADDITIONAL EXPERIMENTAL DETAILS

C.1

Datasets

We evaluate on three 10-class image-classification benchmarks. CIFAR-10 contains natural images of objects and animals. SVHN Netzer et al. (2011) contains digits cropped from street-view house numbers. CINIC-10 Darlow et al. (2018) extends CIFAR-10 with downsampled ImageNet images of the same ten classes, which makes it more diverse and harder. We use the standard training and test splits. All images are 32 × 32 RGB and are normalized per channel using standard dataset-specific mean and standard deviation values. Client partitioning. We distribute each training set among K clients using a Dirichlet label split. Each client k draws a class-preference P vector qk ∼ Dirichlet(α110 ), and the examples of each class c are divided among clients in proportion to qk,c / k′ qk′ ,c . This produces heterogeneous label distributions and unequal client dataset sizes, where smaller α concentrates each client on fewer classes. The partition is generated with a fixed seed and shared by all methods, so every method sees identical client data. Test accuracy is reported on the full, class-balanced test set. Server validation set. We reserve 500 examples as a server-side validation set. These examples are removed from the clients’ data and are therefore excluded from both local training and Fisher estimation. The same validation set is used for hyperparameter tuning for all methods (Appendix C.2). C.2

Training Details and Hyperparameter Tuning

Local training. All clients start from a shared Xavier-uniform initialization and train with SGD (momentum 0.9, weight decay 2 × 10−3 , batch size 64) and an initial learning rate of 0.02 with cosine annealing over the local epochs. We apply a random 32 × 32 crop with 4-pixel padding and a random horizontal flip to each minibatch. FedBEns trains its own mixture components as described below. Model selection. For every method, setting, and seed, we sweep the number of local epochs over {25, 50, . . . , 200} and select the value with the highest accuracy on the 500-example server validation set (Appendix C.1). This is to ensure that every method is compared at its best local training length, since longer local training improves each client model but also increases client drift, and methods differ in how well they tolerate this drift. Method-specific hyperparameters are selected on the same validation set. We report the mean and sample standard deviation over three training seeds on a fixed data partition. The seeds are shared by all methods within a setting. FedFiOT. Each client computes the diagonal empirical Fisher by averaging squared per-example gradients of the observed-label log-likelihood over its entire local dataset (Appendix B). We sweep the tempering exponent over ρ ∈ {0.2, 0.3, 0.4, 0.5, 0.6, 0.7} and the damping over λ ∈ {10−3 , 10−4 , 10−5 }, and select the pair with the highest server validation accuracy. All candidates are computed from the same uploads, so the sweep requires no additional communication. In the one-shot setting, the alignment reference wℓ,[0] is the trained model of the client with the most data, since the broadcast model is the random initialization and its neurons encode no learned features to match against. In later rounds, the sweep additionally selects the reference from this model

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

and the broadcast model. Each hidden-layer assignment is solved exactly with the Hungarian algorithm, with at most S = 8 alternating updates per layer and early stopping once all assignments are unchanged. Clients upload weights and Fisher entries quantized layer-wise to 16 bits. Baselines. • FedFisher (K-FAC) has each client upload a K-FAC approximation of its Fisher, compressed by low-rank SVD with 8-bit factors, together with its weights quantized to 16 bits. Starting from the size-weighted average of the client models, the server minimizes the Fisher-weighted distance to the client models using 2,000 steps of an adaptive momentum optimizer with learning rate 0.01. Every 100 steps, it evaluates the current model on the validation set and keeps the checkpoint with the highest validation accuracy. • RegMean has each client upload the Gram matrix of the inputs to each fully connected layer, compressed in the same way as the FedFisher factors. The server merges each fully connected layer in closed form by solving the least-squares problem that best reproduces all clients’ layer outputs, and averages the convolutional layers by client data size. The off-diagonal entries of the Gram matrices are scaled by α ∈ {0.1, 0.9}, selected on validation data. • DENSE distills the ensemble of client models into a single model, starting from the FedAvg model. Over 20 stages, a generator is trained for 30 Adam steps to synthesize a batch of 256 images on which the ensemble is confident but the global model disagrees with it. These images are added to a replay pool, and the global model is trained to match the ensemble’s predictions on the pool (KL divergence, temperature 20, SGD with learning rate 0.01). After each stage, the model is evaluated on the validation set, and the stage with the highest validation accuracy is kept. • FedBEns trains five models per client, each from one of five random initializations shared across clients, and fits a Kronecker-factored Laplace approximation (temperature 0.1) around each of them. Each client’s posterior is therefore a mixture of five Gaussians. The server multiplies the client posteriors and searches for modes of the product: for each of the five initializations, it starts from the coordinate-wise median of the client models trained from that initialization and maximizes the global posterior with Adam (learning rate 0.001) for 300 steps. Every 30 steps, the current model is evaluated on the validation set, and the checkpoint with the highest validation accuracy is kept. The five resulting models form the final ensemble, whose predicted class probabilities are averaged. C.3

Practical Implementation Considerations

Communication efficiency. We assume that a scalar is represented with 32 bits by default, so a FedAvg client uploads its local model Wk at a cost of 32d bits. A FedFiOT client additionally uploads its diagonal e k , which also has exactly d entries, doubling the uncompressed upload to 64d bits. To match the cost of Fisher F FedAvg, we apply the same uniform quantization as FedFisher Jhunjhunwala et al. (2024). For a vector x (one e k ) and lq = 215 − 1 levels, each entry is quantized as layer of Wk or F [Q(x)]i = ∥x∥∞

round(lq xi /∥x∥∞ ) . lq

(107)

Each quantized entry takes one of 2lq + 1 < 216 values and is therefore sent with 16 bits (15 for the magnitude and 1 for the sign), plus 32 bits per layer for ∥x∥∞ . Quantizing both the weights and the Fisher layer-wise gives a total upload of 2 · 16d + 2 · 32L = 32d + 64L bits, where L is the number of layers. For our CNN, L = 5 and d ≈ 0.94M, so the 320-bit overhead relative to FedAvg is negligible. The server applies the tempering exponent ρ and damping λ after receiving the uploads, so the hyperparameter sweep of Appendix C.2 requires no additional communication. Table 2 compares FedFiOT with full-precision and quantized uploads. Quantization changes the mean accuracy by at most 1.3 percentage points in all settings except SVHN with α = 0.1, where it improves accuracy upon the uncompressed version which can be attributed to an additional regularization effect of compression, as also observed for FedFisher Jhunjhunwala et al. (2024).

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Table 2: One-shot test accuracy (%) of FedFiOT with full-precision (32-bit) and quantized (16-bit) uploads, K = 6 clients. Mean ± std over three seeds. Dataset α 32-bit 16-bit Difference SVHN

CIFAR-10 CINIC-10

0.2 0.1

78.2±2.0 62.1±2.1

78.9±1.9 66.8±0.8

+0.7 +4.7

0.2 0.1

51.1±0.9 44.6±1.3

52.0±0.9 43.3±1.8

0.2 0.1

45.5±2.5 38.6±0.4

45.6±1.4 39.6±0.3

+0.9 −1.3 +0.1 +1.0

Computation efficiency. To compute its diagonal Fisher, each client performs one additional pass over its local data, computing per-example gradients of the observed-label log-likelihood. This is a modest cost compared with local training, which takes multiple local epochs in our experiments. On the server, each inner iteration for layer ℓ requires, for every client, building the cost matrix in O(m2ℓ mℓ−1 ) time and solving the assignment in O(m3ℓ ) time with the Hungarian algorithm. With at most S = 8 inner iterations, the merge is cheap compared with client training. Computing Infrastructure. Experiments were conducted on a Linux workstation running Ubuntu 20.04 LTS, equipped with an Intel Xeon E5-2697 v4 CPU @ 2.30 GHz (36 cores, 72 threads), 384 GB RAM, and two NVIDIA RTX A6000 GPUs (48 GB VRAM). Our implementation is based on PyTorch 2.8 with CUDA 12.8.

D

ADDITIONAL EXPERIMENTAL RESULTS

D.1

Extreme non-i.i.d. setting

We further evaluate one-shot performance under more extreme heterogeneity by partitioning CIFAR-10 among K = 6 clients with Dirichlet parameter α = 0.05, so that each client’s data is concentrated on very few classes. All other settings follow Section 5.1. Table 3 reports the results.

Table 3: One-shot test accuracy (%) on CIFAR-10 with K = 6 clients and α = 0.05. Mean ± std over three seeds. Method Accuracy FedAvg OTFusion PFNM RegMean DENSE FisherMerge FedFisher (K-FAC)

25.5±1.7 27.8±2.1 20.5±3.0 26.6±2.6 26.4±0.7 28.0±3.0 31.9±1.1

FedFiOT (Diag)

41.4±1.7

FedFiOT outperforms the strongest baseline, FedFisher (K-FAC), by 9.5 percentage points and FedAvg by 15.9 points. The advantage of FedFiOT grows with heterogeneity: on CIFAR-10, its gap over FedFisher (K-FAC) increases from 3.3 points at α = 0.2 and 4.7 points at α = 0.1 to 9.5 points at α = 0.05. Moreover, FedFiOT loses only 1.9 points when α decreases from 0.1 to 0.05, whereas FedFisher (K-FAC) loses 6.7 points and PFNM loses 12.3 points. These results are consistent with our motivation: under stronger heterogeneity, client models drift further apart, neuron misalignment becomes more severe, and jointly aligning neurons and weighting parameters by their Fisher importance becomes more beneficial.

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

D.2

Comparisons with Ensemble-based Methods

We compare FedFiOT with FedBEns Talpini et al. (2025), which returns an ensemble rather than a single model. Each FedBEns client trains five models, one from each of five shared random initializations, and fits a Laplace approximation around each. The server then constructs five global models from the resulting mixture posteriors (Appendix C.2). We evaluate three variants that differ in their curvature approximation: Kron (Kroneckerfactored), Diag+Full (diagonal, with a full covariance for the last layer), and Diag (diagonal). For Kron, we report the validation-selected version, which is the strongest. Figure 3 reports the results. FedFiOT outperforms FedBEns Diag and Diag+Full in every setting. Against the strongest variant, FedBEns Kron, FedFiOT is more accurate in three of the six settings (by 2.5 to 4.6 points on SVHN with α = 0.2 and on CINIC-10) and less accurate in the other three (by 0.4 to 1.9 points). FedFiOT achieves this comparable accuracy at a substantially lower cost, summarized in Table 4. FedBEns requires each client to train five models instead of one and to upload all five models together with their curvature estimates, and it requires five forward passes for every prediction. In contrast, FedFiOT trains a single model per client, matches the upload cost of FedAvg, and returns a single model.

Figure 3: One-shot test accuracy (%) of FedFiOT and FedBEns with six clients, mean ± sample standard deviation over three training seeds. FedFiOT returns a single model; each FedBEns variant returns a five-model ensemble. FedBEns Kron selects server checkpoints on validation data.

Table 4: Cost comparison between FedFiOT and FedBEns per client and at inference, relative to FedAvg. FedFiOT FedBEns Models trained per client Upload per client Models stored at inference Forward passes per prediction

D.3

1 1× FedAvg 1 1

5 > 5× FedAvg 5 5

Different client initialization

In the main experiments, all clients start local training from the same broadcast model, which partially aligns their neurons before merging. Here, each client instead initializes its model independently at random, as in model merging of independently trained networks (Singh and Jaggi, 2020; Ainsworth et al., 2022). This setting arises when clients hold pre-existing models or when the server cannot coordinate the initialization. All other settings follow Section 5.1. Table 5 reports the results. Without a shared initialization, methods that average parameters coordinate-wise degrade sharply. On CIFAR10 and CINIC-10, FedAvg, FisherMerge, and RegMean fall to 13–21% accuracy, close to the 10% of random guessing. FedFisher (K-FAC), the strongest baseline with a shared initialization, loses 5.2–19.0 percentage points. Neuron-matching methods are more robust, and PFNM becomes the strongest baseline in four of the six settings. FedFiOT, in contrast, retains nearly the same accuracy as with a shared initialization (Table 1) in all settings except SVHN with α = 0.1. As a result, its margin over the strongest baseline grows from 3.3–7.6 points with a

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

Table 5: One-shot test accuracy (%) with independently initialized clients, K = 6. Mean ± std over three seeds. SVHN

CIFAR-10

CINIC-10

Method

α = 0.2

α = 0.1

α = 0.2

α = 0.1

α = 0.2

α = 0.1

FedAvg OTFusion WB Fusion PFNM RegMean DENSE FisherMerge FedFisher (K-FAC)

41.2±4.3 55.6±0.5 59.7±2.6 68.9±3.1 32.0±1.6 65.5±0.9 30.8±3.0 57.3±1.4

17.6±1.3 30.2±1.4 34.3±0.3 39.3±0.3 25.4±2.6 47.0±1.1 18.9±4.6 42.6±5.5

16.3±2.5 36.5±1.0 38.6±1.6 42.8±2.7 20.8±1.4 34.0±2.3 14.3±1.7 34.2±3.4

18.9±2.8 30.6±0.4 27.4±1.7 33.8±1.4 18.2±1.1 28.4±0.6 18.1±3.6 30.6±3.2

14.9±0.8 33.8±0.5 35.7±0.9 38.4±1.5 16.7±1.1 26.0±1.8 15.1±1.7 29.3±0.5

19.6±2.6 28.4±1.3 25.8±1.1 25.4±0.9 17.8±1.1 22.9±1.1 13.2±1.2 26.8±1.1

FedFiOT

79.7±0.3 58.4±4.1 52.1±2.0 44.8±1.5 45.2±2.2 38.6±1.4

shared initialization to 6.8–11.4 points. Alignment is therefore essential once clients do not share an initialization, but alignment alone is not enough: OTFusion and WB Fusion also align neurons, yet with a Euclidean metric they trail FedFiOT by a significant margin. D.4

Fraction of Neurons Permuted per Round

To understand how the role of alignment changes across rounds, we record the fraction of hidden neurons that FedFiOT permutes in each round of the few-round experiments (Section 5.3). Figure 4 reports this fraction.

Figure 4: Fraction of hidden neurons (%) permuted by FedFiOT in each round, K = 6 clients and α = 0.1. Lines show the mean over three seeds and shaded regions show ± one standard deviation. In the first round, FedFiOT permutes about 84% of hidden neurons on both datasets. This indicates that client models trained from a shared random initialization are substantially misaligned after local training. In later rounds, the permuted fraction decreases. On CIFAR-10, it drops sharply to 14–23% from the second round on. On CINIC-10, it decreases more gradually, from 84% to 44% by the fourth round.

E

FEDFIOT’S PSEUDOCODE

Algorithm 1 summarizes the full FedFiOT procedure for one-shot and few-round FL.

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

Algorithm 1: FedFiOT for One-shot and Few-round FL (0) Input: Client datasets {Dk }K ; communication rounds T ; local steps E; k=1 ; initial model W maximum inner iterations S ≥ 1; damping λ > 0; tempering exponent ρ ∈ (0, 1] Output: Global model W(T ) 1 for t = 0, . . . , T − 1 do 2 Select a nonempty client set K(t) 3 Broadcast W(t) to clients in K(t) 4 foreach k ∈ K(t) in parallel do 5 Wk ← LocalTrain(W(t) , Dk , E) e k ← DiagonalEmpiricalFisher(Wk , Dk ) 6 F e k ) to the server 7 Upload quantized versions of (Wk , F e k ← (F e k )ρ for all k ∈ K(t) 8 Set πk0 ← Im0 for all k ∈ K(t) . Set F 9 for ℓ = 1, . . . , L do // Reorder incoming coordinates using the previous layer 10 foreach k ∈ K(t) do 11 for j = 1, . . . , mℓ do ℓ ℓ 12 wk,j ← πkℓ−1 wk,j ℓ−1 ℓ ℓ 13 fe ← π fe (π ℓ−1 )⊤ k,j

14 15

16 17 18 19 20

ℓ,[0]

k,j

k

k

ℓ ← wref

w // Largest-client reference or broadcast reference for s = 1, . . . , S do // Fisher-weighted neuron alignment foreach k ∈ K(t) do if ℓ < L then for i, j = 1, . . . , mℓ do  ℓ,[s]  ℓ,[s−1] ℓ,[s−1] ℓ ℓ ℓ Ck ← (wi − wk,j )⊤ (fek,j + λI)(wi − wk,j ) ij ℓ,[s]

ℓ,[s]

πk ← Hungarian(Ck else ℓ,[s] πk ← Imℓ

21 22

)

for i = 1, . . . , mℓ do mℓ X ℓ,[s] ℓ ℓ w bk,i ← πk,ij wk,j

23 24

j=1

ℓ,λ fbk,i ←

25

26

27

mℓ X j=1

ℓ,[s] ℓ πk,ij fek,j + λI

// Fisher-weighted neuron aggregation for i = 1, . . . , mℓ do  −1 X ℓ,λ X ℓ,λ ℓ,[s] ℓ wi ←  fbk,i  fbk,i w bk,i k∈K(t)

28 29 30 31 32 33

34 35 36

k∈K(t)

if ℓ = L then break

ℓ,[s]

if s > 1 and πk break

ℓ,[s−1]

= πk

for all k ∈ K(t) then

wℓ,(t+1) ← wℓ,[s] ℓ,[s]

Set πkℓ ← πk

for all k ∈ K(t)

// Keep the alignment only if it does not increase the merge objective
if J (t) W(t+1) , (πkℓ )k,ℓ > minW J (t) (W, I) then W(t+1) ← arg minW J (t) (W, I) return W(T )

Quan Nguyen, Hoang Ngo, Bhakti Chowkwale, Y. Keshtkarjahromi, My T. Thai

F

ABLATION STUDY

To isolate the contribution of each component of FedFiOT, we evaluate variants that remove or replace one component at a time: neuron alignment, the Fisher-weighted matching cost and Fisher-weighted aggregation. The variants are: (A) FedFiOT: Fisher-weighted matching and aggregation, S = 8. (B) No alignment: identity permutations, with the same Fisher-weighted aggregation. (C) Euclidean matching: unit weights in the matching cost, with the same Fisher-weighted aggregation. (D) Uniform aggregation: Fisher-weighted matching, with the aligned clients averaged uniformly. (E) WB Fusion: Euclidean matching and uniform aggregation, i.e. (C) and (D) combined (Akash et al., 2022).

Table 6: Ablation of FedFiOT’s components on K = 6. Test accuracy (%), mean ± std over three seeds. SVHN

α = 0.1

CIFAR-10

α = 0.2

(A) FedFiOT

78.9±1.9 66.8±0.8 52.0±0.9 43.3±1.8 41.4±1.7 45.6±1.4 39.6±0.3

(B) No alignment (C) Euclidean matching (D) Uniform aggregation (E) WB Fusion

60.4±2.7 77.3±0.2 61.3±1.7 60.8±1.0

42.4±5.6 49.9±4.7 38.1±1.0 34.7±2.3

α = 0.2 40.1±4.2 50.0±1.3 38.2±1.9 40.3±1.2

α = 0.1

CINIC-10

Variant

33.2±1.5 34.4±1.8 31.7±1.2 29.0±1.1

α = 0.05 27.6±3.9 27.7±4.1 21.1±0.7 18.3±0.7

α = 0.2 33.1±2.0 42.2±1.2 34.4±1.0 36.4±1.1

α = 0.1 28.2±1.6 32.0±0.5 26.2±0.9 25.4±0.8

Each component of the merge is necessary. Removing alignment (B) reduces accuracy by 10–24 points, and replacing Fisher-weighted with uniform aggregation (D) by 11–29 points. Replacing the Fisher-weighted matching cost with a Euclidean one (C) costs 2–3 points at α = 0.2 but 8–17 points at α ≤ 0.1: the more heterogeneous the clients, the more it matters that neurons are matched under the same metric in which they are averaged. Removing both Fisher weightings yields WB Fusion (E), the weakest variant in four of the seven settings. Hence neither alignment nor Fisher-weighted aggregation alone explains FedFiOT’s gains, but using the Fisher metric consistently in both steps does.

G

NUMERICAL EXAMPLE

We illustrate the alternating updates of Section 3.3 on a hidden layer with three neurons, merged from two clients. Each neuron has one scalar incoming weight, and we suppress ℓ. Let w[0] = w1 = (0, 4, 8)⊤ ,

w2 = (10, 2, 6)⊤ .

Client 2 learns the same three features as client 1, shifted by +2 and stored in a different order: its first neuron corresponds to global neuron 3, its second to global neuron 1, and its third to global neuron 2. Given λ = 1 and Fisher values (2, 0, 2) for client 1 and (0, 0, 2) for client 2, the damped Fisher values are (3, 1, 3) and (1, 1, 3), respectively. First alignment step.

Equation (4) gives 

 0 16 192 [1] 0 48  , C1 =  48 192 16 0

 100 4 [1] C2 =  36 4 4 36

 108 12  , 12

where rows index global neurons and columns index client neurons. Client 1 coincides with w[0] , so it retains its order at zero cost. For client 2, write an alignment as the client neurons assigned to global neurons 1, 2, 3. The

FedFiOT: Fisher-Weighted Optimal Transport for Model Merging

identity (1, 2, 3) costs 116; the swaps (2, 1, 3), (3, 2, 1), and (1, 3, 2) cost 52, 116, and 148; and the cycles (2, 3, 1) and (3, 1, 2) cost 20 and 180. The assignment (5) therefore selects the cycle (2, 3, 1),   0 1 0 [1] π2 = 0 0 1 , 1 0 0

which recovers the true correspondence; every swap leaves two neurons mismatched. Client 2’s Fisher values move with its neurons, so by (6) its aligned weights and damped Fisher values are (2, 6, 10) and (1, 3, 1), respectively. First aggregation step.

Equation (7) gives     (3 · 0 + 1 · 2)/(3 + 1) 0.5 w[1] =  (1 · 4 + 3 · 6)/(1 + 3)  = 5.5 . (3 · 8 + 1 · 10)/(3 + 1) 8.5

Each merged weight moves toward the client with greater precision for that neuron: toward client 1 for global neurons 1 and 3, and toward client 2 for global neuron 2. Second alignment step. With w[1] as the reference, the cheapest assignments are again the identity for client 1 (cost 3.75; next best 39.75) and the cycle (2, 3, 1) for client 2 (cost 5.25; next best 41.25). Since no alignment changes, the layer merge stops after two iterations. The objective decreases from J [0] = 116, the cost of the identity alignments at w[0] , to J [1] = 3.75 + 5.25 = 9, as Theorem 1 guarantees. Comparison. Table 7 compares FedFiOT with merges that omit alignment, Fisher weighting, or both. Without alignment, global neuron 1 combines client 1’s weight 0 with client 2’s weight 10, although the two neurons encode different features. Fisher merging without alignment attains J = 83. In this single-layer example the full-model objective equals J , so the identity-aligned comparison at the end of the merge keeps the aligned result. Table 7: Merged layer and objective J for the three-neuron example. Methods without alignment use identity alignments for both clients. Merge

Alignment

w⊤

Average by index Fisher merging Average of aligned weights FedFiOT

identity identity (2, 3, 1) (2, 3, 1)

(5, 3, 7) (2.5, 3, 7) (1, 5, 9) (0.5, 5.5, 8.5)

J

108 83 12 9

Related documents

• Federated learning with continual update for privacy-preserving clinical event prediction across distributed hospitals using MCN-GNN.#30788
• TRANSFER LEARNING UNDER HETEROGENEITY: METHODS AND THEORY FOR PREDICTION AND INFERENCE#929341
• Weight Space Learning for Certifiable Few-shot Transfer Learning#1402124
• Linear Independent Component Analysis via Optimal Transport#370313
• GeoTTER: Leveraging Local Geometry of Optimal Transport for Zero-Shot Classification#1071302