Leveraging Function Space Aggregation for Federated Learning at Scale
Authors: Nikita Dhawan
Nicole Elyse Mitchell
Zachary Charles
Zachary Garrett
Gintare Karolina Dziugaite
Venue: Accepted by TMLR ()

Abstract: The federated learning paradigm has motivated the development of methods for aggregating multiple client updates into a global server model, without sharing client data. Many federated learning algorithms, including the canonical Federated Averaging (FedAvg), take a direct (possibly weighted) average of the client parameter updates, motivated by results in distributed optimization. In this work, we adopt a function space perspective and propose a new algorithm, FedFish, that aggregates local approximations to the functions learned by clients, using an estimate based on their Fisher information. We evaluate FedFish on realistic, large-scale cross-device benchmarks. While the performance of FedAvg can suffer as client models drift further apart, we demonstrate that FedFish is more robust to longer local training. Our evaluation across several settings in image and language benchmarks shows that FedFish outperforms FedAvg as local training epochs increase. Further, FedFish results in global networks that are more amenable to efficient personalization via local fine-tuning on the same or shifted data distributions. For instance, federated pretraining on the C4 dataset, followed by few-shot personalization on Stack Overflow, results in a 7% improvement in next-token prediction by FedFish over FedAvg.

Published in Transactions on Machine Learning Research (02/2024)
Leveraging Function Space Aggregation for Federated Learning at Scale
Nikita Dhawan*
wilkia@@cs.oronto.ccom edu
University of Toronto and Vector Institute.
Nicole Mitchell
nicolemitchell@ggogle.com
Google Research
Zachary Charles
machcharles@googleee
Google Research
Zachary Garrett Google Research
zachgareett@@goole.comm
Gintare Karolina Dziugaite
gkd=:ggoole..omm
Google Deed and Mila Quebec AT Institute
Reviewed on OpenReview: https: // openreview. net/ farum?? id= 1tp9KFKZfZ
Abstract
The federated learning paradigm has motivated the development of methods for aggregating multiple client updates into global server model, without sharing client data. Many fed- erated learning algorithums, including the canonical Federated Averaging (FEDAvG), take direct (possibly weighted) average of the client parameter updates, motivated by results in distributed optimization. In this work, we adopt. a function space perspective and propose new algorithm, FEDFISH, that aggregates local approximations to the functions learned by clients, using an estimate based on their Fisher information. We evaluate FEDFISH on realistic, large-scale cross-device benchmarks. While the performance of FEDAVG can suffer as client models drift further apart, we demonstrate that FEDFISH is more robust to longer local training. Our evaluation across several settings in image and language benchmarks shows that FEDFISH outperforms AVD as local training epochs increase. Further, FED- FISH results global networks that are more amenable to efficient personalization via local fine-tuning on the same or shifted data distributions. For instance, foderated pretraining on the C4 dalaset, followed by few-shot personalization Stack Overflow, results in 7% improvement in next-token prediction by FEDFISH over FEDAVG.
1 Introduction
Methods for aggregating separately trained noural networks have received renewed attention machine learning models and data reach ever larger scales (Wortsman et al. 2022 Li al. 2022 Rame el. al.
2023 Parallel training can yield gains in computational efficiency AAssran et al. 2020 Li et al. 2022) and
meet constraints on access to private data, as in Federated Learning (FL; McMahan et al. 2017). Model aggregation plays a central role in how clients collaboratively train model in distributed manner via FL without sharing data among each other or with any orchestrating server. Given that the cross-device FL setting is characterized by client data heterogeneity, unreliable client availability and network constraints [Karrouz al. 2021). FL is typically carried out over multiple communication rounds in which updates from
"Equal contribution.
Work done as Student Researcher Google Research.

---

Published in Transactions on Machine Learning Research (02/2024)
local training are aggregated to iteratively improve the global model. The canonical approach to aggregation, implemented by the FEDAVG method and its adaptive variants (Reddi et al. 2020), is to combine client model parameter updates by averaging them, weighted in proportion to their respective dataset sizes.
In this work, we take a function space perspective (Benjamin et al., 2018) of model aggregation in FL, where we aim to obtain a global model that simultaneously matches each client model's logit outputs on that client's data. One motivation for this is to allow for, and reap the benefits of, more local training between global updates. The performance of algorithms like FEDAVG depends heavily on the number of local training iterations. When the data are heterogeneous across clients, training too long between communication rounds leads to updates that hurt the global model's performance, phenomenon known as client drift (Karimireddy et al. 2020). Indeed, prior work has shown that the number of local training steps dictates a trade-off ween the speed of convergence and the quality of the resulting model (Charles & Konečný] 2021 Malinovskiy et al. 2020 Pathak & Wainwright, 2020). These results imply that selecting the number of local steps is critical. It also challenging, in part, duc the difficulties of hyperparameter tuning in federated settings (Kuo et al.) 2023), and in part. because the number of local steps has wide ranging effects. These include not only the speed of convergence (Pathak & Wainwright 2020 Mitra et al., [021], but also the optimization dynamics (Charles & Rush 2022) and even whether the method acts as a mcta-learnes (Collins et al.. 2022: Charles et al. 2023). While there are a variety of FL methods aimed at mitigating client drift (see Wang et al. 2021) for an overview), many of these introduce extra hyperparameters to tune and remain sensitive to the number of local steps (Mitra et al.) 2021 Charles & Konečný 2021p. Instead, we argue that the choice of model aggregation technique plays a role in the client drift problem. Taking a function space view of the client models sidesteps the drift problem by aiming for global model that in the function space more accurately reprosents cach client model.
A key obstacle to our function space approach, matching the client models' outputs on client data, is its dependence on client data. FL constrains access to client data, preventing a direct approach where the global server uses client data to match the corresponding model outputs. As a step towards parametric function space model aggregation which does not require direct access to client data, we propose and implement a Fisher-weighted federated averaging algorithm, called FEDFISH. This method is derived from an objective that minimizes an approximate function space distance between local client models and the global model. The closed-form solution of this objective depends the networks' Fisher Information (Cover 1999), , which are typically too expensive to compute and store for large models. The approximations required to implement our method in practice lead to a parametric aggregation scheme which accounts for the client data distributions via the functions represented by their local models. We investigate the advantage this confers upon FEDFISH over simple averaging AVEDAvG) in regression, image classification and language modeling benchmarks.
Our extensive evaluation includes domain-specific criteria as well as metrics specific to FL. We demonstrate settings in which FEDFISH outperforms FEDAVG, especially as the amount of local training is varied. Image and language experiments with varying levels of client data heterogeneity show improved post-personalization performance of FEDFISH throughout training, when the global model is locally finc-tuned for a few steps by clients that were held out during training. This observation also holds when measuring transfer performance by drawing the evaluation clients from a shifted data distribution. For instance, in an experiment with federated pretraining on the large and hetergenous C4 dataset, followed by few-shot personalization on Stack Overflow clients, FEDFISH is able to improve upon FEDAVG's next-token prodiction performance by 5-7,, depending on the amount of personalization data available. We provide insight into these gains by assessing a measure of deviation between global and local models, coined the Client-Server Barrier. Finally, we discuss the impact of these methods and settings on the cost of communication between clients and the server.
Contributions.
We formalize a function space perspective of federated learning to motivate scalable algorithm, FEDFISH, which aims to match client input-output functions during aggregation.
Via a synthetic example, we demonstrate that FEDFISH outperforms FEDAVG as client data hetero- geneity increases. We then investigate this performance at larger scales than have been explored by previous works.
2

---

Published in Transactions on Machine Learning Research (02/2024)
Given two functions modeled over disjoint supports (left), direct parameter average fails to represent either function well (center), while function space aggregation aims to preserve both functional relationships (right)..
Figure 1:
Our thorough empirical results show that FEDFISH allows for longer local client training compared to FEDAVG. We find that the global models learned via FEDFIIH have greater ability to be per- sonalized via fine-tuning on the same or shifted data distributions, indicating they provide better initialization for local training in each round.
We propose to evaluate effects of aggregation via a Client-Server Barrier, leveraging the function space perspective gain further insight into the observed results.
2
Federated Learning in the Function Space
We now define the federated learning problem from function space perspective and describe the approxi- mations that lead to practical and parametric objective.
2.1 Problem Setting
Let network parameterized by 0 be trained on a dataset of input target pairs, D = (X,y), optimize loss function C, such that it reprosents function f(X;0) Z, where Z denotes the network's outputs across all inputs. Consider the canonical federated learning setting, where a global model'ss parameters lG are broadcast to N clients for local training. Each client i trains their model on local data D; (with corresponding set of input data X,) for a fixed number of iterations to produce trained parameters 0,. These parameters are then communicated back to a global server where they are aggregated using a specific aggregation technique. This procedure is repeated over multiple rounds, where the aggregated model from each round serves as the initialization for local training in the subsequent round.
Viewing FL from function space perspective, the aggregated model should ideally match the input-output output relationships learned by each client so far. More formally, for each federated round, define the oplåmal global model as the one that produces outputs closest to each client's outputs when evaluated on the corre- sponding input data. Let us denote this function space distance as D (s).). Averaging this quantity over all clients, we obtain the following objective:
0 = arg min
Σ ,((;,))),((x;,0)
(1)
N
We depict this idealized objective in fig. where two client functions are learned different supports (left) and we wish to aggregate them into model that preserves both functional relationships (right), which direct parameter averaging cannot achieve (center).
The objective in eq. (1) depends on function outputs, which in turn rely on client-specific inputs X,. Exactly implementing this would require global access to the local client data, which violates a fundamental constraint in federated learning. This necessitates a parametric approximation to the function space distance such that it may be estimated without actual data points.
3

---

Published in Transactions on Machine Learning Research (02/2024)
2.2 Approximating Function Space Distance
The function space distance can be estimated with second-order Taylor approximation with respect to model parameters 0, centered at 0,, the client network whose outputs are to be matched. This is useful in a fedorated context because appropriate approximations to this estimate lead to a paramctric method which does not directly depend on client data.
Setting D (-:-) to be the
(KL) divergence between softmax outputs of the networks,
f(x( ~ (0 0.)² F, 0,)
(2)
0,1 ( (B0) -((0)²
~
(3)
2
where F; the Fisher Information matrix corresponding to 8 In eq. 2 the zero-th and first order terms vanish because D (..)) is a non-negative function that evaluates to zero when its arguments are equal. Hence, its value and gradient both vanish
0,, leaving only the second order term. Note that this approximation does not require 0g to be optimal and can be used at intermediate stages of training in multi-round FL. We defer complete details of this derivation lo appendix A.2
The full Fisher Information matrix is expensive to compute and store for large scale networks. Further, the corresponding closod-form solution to this optimization problem would involve the inverse sum of Fisher Information matrices, which need not be invertible in practice. Common approximations (Kirkpatrick et al. 2017) involve using the diagonal empirical Fisher Information matrix, shown in eq. (3) where F. is the j-th diagonal entry of F,.
3 FedFish Algorithm
Given function space distance approximations in section 2.2 we now obtain a parametric aggregation scheme that can be practically implemented in federated settings. Plugging eq. optimization problem gives:
3 into eq.
(T) and solving the convex
[0.1
Σ == (F)
3 (0() - -
0 arg min
(4)
2N
Di=1 diag( F.)
where diag(Fi) represents the diagonal of F.. Hence, the global model at each round is the Fisher-weighted average of client models, normalized by the sum of all Fisher diagonals. In this form, FEDFISH is simple to implement and efficient to deploy in cross-device federated seltings. The empirical Fisher diagonal can be computed during local training, via an average of squared gradients of the training objective with respect to mini-batches of data, as shown in algorithm [ Similar to common FEDAVG implementations, the clients communicate model deltas to the global server, which are then aggregated and used as pseudo-gradients in one step of global model optimization. FEDFISH can be combined with adaptive optimization ion techniques for different choices of global optimizers. We show this procedure for SGD with global learning rate Ty and local learning rate ne in algorithm ! See appendix A.5 for a discussion efficiency improvements.
Note that we recover the FEDAVG algorithm if the Fisher coefficients for each parameter are the same across client networks. Such a condition is unlikely to hold even approximately, especially in settings with increasing data. heterogeneity. Fisher coefficients account for the influence of each parameter its network's predictions with respect to corresponding inputs. Intuitively, FEDFIIH has the following advantages: (1) Each local parameter's contribution to the global aggregate is proportionate to mcasure of its importance for making predictions on its training data, which improves upon the FEDAVG global model. (2) Derived from a function space perspective, FEDFISH simply aims to match the functions learned by clients so far, without relying on optimality assumptions. This makes compatible with heterogencous datasots, longer local training as well as multiple rounds of FL. We empirically test for these gains across several settings ranging from toy to large scales.
4

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm FEDFISH (SGD)
Algorithm 2 FedLocallhrin (SGD)
Require: rounds R, local epochs E, fg. client
Require: E, 0, Di.Te 1: ModelDelta, SumFisher 0,0
datasets D, global Ir fl.. local Ir T/
for r 1 to R do
for E f- to E do
Sample cohort of clients, C
for b E D; do
3: for E C' in parallel do
4: - b),(0,,b) 0g + 0, -
4: Compute client weights, w) = Di]
5:
5:
0, + fG
6:
ModelDelta + ModelDelta + 1/
20,, F; FedLocal
D.. mo)
7: end for end for
end for
N W,FT A0:
for b E D; do
06 t 8G 1l
1==1
10: g (oL(O.,)))
E WiF:
11: SumFisher SumFisher + g²
end for 10: return 0g
12: end for
13: return ModelDelta, SumFisher
4
Evaluation and Client-Server Barrier
While there are natural domain-relevant evaluation metrics for our benchmarks, here, describe specific criterion relevant to FL. The commonly used global and personalization performance useful indicators of successful FL algorithms. However, they may be confounded by local optimization choices. Hence, we formalize the Slivet-Server Barrier below and additionally evaluate it in section 5 as more direct measure of the quality of given aggregation method. The criteria below are defined in terms of a performance metric Lz(+) that measures a quantity of interest, such as loss or prediction error, on client data D,.
Global Performance. The most natural measure of success in FL is the performance of the global model on held-out client data, averaged over clients. Using our notation from before, this is given by Σ (,(G)..
Client Personalization Performance. While global performance is akin to network's zero-shot abilities, we are often interested in its personalization ability to unseen clients after a few steps of fine-tuning, Quick adaptation of the network to particular clients or downstream use cases critical in the compute-constraincd settings that FL targets. We measure client personalization by fine-tuning fG on a portion of each held-out client datasct and then evaluating each fine-tuned model on the remaining unseen client data (the same portion used for measuring global performance), averaging over clients before.
Client-Server Barrier. For given performance metric L and aggregation technique, we define the Client- Server Barrier as the difference in this performance metric between each client model 0 and the aggregated global model Og with respect to the client data D., averaged over all the clients involved in the aggregation. Mathematically, this is given by
N
Σ (L((0)) N Σ (.(g))
(..)).
(5)
N
This is simple and direct measure of the impact of aggregation on model performance and can be computed in federated manner: averaging the broadcast global model's performance across all client data in the sampled cohort (first term), and averaging trained local models' performance on their respective client data across the sampled cohort (second term).
5 Experiments
Using the criteria described section 4 we now conduct a systematic empirical evaluation of FEDFISH in variod scttings, compared to the best performing variant of FEDAVG. We first demonstrate the advantage of FEDFISh as client data heterogeneity increases in toy regression problem. We then assess its performance across settings in larger scale image and language benchmarks.
5

---

Published in Transactions on Machine Learning Research (02/2024)
Client 1
FgdAvg FedFish
Client
FedAvg
350
Client 2
3.0
Client
FedFish
1.5
0.3
0.
2.5
0.31
Client I
FodAvg FedFish
FodAvg Client- Ser-errrrrrrr Barrier
Client 2
0:23
3.5
8:30
2.3
x=y
0.19
Client
Client
6:10
1.2
Full Overlap
6:05
Partial Overlap
Overlap
10
0.56
0.0
0.18 a
0.5
0.20
FedFish Clivet-Server Barrier
Figure 2: As beterogencity across clients increases (top left + top right + bottom left), FEDAVG deteriorates, while FEDFISH matches predictions of both client models. For each setting shown and each client within it, the FEDFISH global model has lower barrier the clients (bottom right).
5.1 A Toy Regression Demonstration
Figure ssoows non-lincar regression problem with clients across which data is distributed with varying heterogeneity, including full (+), partial (o) and no overlap (+). We plot the local functions learned by each client, as well the global functions produced by aggregation via FEDAVG and FEDFISH after one round. For completely homogeneous client data, FEDAVG and FEDFFISH fit similar functions. When there is partial overlap, FEDAVG scems to reasonably retain predictions on one client's data while poorly fitting the other, while FEDFISH fits both datasets well. In the extreme case of completely disjoint supports, FEDAVG fails fit either client dataset, but FEDFISH matches the locally learned functions of both clients on their respective input data. The Client-Server Barrier (CSB) defined in eq. (5) is computed in terms of mean squared error for each client on their corresponding data. As shown by all the points above the y line in fig, 2 (bottom right), the CSB is lower for FEDFISH than FEDAVG in each of these settings, with more significant difference as data heterogeneity increases. We hypothesize that accounting for the functions learned by local models confers this advantage upon FEDFISH.
5.2 Image Classification and Text Benchmarks
Datasets and architectures. We consider a variety of federated benchmarks for image classification (EMNIST (Cohen et al.) 2017), CIFAR100 (Krizhevsky al.)) and language modeling (Stack Overflow AAthhor 2019), CC-News HHmmorr et al. 2017) and C4 (Raffel et al.] 2020)). In particular, C4 is large- scale and significantly heterogenous dataset. For these domains, we use standard classifier and transformer architectures, respectively.
For EMNIST, we partition the handwritten characters according to their author, as proposed by Caldas et (2018). For Stack Overflow, posts on the eponymous web forum are partitioned by their author well. For CIFAR100, we partition the examples over 100 clients in heterogeneous fashion using the

---

Published in Transactions on Machine Learning Research (02/2024)
Post Pww.Peerrzzzaatiiioo Accuracy
Global Accuracy
6.5
2-4
FedAvg
FedAvg
FedFish
FedFish
local epochs
local epochs
local epochs
local epochs
local epochs
local spochs
Training Iterations
Training Tterations
Figure 3: Training on EMNIST with FEDFISH converges faster and to a higher global accuracy (left) and post-personalization accuracy (right) than training with FEDAVG, across varying numbers of local epochs. Results are shown with fixed compute across configurations: each training iteration corresponds to local epoch, and each marker indicates 100 federated communication rounds.
two-level latent Dirichlet allocation scheme proposed by Reddi al. (2020). For CC-News and C4, we use Dataset Grouper (Charles al.] (023) to partition web-crawled examples according to their base URL (e.g. mytimes. com). More details about data splits, architectures and hyperparameters are included in appendix ix A.3
Performance metrics. We evaluate global performance, client personalization performance and Client- Server Barrier (see section 4 using standard domain-relevant performance metrics. These include classifi- cation accuracy for images and next-token prediction accuracy and perplexity for language modeling. Since C4 is very large scale dataset that may generally be used as pretraining corpus, we evaluate its global model on held-out clients from C4 itself as well as on the shifted Stack Overflow and CC-News datasets. This tests the methods' transfer performance in addition adaptability to new clients that were not seen during training. For the C4 experiments, we also vary the amount of fine-tuning data used for personalization 25% or 50% of each held-out elient's data- to assess few-shot performance. Additional details are reported in
5.2.1 Effect of Local Training Global Model Performance
We study the effect of local training on the global model's performance by varying the number of epochs of training clients perform in between rounds. Since increasing the mumber of local epochs for fixed number of rounds increases computational costs, we present our results by separately fixing compute (or total number of training iterations) and number of aggregation rounds. We provide complete table of results covering all settings in table 3 of appendix A.4 and discuss representative experiments here. With fixed compute, fig. 3 (left) shows a decline in global accuracy as number of local epochs is increased for both FEDAVG and FED as expected by the client-drift phenomenon. However, within each setting, and across datasets, FEDFFis outperforms FEDAVG, suffering more graceful decline in performance with more local training and converging to higher performance faster. Figure similarly shows FEDFISH outperforming FEDAVG in terms of classification accuracy and next-token prediction accuracy for CIFAR100 (left) and Stack Overflow (right) datascts, respectively. In both of these settings, FEDFISH is much more robust to longer periods of local training that FEDAVG, whose performance suffers as local training increases. In fig. 5 (left), we use the heterogeneous C4 dataset for federated training and evaluate its zero-shot performance on unseen clients from C4 itself. The bottom row of the plot shows that the global model, without any personalization data, benefits considerably from increased local epochs during federated pretraining for both algorithms, with FIDFFih putprrrorinng FEDAVG:

---

Published in Transactions on Machine Learning Research (02/2024)
Paziiaaaatiie Training Dau
Prosoaali Pasoraliaatie nasiver Training Data
Podira FolFidd
FalFok
hoad
Accuracy
Acceracy
Figure Global and post-personalization performance in terms of classification accuracy on CIFAR100 (left) and next-token prediction accuracy on Stack Overflow (right). Varying number of local training epochs can significantly impact FEDAVG performance while FEDFISH remains relatively robust this.
Traning Denrs)
Panseaaiaann weaaiiratiie Traning Datat
attto Tralling Data
FodAvg
baxi gpach
local spochs
Next token Prediction Accuracy
Nest token Preduction Accuray
Nest soken Prediction Accuncy
Figure Transfer (global and post-personalization)) performance terms of next-token prediction after lederated pretraining on C4 and evaluating on C4 (left), Stack Overflow (center) and CC-News (right). Personalizing with 0%, 25% or 50% of held-out client data results in FEDFISH outperforming FEDAVG, especially with longer local training.
5.2.2 Post-Personalization Performance
At scale, pretrained models are often personalized or fine-tuned for a small number of steps using local client data before deployment. Accordingly, we fine-tune the global models for few steps on limited datapoints from held-out clients and evaluate the metrics discussed earlier. Consistent with global performance reported above, we observe across tasks that FEDFISH yields models with higher post-personalization performance than those trained with FEDAVG. This is shown on EMNIST in fig. (right), CIFAR100 in fig. (left), Stack Overflow in fig, 4 (right), and C4 in fig. (left). Notably performance on CIFAR100 trained with FEDAVG using 16 local epochs, while substantially improving CI- FAR100 trained with FEDFISH using the same configuration. By contrast, FEDAVG with 16 local epochs on Stack Overflow improves dramatically with personalization, despite still underperforming all other models. Interestingly, we see in the case of C4 that more than aggregation algorithm the amount of local training seems to impact personalization performance. While both FEDFISH models trained with 16 local epochs have higher zero-shot performance than either FEDAVG model, the 16 local epoch FEDAVG model's post- personalization performance surpasses that of FEDFISH with 1 local epoch as the amount of fine-tuning data increases. Overall, we find that models trained with FIDFISH using longer periods of local training tend to be more amenable to personalization than models trained with FEDAVG.
that personalization worsens
5.2.3 Transfer Performance
Considering FEDAVG and FEDFIIH as methods of federated pretraining, we further evaluate the networks trained on C4 in terms of their transfer performance (fig. 5) on Stack Overflow (center) and CC-News (right). Here, performance is in lerms of next-loken prediction accuracy. We similarly report the perplexily for each of these settings in appendix [A.3.3] We vary the amount of data available for personalization, between 0%, 25% and 50% of each held-out client datasct. The reported performance is always evaluated on the unseen 50% of the data. In each of these settings, we find that transfer performance benefits from longer local training for both methods and FEDFISH yields better zero-shot, few-shot. and post-personalizationn
8

---

Published in Transactions on Machine Learning Research (02/2024)
0.25
1200
0.7
2400
FedAug Client Server Barrier
1000
Client-Servel Server Barrier
0.6
1000
0.15
0.4
800
600
Fed/vg Client
0.10
0.3
600
400
0.05
0.1
0.00
0.0
0.00
0.02
0.04
0,06
0.08
0.10
0.00
0.02
0.04
0.06
0.08
0.10
0.12
0.14
FedFish Client-Server Barrier
FedFish Client-Server Barrier
0.10
10000
1000
FedAvg clent-Server Barrier
0.8
FedAvg Server Barrier
0.08
8000
800
0.6
0.06
600
0.4
4000
200
0.02
2000
0.0
0.00
0.0
0.2
0.4
0.6
0.8
1.0
0.00
0.02
0.04
0.06
0.08
0.10
FedFish Cient-SSrver Barrier
FedFish Client-Server Barrier
Figure 6: Measuring the Client-Server Barriers throughout training illustrates differences in how aggregating via FEDAVG FEDFFISH influences the training trajectory. Colors indicate the federated round. Stack Overflow with 8 local epochs (top left), Stack Overflow with 16 local epochs (top right), CIFAR100 with 16 local epochs (bottom left), and C4 with 16 local epochs (bottom right).
performance than FEDAvG. We observe largest gains in the case of federated pretraining C4, followed by few-shot personalization Stack Overflow clients, where FEDFISH improves upon FEDAVG's next-token prediction performance by 5-7%, depending the amount of personalization data available. These results are promising since they encourage longer local training, which connotes parallelism and efficiency gains. Note, the results in fig. 5 correspond to a fixed number federated rounds, R; we similarly report performance at round R/2 in table of appendix A.4
5.2.4 Client-Server Barrier
To gain more insight into the performance FEDFISH, which only differs from FEDAVG in the aggregation step, we measure the Client-Server Barrier (CSB) defined eq. ( using accuracy as the metric L, for different model checkpoints across rounds of federated training. Note that fix seeds control cohort sampling in each federated round, such that the dataset iteration for FEDAVG and FEDFISH training runs match. In fig. 6 plot this quantity for FEDFIIS networks trained on Stack Overflow with 8 local epochs (top left), Stack Overflow with 16 local epochs (top right), CIFAR100 with 16 local epochs (bottom left) and C4 with 16 local epochs (bollom right). Points lying above or below the x = y line on these plots indicate whether FEDFISH or FEDAVG, respectively, achieves lower CSB in those rounds. We observe the relative round-over-round aggregation performance of FEDAVG and FEDFISH vary significantly for different datasets and settings throughout training. For Stack Overflow, the CSB steadily increases with rounds of training, with FEDFISH achieving lower barrier than FEDAVG
the r-axis and for FEDAVG the y-axis, using
9

---

Published in Transactions on Machine Learning Research (02/2024)
during later stages of federated training. This difference is more stark when training with 16 local epochs compared to 8 local epochs. In the case of C4, while the barrier generally increases with rounds of training, we see that FEDFFIH tends to have lower values in the beginning stages of training while FEDAVG obtains lower CSB towards the final 20% of training rounds. This indicates connections to linear mode connectivity in later stages of training. We discuss this connection in appendix [A.6 and leave deeper explorations to future work. Interestingly, CIFARI00 CSB values decrease with rounds of federated training with FEDFISH achieving lower barrier than AVD throughout. On investigating further, we find that the client and data splits on CIFAR100 are such that each local model achieves very high performance right from the beginning of training, and maintains that performance throughout. Hence, the reduction in CSB as rounds increase is indicative of the improvement of the global model as it bridges its gap to local models. This is in contrast to the other datasets we present, wherein the local models often improve their performance more gradually.
5.3 Communication Cost
In general, cost of communication between clients and the server is directly related to the number of rounds of federated training as well as the number of units of (parametric) information to be exchanged. So far, we have demonstrated that FEDFISH can reduce the number of communication rounds by allowing longer local training. However, the procedure described in algorithm | requires clients to communicate their parameters as well as Fisher diagonals to the global server. Relative to FEDAVG, this increases communication cost of per round by factor of two for FEDFFSH.. However, as demonstrated in figs. 3 and 9 (see appendix AA.5 the advantage of training with FEDFISH for more local epochs can eliminate the communication overhead from our FEDFISH implementation when achieving comparable accuracy to FEDAVG in half the number of communication rounds (compare EMNIST FEDAVG with 8 local epochs after 800 communication rounds to EMNIST FEDFISH with 16 local epochs after 400 communication rounds). Alternatively, our method can be combined with adaptive federated optimization techniques (Reddi et al. 2020) so that clients only have to communicate their weighted parameters, (i,g(FF)" and the normalization step is folded into the server optimization. In this case, the communication cost of FEDFFISH would be the same, per round, as that of FEDAVG. We leave this adaptive extension of our method to future work.
6 Related Work
Federated Learning. Federated Learning (FL.) is well-established paradigm, introduced in McMahan et al. (017)] and advanced through variants that account for adaptive optimization (Reddi et al. 2020).
client drift (Karimiredly et al., 2020: Dandi et. al.] 2022) or helerogeneity (Li et al.] 2020). Recent works have also explored its connections to representation learning (Collins et al. (022) and meta-learning (Charles
ctt al. 2023). Relevant to our function space perspective of FL are frameworks that view FL as a distributed
inference problem. Al-Shedivat et al. (022)) and Guo et al. (2023) aim to approximate local posterior distributions over client parameters, deriving MCMC-based and variational inference objectives, respectively. Performance of both these methods rely on a "burn-in" period of FEDAVG training, after which the proposed algorithms are applied. The amount of burn-in training is a crucial hyperparameter and given this setting, Hou et al. 2022) find that simply chaining FEDAVG and FEDSGD is actually a theoretically sound, efficient alternative.
This work evaluates performance in settings with varied amounts of local training. At the extreme of the local training spectrum are methods that operate in the one-shot setting (Guha et al. 2019), to take advantage of local models that are independently trained to convergence and aggregated only once. This single-round training is a special case of FL and leads to aggregation objectives derived for optimal local models. In practice, not all clients are available simultaneously and coordinating single round of FL is unrealistic when there are millions of clients. In contrast, the more general multi-round setting has the advantage of allowing for new clients and for clients to benefit from each other indirectly. This is the setting our work has locused on, where at each round, client. models are inltiel from the global model obtained in the previous round, intuitively allowing future clients to leverage information aggregated previously.
Concurrent to our work, Jhunjhunwala et al. (2023) also experiment with the one-shot setting and propose to optimize Fisher-weightel objective to train the global model for several epochs after all clients converge
10

---

Published in Transactions on Machine Learning Research (02/2024)
independently. In contrast, our work proposes a method for general multi-round federated training from scratch, with iterative global aggregation that is computationally equivalent to one iteration per round. The function space aggregation perspective makes no assumptions about the optimality of client models and motivates application resulting algorithms to FL settings with multiple rounds. The resulting method is constrained to neither lew local steps nor full local convergence. It. is robust. local hyperparameter choices, which can otherwise be expensive and tedious to tune.
Model aggregation. Model aggregation has recently reccived attention in a number of works, most of which differ in their data and training choices during pretraining and fine-tuning, as opposed to the aggre- gation technique itself. For example, Wortsman et al. (2022) average parameters of models trained using different hyperparameters, random seeds, etc. Rame et al. (0023 build on this to reuse foundation models fino-tuned on an auxiliary task. In these works, fine-tuning started from the same model is likely to yield networks in the same loss basin for new task, thus enabling parameler-space averaging and exploiting model diversity to improve performance. Model averaging has also shown up in the empirical study of [i] ct al. (2022) showing benefits of parallel fine-tuning of large language models on diverse data over monolithic single-mode training. Similar motivations appear in Gu et al. (2023). Matena & Raffel (2022) also imple- ment a specific Fisher-weighting, but only evaluate it in the one-shot setting to merge converged or optimal models. In fact, model averaging has been in practical use for large models at least since Vaswani et al. (2017), where final models are the result of averaging previous checkpoints. We believe our motivations for model aggregation are general and their application to these varied settings exciting future work.
Additional literature relevant to the Client-Server Barrier evaluation criterion (section 4 is discussed in appendix A.6
7 Discussion and Outlook
In this work, provided function-space perspective of federated learning and proposed an aggregation technique for locally trained client models based on the input-output functions they parameterize. FEDFISH is a parametric, iterative algorithm that is robust to longer local training and client data heterogeneity,
While we have highlighted the settings where FEDFISH has advantages over FEDAVG, we now discuss its limitations and possible extensions. First,, FEDFISH derived from a second-order Taylor expansion of a function space distance. There is scope to go beyond this quadratic form for better approximations (for example, see Dhawan et al. (2023)) ) to derive parametric model aggregation objectives. Second, we make diagonal approximation to the Fisher Information matrix that effectively treats each parameter as independent. It is well-understood that deep neural network parameters are highly correlated. Better approximations to the Fisher Information matrix, such as K-FAC (Martens & Grosse 2015) or FishLeg (Garcia et al. 2023), could further boost Fisher-weighted model aggregation. However, higher dimensional approximations to the Fisher Information matrix present new challenges if applied naively, since they often increase computational burden and/or communication costs between clients and the server. Critical compatibility with FL systems in practice, FEDFISH can also be easily combined with adaptive optimization techniques as well differentially private training.
The demonstrated advantage of FEDFISH over FEDAVG across various large-scale settings, in terms of global performance, personalization, transfer to shifted distributions and a Client-Server Barrier metric, presents a compelling case for applying this aggregation technique more broadly. While our primary focus is FL at scale, a function-based model aggregation method and the barrier-based evaluation can be applied in any other setting where multiple models of the same architecture are trained with potentially different optimization algorithms, randomness, hyperparameterss data splits, etc. These use cases may allow for more flexibility in the use of data for merging, presenting new opportunities for improved function space aggregation.
Acknowledgments
The authors would like thank Daniel M. Roy and Sewoong Oh for feedback on various drafts; Sean Augenstein for helpful discussions and Keith Rush for experiment infrastructure support.
11

---

Published in Transactions on Machine Learning Research (02/2024)
References
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Represenlations, 2021. URL https:://pperreviitt http://tp...pppaarrrreeee//////a mew.://oorrrrrii--ffs88800000
Mahmoud Assran, Arda Aytekin, Hamid Reza Feyzmahdavian, Mikael Johansson, and Michael G Rabbat. Advances in asynchronous parallel and distributed optimization. Proceedings of the IEEE, 108(11):2013- 2031, 2020.
The TensorFlow Federated Authors.
TensorFlow Federated Stack Overflow dataset, 2019.
URL https:// http::///wwwtttee wttp://wwww.f....ooooo/ http::///wwoooooooooooo///o/aaa/ http:///wwpppppp...//o/o///////a /ttpa///w/aaa//o//////o////////a
stackoverflos/load data
Ari S Benjamin, David Rolnick, and Konrad Kording. Measuring and regularizing networks in function space. In ICLR, 2018.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01099 2018.
Zachary Charles and Jakub Konečný. Convergence and accuracy trade-offs in federated learning and meta- learning. In International Conference on Artificial Intelligence and Statistics, pp. 2575-2583, PMLR, 2021.
Zachary Charles and Keith Rush. Iterated vector fields and conservatism, with applications federated learning. In International Conference on Algorillmicc Learning Theory, pp. 130 147. PMLR, 2022.
Zachary Charles, Nicole Mitchell, Krishna Pillutla, Michael Reneer, and Zachary Garrett. Towards fed- erated foundation models: Scalable dataset pipelines for group-structured learning. arXiv preprint arXiv:2307.09611 2023.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. 2017 International Joint Conference on Neural Networks (IJCNN), 2017. doi: 10.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. Advances in Neural Information Processing Systems, 35:10572- 10586, 2022.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 6454-6462, 2022. Yatin Dandi, Luis Barba, and Martin Jaggi. Implicit gradient aligmment in distributed and federated learning,
Nikita Dhawan, Sicong Huang, Juhan Bac, and Roger Baker Grosse. Efficient parametric approximations of neural network function space distance. In International Conference on Machine Learning, pp. 7795 7812. PMLR, 2023.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel lcarning: an empirical study of loss landscape geometry and the time evolution of the neural langent kernel. Advances Neural Information Processing Systems, 33:5850- 5861, 2020.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Lincar mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. 3259-3269. PMIRR 2020.
12

---

Published in Transactions on Machine Learning Research (02/2024)
Jezabel R Garcia, Federica Freddi, Stathi Fotiadis, Maolin Li, Sattar Vakili, Alberto Bernacchia, and Guil- laume Hennequin. Fisher-legendre (fishleg) optimization of deep neural networks. In The Fleventh In- ternational Conference on Learning Representations, 2023. URL https://openreviers net G91ADPvQHS
htt.:pponrrreeeeeeeeeoorrrriiiaa net/forunfi=
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local sgd generalize better than sgd? arXiv preprint urXiv:2303.0122 012111 2023.
Neel Guha, Ameet Talwalkar, and Virginia Smith. One-shot federated learning. arXiv preprint
arXiv:1902. 2019.
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning as variational inference: A scalable expectation propagation approach. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://openreviest meet/forumniddd2z2r777000
Felix Hamborg, Norman Meuschke, Corinna Breitinger, and Bela Gipp. news-please: A generic news crawler and extractor. In Proceedings of the 15th International Symposium of Information Science, pp. 218 223, March 2017. doi: 10.5281 22200/........1116.....
Charlic Hou, Kiran Koshy Thekumparampil, Giulia Fanti, and Sewoong Oh. Fedchain: Chained algorithms for near-oplimal communication cost in federated learning. In Internalional Conference on Learning Representations, 2022. URL htttt::/
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Towards a theoretical and practical understanding of one-shot federated learning with Fisher information. In Federated Learning and Analylics in Practice: Algorithms,, Systems, Applications, and Opportuníties, 2023.
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 Trendss 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 on machine learning, pp. 5132-5143. PMLR, 2020.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaune Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka BarabkkaBBarwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, (14(13):3521-3 2017.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-100 (canadian institute for advanced rescarch).
URL http:://ww.....
cww...c..oooo www.oooooooooooooooooooooooo.. cetronto.edu//a www.
Kevin Kuo, Pratiksha Thaker, Mikhail Khodak, John Nguyen, Daniel Jiang, Ameet Talwalkar, and Virginia Smith. On noisy evaluation in federated hyperparameter tuning. Proceedings of Machine Learning and Systems, 5, 2023.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-233 1998. doi: 1..1109/5577679)
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zettle- moyer. Branch-train--errge Embarrassingly parallel training of expert language models. arXie preprinl arXiv:2208.0330 2022.
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, 450, 2020.
Grigory Malinovskiy, Dmitry Kovalev, Elnur Gasanov, Laurent Condat, and Peter Richtarik. From local SGD to local fixed-point methods for federated learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 97th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6692-6701. PMLR, 13-18 Jul 2020. URL press//111/mml1hookkiik00.....
mlr.
13

---

Published in Transactions on Machine Learning Research (02/2024)
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curva- ture. In International conference on machine learning, pp- 2408-220220 2417. PMLR, 2015.
Michael Matena and Colin A 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. Communiattion-effiiiie learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 2273-1282 PMLR, 2017.
Aritra Mitra, Rayana Jaafar, George J Pappas, and Hamed Hassani. Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients. Advances in Newral Informalion Processing Systems, 34:14606-1466 2021.
Reese Pathak and Martin J Wainwright. FedSplit: An algorithmic framework for fast federated optimization. Advances neurol informalion processing systems, 33:7057-7066, 2020.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified lext-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-553 2020.
Alexandre Rame, Kartik Ahuja, Jianyu Zhang, Matthicu Cord, Léon Bottou, and David Lopez-Paz. Model ratatouille: Recyeling diverse models for out-of-distribution generalization. 2023.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koncčný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deeposh Data, et al. A field guide federated optimization. arXiv preprint axXic:2107.06117, 2021.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Mor- cos, Hongsook Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Interna- tional Conference on Machine Learning, pp. 23965 23998. PMLR, 2022.
14

---

Published in Transactions on Machine Learning Research (02/2024)
A
Appendix
A.1 Notation
f
Function represented by a neural network
Network parameters
D (3-)
Function space distance
X
Input data
Target data
Z
Network outpuls
0c
Global aggregated model
N
Number of clients
D.
Dataset y:) corresponding to client i
0,
Locally trained model for client
J(())
Output Jacobian matrix
F
Fisher Information matrix
diag(F)
Diagonal of Fisher Information matrix
L;
Performance metric of interest, evaluated on data D; for client i
A.2 FedFish Derivation
We present the complete details of our derivation for FEDFISH here.
Given N client models with locally trained parameters (0,) and a function space distance as D ((,.) we defined the optimal global model to be:
N N
6 = arg min
((((0;:)),((0000)
(6)
We make a second-order Taylor approximation to the function space distance above with respect to model parameters 0, centered at cach client's 0,, the network whose outputs are to be matched. Sctting D (:.)) to the Bellback-Leiblee (KL) divergence between softmax outputs of the networks, we have for each client i,
D D(((X;;)),((;;;:))) D(((
(7)
+ (0 - 8.)7 "
(8)
8=8,
1 0,) (1) + 2
(0 (0-08 0:),
(9)
0=0,
where j(() is the output Jacobian, arising in eq. from the chain rule.
For any distance measured between the softmax outputs of the networks, (r,)), :), the first term in eq. is
0. Since the distance is minimized at f 6,, the gradient, VzD (f(X;;) f( 0.)) in eq. (8) also 0 when evaluated at. 0 0,, causing this term vanish.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Finally, applying the chain rule to the Hessian in eq.
9
yields
=
(10)
8=0;
AmB, (11)
Again, the second term in eq. 11 vanishes
0.)) = when evaluated at 0 = 0,.
Since we consider models that are trained with cross-entropy loss, natural measure of difference between outputs of two models is the KL divergence. In this case, DzLKLLff ():::)..(0 f(x;;;)) the Fisher Informa- tion matrix for outputs, Fz. Via chain rule, the first term above (J/²) P2- Fg is simply the Fisher Information matrix for the network parameters. Henceforth, we simply denote this as F or F3 to indicate the Fisher Information matrix corresponding to parameters 0, of client i.
Our final approximation function space distance reduces to
f(xx
(0 0,) F, 0.)
(12)
0.1
F (1) (g()) 0(
(13)
Here, eq- 13 makes diagonal approximation to the Fisher Information matrix, with p()) denoting the j-th diagonal entry of F,. Plugging eq.
13 into the optimization problem of eq.
gives
N 10.
0g
arg min
ΣΣ
(1)) (0
(00))
-00))
(14)
0 2N
Equation is now a convex optimization problem, which we can solve by taking its gradient with respect to and setting it to has the following closed-form solution:
Σ iiml diag( diag(F.)
0g
(15)
N diag(F)
A.3 Experimental Details
A.3.1 Datasets, Tasks & Models
We use four datasets for training models using FEDAVG and FEDFFSH:: the federated extended MNIST dataset (EMNIST) (Cohen et al. 2017), the CIFAR100 dataset (Krizhevsky et al.), the Stack Overflow dataset Authors 2019) and the C4 dataset (Raffel et al. 2020 We additionally use the CC-News (Hamborg al. 2017) and Stack Overflow datasets for evaluating transfer and post-personallzation performance of
models trained C4 using each algorithm of study.
Each of these datasets is publicly available. EMNIST is licensed under Standard Reference Data by NIST. CIFAR100 is published by the authors. Stack Overflow is licensed under the Creative Commons Attribution- ShareAlike 3.0 Unported License. C4 and CC-News are hosted by commoncrawl..rg org and we access both through HuggingFace datasets.
Table lists the scale of each dataset, the associated task and the model used for training. We include additional details dataset preprocessing and model configuration for each experiment, setting below.
16

---

Published in Transactions on Machine Learning Research (02/2024)
Table 1: Datasets, Tasks & Models
Dataset
Num Clients
Num Examples
Task
Model
Train
Test
Train
Test
CHARACTER RECOGNITION
EMNIST
3.4K
3.4K
672K
77K
CNN
IMAGE
RESNET-18
CIFAR100
500
100
50K
10K
RECOGNITION
WITH GROUPNORM
STACK
NEXT-TOKEN
350M PARAMETER
342K
204K
135.8M
16.6M
OVERFLOW
PREDICTION
DECODER-ONE ONLY TRANSFORMER
NEXT-TOKEN
1.5B PARAMETER
C4
15.6M
8.5K
364.9M
365K
PREDICTION
DECODER-ONLY TRANSFORMER
NEXT-TOKEN
1.5B PARAMETER
CC-NEWS
8.8K
708K
PREDICTION
DECODER-ONLY TRANSFORMER
The EMNIST dataset is comprised of 28x28 grey-scale pixel images of alphanumeric handwritten characters. There are characters represented. The dataset has natural heterogeneity stemming from characters being written by different authors. We partition the handwritten characters in EMNIST according to their author, as proposed by Caldas et al. (2018). We train two-layer LeNet CNN model (Lecun et al., 1998) for character recognition: two convolutional layers with 3x3 kernels and strides of length 1, pooling layer using dropout with = 0.25, dense layer with 128 units and dropout with p 0.5, and a final dense output layer.
EMNIST
CIFAR100 The CIFAR100 dataset consists of 32x32x3 pixel images with one of 100 labels. We preprocess the images using standard data augmentations, including padding to 36x36 dimenssonss, randomly cropping to 32x32, randomly flipping along the vertical axis and applying normalization. We partition CIFAR100 according to the two-level Dirichlet allocation scheme proposed by Reddi al. (2020). We train a standard R8 model with the batch normalization layers replaced with group normalization layers, following
Reddi al. (2020)
The Stack Overflow dataset is language-modeling dataset consisting of question-answe pairs from stackoverflox coml Each client corresponds to user on the platform. The data is split into train, test and validation: train client examples are from before 2018-01-01 UTC, test client examples are from after 2018-01-01 UTC, and validation clients are held out from both train and test splits. We train a 350M parameter decoder-only transformer model on the train split of Stack Overflow. We use the validation client split in our evaluations both the Stack Overflow base model and the C4 base model (for assessing transfer performance, see below).
Stack Overflow
C4 The Colossal Clean Crawled Corpus (C4) dataset is a cleaned version of Common Crawl's web crawl corpus (Raffel et al. 2020 We the federated version of this dataset presented in Charles et al. (0223
17

---

Published in Transactions on Machine Learning Research (02/2024)
Dataset
EMNIST
CIFAR100
Stack Overflow
C4
Number of clients per round
64
64
16
Number of training rounds
1500
1500
1500
10000
Batch size for local training Max client dataset size per round
10
25
4
4
100
100
16
16
Sequence length (training)
128
1024
Sequence length (personalization and evaluation)
128
128
Number of personalization epochs
4
4
Global learning rate
1e-3
1e-3
le-2
le-2
Local learning rate
1e-3
1e-3
1e-3
1e-3
Table 2: Final hyperparameter configurations for all datasets.
where each client corresponds to a different domain name (i.e., mytimes com decoder-only transformer model on the train split of federated C4.
We train a 1.5B parameter
We evaluate the C4 base model on federated version of the C4 test split, well federated CC-News and Stack Overflow to assess transfer performance. CC-News is similarly split by domain name, using Dataset Grouper (Charles et al. 2023). We further measure few-shot performance of the C4 base model by conducting a personalization evaluation on 25% held-out client data for each dataset of interest (C4, CC-News, Stack Overflow). Because of this specific personalization evaluation, we filter C4 evaluation datasets to have at least 4 examples per client.
A.3.2 Federated Algorithm Configuration and Hyperparameters
We implement the federated algorithus,, FEDAVG and FEDFISH, such that a global model is broadcast number or clients at each round, client train their models locally, model deltas are returned pseudo- gradients to the global server for aggregation, and global optimizer is used to make single update the global model using the aggregated pseudo-gradient as its own gradient.
Optimizers. We follow standard configurations used in Charles al. (2023), with stochastic gradient descent. the local optimizer and Adam the global server optimizer, across experiments.
Hyperparameter Tuning We fix hyperparameters like number of clients per round, mumber training rounds, local batch size, maximum dataset size for any client and sequence length for language models reasonable values based on previous literature. For local and global learning rates, we conducted grid search over [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, le-1], and chose the best performing hyperparameters. Final hyperparameter configurations used to obtain the results in the paper are listed in table 2 These are held consistent for FEDAVG and FEDFISH experiments.
A.3.3 Hardware Configuration
We run image-classffiattion experiments a TPC Pod slice consisting of 4 TPU v2 chips in 2x2 topology, interconnected on a single machine. Each TPU v2 chip contains two TensorCores, 16 GiB of high-bandwidth memory, and 90.75 75 GiB RAM. We run language-mode 4x4 topology, configured to use a multi-machine inter-chip interconnect mesh. Each TPU v3 chip contains two TensorCores, 32 GiB of high-bandwidth memory, 87.25 GiB RAM, 900 GBps bandwidth, and 123 teraflops peak compute.
experiments on a TPU Pod slice consisting of
TPU v3 chips
18

---

Published in Transactions on Machine Learning Research (02/2024)
Global Accuracy
Personalized Accuracy
Dataset
Local epochs
Method
Fixed compute
Fixed rounds
Fixed compute
Fixed rounds
Avd
81.38
81.38
83.09
83.09
4
FedFish
82.64
82.64
83.44
83.44
FedAvg
82.05
82.24
83.46
83.82
EMNIST
FedFish
83.19
84.42
84.37
85.6
FedAvg
76.45
77.52
77.96
79.63
16
FedFish
79.01
82.86
81.05
84.96
Fed
35.75
35.75
37.06
37.06
FedFish 36.88
36.88
41.31
41.31
CIFAR100
FedAvg 28.56
33.28
36.00
32.13
16
FedFish
31.63
37.72
40.53
44.69
Table 3: Longer local training improves overall performance in terms of both global model accuracy as well as personalization. When amount of compute, i.e. total number of training epochs, is fixed, FEDFISH can achieve better performance than FEDAVG with fewer rounds of communiattin between the server and clients.
(9)
Peroralizatoo Thannng Data (S)
Peronalizato Training Data
Feddog FolFFils bozi cauch loralepoch
Perplesit
Perplesity
Paplesity
Figure 7: Transfer performance on perplexity (lower is better) after federated pretraining on C4 and evalu- ating on C4 (left), Stack Overflow (center) and CC-News (right...
A.4 Additional Results
Here, we include additional results and report the performance visualized in section Table 3 shows results image benchmarks where performance is computed by keeping either the computational requirement the total number of rounds fixed. For the former, we limit the number of training rounds according to the number of local epochs so that total amount of client-side training remains constant. Unsurprisingly allowing for more rounds of training helps performance. However, at fixed compute, we see that. FEDFFIS is able to achieve better performance than FEDAVG with fewer rounds of communication. Table 4 and table S inchude full results on our language modeling experiments, including different numbers of local training epochs, different datasets used for federated pretraining or personalization different, performance metrics and different amounts of data used for personalization. Finally, similar to fig. with shows next-token prediction performance, present perplexity performance for the C4 transfer experiments in appendix A.3.3
A.5 Efficiency of FedFish
We recognize that there are opportunities improve the efficiency of FEDFISH, both in terms of computation and communication. As presented in algorithm each client takes additional pass its data to estimate
19

---

Published in Transactions on Machine Learning Research (02/2024)
Training
Training Local
Persoualization
Global
Personalized (25%)
Personalized (50%)
Method
Epochs
Token
Perplexity
Perplexity
Perplexity
FedAvg
FodFish
32.98
3.77
34.25
Stack Overflow
Stack Overtlow
FedAvg
21.77
5.56
25.46
1.61
FedFFish
31.27
3.96
32,60
3.81
FedAvg
30. 16
4.26
30.81
4.21
30.95
4.20
FedFish
31.15
4.19
32.10
4.12
32.34
4.11
C4
Avd
31.04
32.73
FedFish
Fed
17.91
5.29
19.85
5.13
21.58
FodFish
19.55
5.101
22.07
4.91
23.38
4,815
Stack Overflow
FedAvg
18.83
5.13
19.63
5,201
22.98
4.69
FedFish
19.54
5.16
26.80
4.32
21.77
4.22
FedAvg 29:03
1.38
29.45
4.31
29.73
4.31
FodFish
29.53
30.61
C4
CC-News
Aod
4.25
16
FedFish
31.00
1.15
31.75
4.06
32.58
3.98
Table 4: Full results on global and post-personalization performance of language models that are pretrained in a federated manner with varying number of local epochs and then evaluated on different datasets. These results correspond to the entire rounds federated training.
Training
Training Local
Porsonalization
Global
Personalized (25%)
Personalized (50%)
Method
Dataset
Epochs
Dataset
Token Pred
Perplexity
Token Pred
Perplexity
Token Pred
Perplexity
FedAvg
1.61
27.
1.60
FedFish
4.54
C4
Fvd
4.30
FedFish
30.33
31.54
4.11
31.61
4.10
FedAvg
17.10
5.38
19.92
5.20
21.78
5.10
FedFish
17.25
5.41
19.62
5.2
21.10
5.08
FedAvg
18.57
5.29
16.23
7.03
18.91
6.09
FedFish
17.80
5.42
20.81
4.77
22. 78
4.58
Fod
26.03
4.79
26.44
4.75
26.72
4.72
FedFish
26.23
4.86
26.98
1.50
27.31
1.76
CC-News
FedAvg
27.55
4.58
27.23
4.59
28.02
4.55
16
FedFish
28.91
4.44
29.86
4.29
30.06
4.27
Table 5: Results global and post-personalizztion performance of language models in different settings, evaluated at R/2 rounds of federated training.
the Fisher diagonal (lines 9-12) and communicates this estimate along with its model update (line 13). This results in an increase in computational cost equivalent to one epoch of training and two-fold increase in communication cost.
A.5.1 Computational Overhead of FedFish
To climinate the need for an additional forward and backward pass through the model compute Fisher information, we can instead use the gradients from the last local epoch of training. This is a further approximation as the gradients with respect to the evolving client model rather than with respect to
20

---

Published in Transactions on Machine Learning Research (02/2024)
Proodiitattin To Traning Duaa
Pessonaitatii Traning
cpoch locallppcccs
loca spocts
Prediction Accunac
Perplesity
Figure 8: Global and post- perssonooaaiintoonn performance in terms of next-token prediction (left) and perplexity (lower is better) (right) after federated pretraining C4 using FEDAVG FEDFISH where the Fisher is estimated using gradients from epoch E.
Training
Training Local Epochs
Personalization
Global
Personalized (50%)
Method
Dataset
Dataset
Token Pred
Perplexity
Token Pred
Perplexity
FedAvg
30.16
4.26
30.95
4.20
FedFish
32.10
4.10
34.86
C4
C4
FedAvg
31.04
4.13
32.72
4.08
16
FedFish
30.28
4.07
44.31
3.42
Table 6: Results from ablation study of FEDFISH without computational overhead, where the Fisher is estimated using gradients from epoch E. Global and post-personaliatton performance of model pretrained and personalized on C4 evaluated at R rounds of federated training.
the fully updated client model that will be merged. The computational overhead FEDFISH, as presented algorithm 2 likely to be more of hindrance as models scale. Given this efficiency concern most relevant for the largest setting consider, we perform an ablation study C4 to investigate the effect of making the approximation described above.
We run FEDFISH C4 experiments using gradients from the last epoch of local training in cach round to com- pute the Fisher diagonal. Results are presented in fig. g and table We find that this further approximation does not reduce performance even when only training with a single pass over the data. Surprisingly, we see a substantial improvement in personalization performance for training with 16 local epochs when using gradients from the 16th epoch rather than from an additional pass after fixing the local model.
A.5.2 Communication Overhead of FedFish
Note that overall communication cost is not just measured in bits communicated per round, but also in the total number of communication rounds required. Depending on the network constraints of these factors may be more critical to minimize. As discussed in section 5.3 the presented implementation of FEDFIIH (see algorithm and algorithm 2 has twice the communication cost (per federated round) of FvDAVG. This overhead can be eliminated by combining our method with adaptive federated optimization techniques Reddi et al. 2020). having clients only send their weighted parameters and folding the normalization into the server optimization. We leave this extension to future work, and provide more critical look at the advantage of FEDFISH given the presented implementation.
Ablation: Communication Rounds Figure 9 depicts the evaluation performance of FEDAVG and FED- FISH on EMNIST across local epochs normalized by the number of communication rounds. FEDFISH con- verges faster and to higher global accuracy and post-personalization accuracy than FEDAVG. Despite the additional cost per round of FEDFISH, convergence over fewer communication rounds useful for setting in which the network is not bandwidth--onstranned but clients are intermittently reachable.
21

---

Published in Transactions on Machine Learning Research (02/2024)
Global Arcuras
Puzsonlization
FodAng
FedAvg
FedFish
FedFish
local epochs
local spochs
local epochs
local spocks
16 lood spodla
local opochs
ab Coumunication Rounds
sb
de
Communication Rounds
Figure 9: When normalizing by communication rounds, training on EMNIST with FEDFISH converges faster and to higher global accuracy (left) and post-personalization accuracy (right) than with FEDAVG. Results are shown for experiments across varying numbers of local epochs. Each marker indicates 100 federated communication rounds.
Prat Personalization Acmuacc
0.82
03
0.75
Globe] Accurocy
37*
0.79
0.8
0.85
FedFish
FedAsg
0.60
TedAvg
4 lood epochs epochs
4 lood spoals
0.88
0.55
epocta
16 local epochs
local spochs
0.5
oin
700
033
100
250
300
400
s00
ebo
700
Communniation Cost parameters sent)
Communicatio
((0) parameters sent)
Figure 10: When normalizing by communication cost, training on EMNIST with FEDFISH matches the performance FVDAVG both in terms of global accuracy (left) and post- post-personalization accuracy (right), with half the number of communication rounds. Results are shown for experiments across varying numbers of local epochs. Each marker indicates 100 foderatod communication rounds.
Ablation: Communication Cost Figure 10 depicts the evaluation performance of FEDAVG and FED- FISH on EMNIST across local epochs normalized by the total communication cost (measured in units of 0). FEDFISH performs comparably to FEDAvG both global accuracy and post-personalization accuracy. Taking the additional cost. per round of FEDFISH into account, FEDFISH no more costly than FEDAVG measured in terms of bits communicated accuracy achieved.
A.6 Additional Related Works
Linear Mode Connectivity. Frankle et al. (2020) The Client-Server Barrier is largely inspired by the error barrier definition in Fort el. al. (2020). There the authors define an error barrier as the maximum increase in error on a linear path in the parameter space between two models. There are few key differences between the error barrier and client-server barrier: they consider models trained on the same training data; using our notation, the losses L; are identical for all indexes over different models 0, (c.g., models trained in the same centralized way but independently). The error barrier corresponds then to the maximum over
22

---

Published in Transactions on Machine Learning Research (02/2024)
{wo IE...... where W, is the weight corresponding to 0 when linearly interpolating between {0s h= .N. Prior to Fort et al. (2020), error barriers appeared under "instability" definition in Frankle et al. (2020) for evaluating how different trajectories of pair of networks connect in the loss landscape. Later, similar metrics were used other model averaging work, such as Wortsman et al. (2022). where the authors consider multiple networks trained with different hyperparameters.
23


Reviews:
[{"review_id": "kPfDDUFoAI", "reviewer": "Reviewer_yAgt", "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\n- The paper is well written.\n- The idea of aggregating models in the function space is well-motivated and novel.\n- The paper contains well-designed experiments for both toy and real datasets to validate the claimed improvements, such as the change in the improvements for varying number of local steps.\n\n# Weaknesses\n\n- The paper lacks any theoretical analysis or guarantees.\n- The ``function space\" metric defined by the KL-divergence of the softmax outputs requires additional motivation. It should further be emphasized that the metric is limited to imposing similarity of the function outputs on the training data, not on the entire input space.\n- The proposed approach isn't compared to algorithms such as SCAFFOLD (Karimireddy et al., 2020) and FedGA (Dandi et al., 2022) which also improve over FedAVG in the presence of client heterogeneity.", "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": "The paper introduces a novel algorithm FedFish for federated learning based on function space aggregation. In contrast to the standard approach of FedAvg that updates the global server model by minimizing the average distance between the parameters (which corresponds to taking an average of the client parameters), FedFish aims to minimize the average distance in function space betweeen the output functions. This is especially relevant for settings with high heterogeneity across clients since the FedAvg approach suffers from divergence of the model parametes as the number of local steps is increased.\n\nFedFish considers the function space metric defined by the KL-divergence of the softmax outputs, averaged over training data. It then approximates this objective through a quadratic approximation defined by the Fisher Information matrix.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Missing references that should be included for securing recommendation:\n\n- A number of references related to the alleviation of ``client drift\" are missing, such as:\n\n     Dandi, Yatin, Luis Barba, and Martin Jaggi. \"Implicit gradient alignment in distributed and federated learning.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 6. 2022.\n\n## Suggested changes that would strengthen the work:\n\n- As an ablation study, it could be useful to include experiments utilizing the full Fisher-information matrix or other function space metrics on toy datasets.", "broader_impact_concerns": "The paper does not posit any broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vjMhWNFuG7", "reviewer": "Reviewer_joh1", "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. The investigated topic is very interesting based on the function space aggregation.\n2. The paper is well written and easy to follow. The presentation is clear.\n3. The empirical studies are extensive and comprehensive.\n\nWeaknesses:\n1. The idea of function space aggregation using Fisher information matrix is existing. Please see [1]. Hence, the novelty in this work seems low. Though in this work the setting is federated learning. The optimal solution $\\theta^*$ is almost exactly the same.\n2. The theoretical contributions in this paper are not strong. If the authors could show some theoretical analysis of how FedFish enables better robustness, that would make the paper more technically solid and sound.\n3. From the experimental results, we know that FedFish performs better if the local training is longer. But the authors only presented the results, failing to discuss why this was the case. Also, for other results, the authors only described the outcomes without clearly discussing the reasons.", "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 proposes a novel federated learning algorithm based on the function space. Instead of taking simple weighted average in the traditional FedAvg method, the proposed algorithm, FedFish,  aggregates local approximations to the functions learned by clients, using an estimate based on their Fisher information. To validate their scheme, the authors utilize several models and datasets for the detailed performance comparison. The results show that FedFish is more robust to longer local training. Also, FedFish benefits the efficient personalization via local fine-tuning on the same or shifted data distributions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please add theoretical analysis for the FedFish.\n2. Please add technical discussion for the empirical results to explain why FedFish outperforms FedAvg.\n3. It would be great to see the results on a larger image dataset, such as ImageNet.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "i3YVRGjFX9", "reviewer": "Reviewer_uX9A", "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. Proposed a new federated learning algorithm referred to as FedFish.\n2. Formalizes a function space perspective of federated learning to motivate FedFish.\n3. FedFish is robust to longer local training and client data heterogeneity.\n4. In the extreme case of completely disjoint supports, FedFish matches the locally learned functions of the clients on their respective\ninput data. In this case, FedFish performs significantly better than FedAvg. \n5. FedFish results in global networks that are more amenable to efficient personalization via local fine-tuning on the same or shifted data distributions.\n6. Experimental results across several settings in image and language benchmarks are presented. \n7. The paper is well-written and easy to follow.\n\nWeaknesses/Questions:\n1. FedFish is a trivial extension of fisher-weighted averaging proposed in [1] to federated learning setups.\n2. FedFish requires $2\\times$ communication (from client to server) as compared to FedAvg in each communication round.\n3. Needs an additional forward and backward pass through the model to compute fisher information. Can we just use the last local iteration's gradients (Epoch E's gradient in Algorithm 2) to compute fisher information?\n4. Theoretical analysis for the convergence rate of FedFish is missing. \n5. In Fig. 3a, the curve for FedAvg with 4 local epochs overlaps with FedFish with 8 local epochs and the curve for FedAvg with 8 local epochs overlaps with FedFish with 16 local epochs. Is it possible that the FedFish converges at the same rate as FedAvg for iso-communication comparison?\n6. The squared gradient is also used as a proxy for sharpness estimate. It would be interesting to correlate FedFish's aggregation mechanism to the sharpness or flatness of an individual client's loss surface. Is this aggregation somehow helping us to achieve a flatter loss region for the global model and hence allowing for longer local training? \n\n[1] Matena, Michael S., and Colin A. Raffel. \"Merging models with fisher-weighted averaging.\" Advances in Neural Information Processing Systems 35 (2022): 17703-17716.", "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": "The paper provides a function-space perspective of federated learning and proposes the FedFish algorithm that uses fisher-weighted averaging for aggregation of the local models. In particular, it replaces the weighted averaging set of FedAvg with fisher-weighted averaging. FedFish is robust to longer local training and client data heterogeneity. Authors demonstrate the advantage of FedFish over FedAvg across various large-scale settings, in terms of global performance, personalization, and transfer to shifted distributions, and a Client-Server Barrier metric presents a compelling case for applying this aggregation technique more broadly.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Plotting Fig. 3 results with respect to communication bits would help in a better understanding of iso-communication comparison.\n2. Comparison with Federated Learning baselines designed for heterogeneous data such as FedProx [1], ScaffNew[2], etc can strengthen the argument that FedFish is more robust to client data heterogeneity.\n\n[1] Li, Tian, et al. \"Federated optimization in heterogeneous networks.\" Proceedings of Machine learning and systems 2 (2020): 429-450.\n\n[2] Mishchenko, Konstantin, et al. \"Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally!.\" International Conference on Machine Learning. PMLR, 2022.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Strengths:
- Introduces a novel function-space aggregation method for federated learning (FedFish).
- Well-motivated and novel idea of aggregating models in function space.
- Extensive and comprehensive empirical studies demonstrating improvement over FedAvg.
- Clear and well-written paper with easy-to-follow presentation.
- FedFish is robust to longer local training and client data heterogeneity.
- Outperforms FedAvg in extreme cases like disjoint client data supports.
- Enables efficient personalization via local fine-tuning.

Weaknesses:
- Lack of theoretical analysis or guarantees for FedFish.
- The "Function space" metric and its limitations need further explanation.
- Potential lack of novelty due to similarities with existing work (e.g., [1]).
- Higher communication cost compared to FedAvg due to Fisher information calculation.
The possible relationship between FedFish's aggregation and loss surface sharpness needs exploration.
- Convergence rate comparison with FedAvg is missing.

Overall: The reviewers offer mixed opinions on the FedFish paper. While they acknowledge its strengths in novelty, empirical performance, and clear presentation, they also raise concerns about theoretical foundations, potential lack of novelty, and higher communication costs. Further work on theoretical analysis, addressing limitations of the metric, and comparisons with existing algorithms could strengthen the paper's contribution.

- Reviewers 1 and 3 suggest comparing FedFish to other relevant algorithms like SCAFFOLD.
- Reviewer 2 recommends discussing the reasons behind observed performance improvements.

The authors have already completed some of the requests of the reviewers; we are asking to satisfy as many as possible of these requests that do not violate the acceptance criteria of TMLR.
Decision: Accept (poster)