FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Authors: Jacopo Talpini
Marco Savi
Giovanni Neglia
Venue: ICML 2025 poster (2025)
Keywords: One-shot Federated Learning
Bayesian Inference
Laplace Approximation
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: This paper proposed a novel approach to the One-Shot Federated Learning, based on Bayesian Inference

Abstract: One-Shot Federated Learning (FL) is a recent paradigm that enables multiple clients to cooperatively learn a global model in a single round of communication with a central server. In this paper, we analyze the One-Shot FL problem through the lens of Bayesian inference and propose FedBEns, an algorithm that leverages the inherent multimodality of local loss functions to find better global models.Our algorithm leverages a mixture of Laplace approximations for the clients' local posteriors, which the server then aggregates to infer the global model. We conduct extensive experiments on various datasets, demonstrating that the proposed method outperforms competing baselines that typically rely on unimodal approximations of the local losses.

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Jacopo Talpini' Marco Savi!
Giovanni Neglia?
Abstract
communication between clients and the server (Guha et al., 2019).
One-Shot Federated Learning (FL) a recent paradigm that enables multiple clients to cooper- atively learn global model in single round of communication with a central server. In this paper, we analyze the One-Shot FL problem through the lens of Bayesian inference and propose FedB an algorithm that leverages the inherent multi- modality of local loss functions to find better global models. Our algorithm leverages mix- ture of Laplace approximations for the clients' local posteriors, which the server then aggregates to infer the global model. We conduct extensive experiments on various datasets, demonstrating that the proposed method outperforms competing baselines that typically rely on unimodal approxi- mations of the local losses.
Several One-Shot FL algorithms have been proposed in the literature. Existing relevant work leverages knowledge dis- tillation at the server (Lin et al., 2020), neuron matching strategies (Singh & Jaggi, 2020) or adopts an optimization approach, trying to directly approximate the global loss at the server starting from the local losses of each client (Jhun- jhunwala et al., 2024; Liu et al., 2024; Matena & Raffel, 2022). Our contribution is in line with the last group of work, which generally employs a unimodal approximation of cach local loss. As an example, Jhunjhunwala et al. (2024) make use of a quadratic approximation of local losses through opportune Fisher information matrices.
However, neural network loss functions are known to be highly non-convex (Izmailov et al., 2021; Wilson & Iz- mailov, 2020; Fort et al., 2019; Lakshminnrayanan et al., 2017) with multiple local minima, some of which may gener- alize better than others. This fact can be extremely relevant for the considered One-Shot FL. setting and motivated us to introduce a novel algorithm that explicitly takes into ac- count the different modes of the loss surface, guided by Bayesian analysis of the FL problem.
1. Introduction
Federated Learning (FL) is a collaborative machine learning paradigm that has recently gained significant attention from the research community. It enables the training of models us- ing decentralized datasets stored locally on multiple clients, without requiring direct access to the raw data (McMahan al., 2017; Kairouz et al., 2021). This paradigm is particu- larly appealing in scenarios involving highly sensitive data, where clients are reluctant to share them, due to privacy or security concerns. Most FL algorithms typically train models using iterative optimization techniques that require numerous rounds of communication between the clients and the central server that aggregates the local models.
Figure shows the importance of considering multiple modes, motivating our proposed approach. The figure illus- trates that the local minimum corresponding to the broader, less sharply-peaken mode of client I's local loss, cru- cial for accurately reconstructing the optimal global model. Moreover, it is worth noting that this represents a relatively simple scenario where the global loss is unimodal. In gen- eral, however, the global loss exhibits complex landscape with multiple optima, and selecting just one of them as global model may lead to suboptimal results (Wilson & Iz- mailov, 2020). We argue that this consideration is especially crucial for learning model in a One-Shot setting while achieving good predictive performance, as we will show in the rest of the paper. The key insight, well-established in Bayesian inference (MacKay, 2003), is that describing the whole posterior distribution should be the main goal of learning from data, while a local description around the posterior's maximum may not be sufficiently informative.
The frequent server-client interactions represent bottle- neck in FL (Guha et al., 2019; Kairouz et al., 2021), which has been addressed by new paradigm called One-Shot FL., where the global model is learned in a single round of
'Department of Informatics, Systems and Communication (DISCo). University of Milano-Bicocea, Milan, Italy: Inria, Uni- versité Côte d' Azur, Sophia Antipolis, France. Correspondence to: Jacopo Talpini <j.talpini @ campus. campus.unimib.le unimib.it>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
Contributions. In this paper. we introduce FndBEns (Feder- ated Bayesian Ensemble), straightforward and principled

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Local loss functions
Global loss function
Our recanstructed optimum
Client
2-
-2-
-2
-4
Local Optima
Reconstructed optimum from unimodal local approximations
-2
2
Figure 1. An illustration of One-Shot FL in toy 2D setting with two clients: with a unimodal loss function (Client 2, left plot) and the other with a bimodal loss function (Client 1, left plot), along with the overall 'ground-truth' global loss (right plot). The left plot also shows several stochastic gradient descent (SGD) trajectories on the two losses (dotted lines), mimic clients' training, starting from different random initializations. In the right plot, the red dots represent the reconstructed global optimum, depending on whether client I approximates local loss as quadratic around its global optimum (a more likely result, represented by the bigger dot) or around the secondary minimum. The yellow star denotes the global optimum inferred by our approach, based on an ensemble of all the optimal solutions found by each SGD run. To calculate the global loss, the secondary minimum of Client more relevant than its global minimum.
one-shot federated learning method based on Bayesian anal- ysis of the distributed inference problem. We derive exact expression combine clients' local model posteriors, yielding the same global posterior one could compute from the union of all local datasets. Guided by this result, we propose a practical algorithm based on a multimodal ap- proximation of the local loss functions- and, consequently, the global loss function constructed using a mixture of local Laplace approximations. We subsequently exploit this multimodal global posterior to approximate Bayesian marginalization, rather than relying solely on its optimum, to derive the global model. Finally, we empirically evaluate the performance of our method on standard benchmarks and we compare with several baselines and state-of-the-art approaches. Our illustrative numerical results demonstrate that our proposal outperforms other state-of-the-an models across different scenarios.
2. Related Work
In this Section, we begin by reviewing seminal work on Bayesian approaches for Federated Learning, followed by an in in-depth discussion of One-Shot FL.
2.1. Bayesian approaches to Federated Learning
Federated Learning is typically framed as distributed op- timization problem characterized by unique challenges, in- cluding unbalanced and non-i.i.d. data distributions across clients and constraints communication. A seminal con- tribution that integrates Bayesian principles and FL is the FedPA algorithm (Al-Shedivat et al., 2021). This algorithm requires multiple communications rounds and it is based on Markov Chain Monte Carlo for approximate inference of local posteriors on the clients. It efficiently communi- cates their statistics to the server, which uses them to refine global estimate of the posterior mode. A further con- tribution is FedEP (Guo et al., 2023), which exploits an expectation propagation strategy, where approximations to the global posterior are iteratively refined through proba- bilistic message-passing between the central server and the clients. For a systematic review of FL with Bayesian models see (Cao et al., 2023).
The remainder of the paper is organized as follows. In Sec- tion 2 we revise several relevant related works, while in Section we analyze the One-Shot FL problem from a prob- abilistic perspective and we describe the proposed approach in Section 4. In Section 5 we describe the experimental setup, the benchmark approaches from the literature, and the considered datasets. In Section 6 we provide and discuss the numerical results, and conclude the paper by highlight- ing the main takeaways and lessons learned in Section 7,
Although the theoretical foundations of previous works are rooted in Bayesian theory, their primary goal is to infer the maximum of the posterior, achieved through multiple rounds of communication. In contrast, our focus is on combining models in a single round of communication and relying on Bayesian marginalization.
2

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
2.2. One-Shot Federated Learning
the local loss function and evaluating its eurvature around the minimum. This information is then transmitted to the server. The server reconstructs quadratic approximations of the local loss functions from the received information and employs stochastic gradient descent (SGD) to minimize the sum of these approximations. A regularization term is included the server's loss function to enhance stability by preventing SGD from drifting too far from the mean of the local model parameters. Our approach aligns with this last group of works, but it differs by providing multimodal approximation of the loss functions and by relying on the marginalization of the posterior rather than its optimiza- tion, thus embracing a full-Bayesiar approach (Wilson & Izmailov, 2020).
Existing work for One-Shot FL can be broadly split into three main categories: (i) knowledge distillation methods, (ïi) neuron matching methods and (iii) loss-based models fusion.
The methods falling in the first group view the client models as an ensemble and the overall idea is to distill the knowl- edge from this ensemble into a single global model. To accomplish that, some works assume the existence of dataset at the server (Gong et al., 2021; Lin al., 2020; Guha et al., 2019) while others mimic this scenario by train- ing generative models, like GAN or VAE, to synthesize an auxiliary dataset at the server side (Zhang et al., 2020; Zhou et al., 2020; Heinbaugh et al., 2023). In another recent paper, Hasan et al. (2024) propose an approach based on a combination of predictive mixture models and Bayesian committee machine. The contribution of each component is weighted by parameter, which is determined at the server by minimizing the log-likelihood using a server-side dataset, also used to distill the committee's 's knowledge into a sin- gle model. The availability of a dataset (whether real or synthetic) raises several concerns. If each client shares a portion of its data, privacy issues arise. Alternatively, a datasets obtained through other means may not accurately represent the clients' data. Additionally, using such a dataset increases the computation cost on the server and requires careful hyperparameter tuning, posing a significang chal- lenge for implementation in One-Shot FL (Jhunjhunwala et al., 2024).
3. Problem Formulation
The One-Shot FI. problem can be formulated as follows: given collection of C' clients, where each client has its own dataset of input-outputs pairs D. = yx) NN the goal is to learn a parametric function (e.g.. a Neural Net- work) that provides the conditional probability distribution p(y) Xx w). for a given input x and model parameters w (Murphy, 2023), in a distributed way. Specifically, we aim to combine the local models trained on each D. such that the resulting global model is equivalent to the one trained on DD = U,D, Local models are trained by minimizing a specified loss function, which typically corresponds to the negative log-likelihood induced by the considered prob- abilistic model p(y|x, w) (Al-Shedivat et al., 2021). For example, the commonly used least-squares loss for regres- sion can be derived from the likelihood under a Gaussian model, while the cross-entropy loss corresponds to the like- lihood under a categorical model. Overall, the local learning process at each client is framed as the minimization of a given loss function Le(w).
On the other hand, methods based on neuron matching rely on the invariance of Neural Networks (NN) w.r.t. permuta- tions of the neurons. Some works leverage this property, by first aligning the weights of the client models according to common ordering (called matching) and then averaging the aligned client models (c.g. (Singh & Jaggi, 2020; Liu et al., 2022: Jordan et al., 2023)). Although this approach has been shown to work well when combining simple models like feedforward NNs, its performance declines when combin- ing more complex ones such as CNNs (Jhunjhunwala et al., 2024).
3.1. Combining models: a probabilistic perspective
A principled way to combine the local models is through the likelihood p(D) w) = p(D,,.,, .., De (w). Assuming that the clients' data are independent given the model, the global
likelihood can be factorized into product of local likeli- hoods as follows:
The last group of works, referred to as model fusion, intends to fuse the capabilities of multiple existing models into a single model (Yadav et al., 2023; Jin et al., 2023; Matena & Raffel, 2022). These approaches were not specifically designed for one-shot federated learning. As result, they overlook the fact that the models to be aggregated may be trained on datasets drawn from different distributions, as is common in federated learning. A relevant contribution to the field, more aligned with our work, is (Jhunjhunwala et al., 2024). Here, the authors propose a One-Shot FL approach based on the Fisher information matrix. Each client com- putes a local estimate of the model weights by minimizing
II p(D./w)
(1)
To obtain the global posterior, instead of combining the likelihoods, it is possible to combine directly the client's posteriors, specified by the following Proposition:
Proposition 3.1. Given C datasets and the posteriors (Dw) of the same model f(w), under the assumptions that (i) the datasets are conditionally independent given
3

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
the model, and (ii) the prior is the same across clients, the global posterior can be written as:
this choice lies in combining lightweight approach to de- scribe the local curvature of the loss (provided by the usual Laplace approximation) with global perspective provided by the different ensembles (Eschenhagen et al., 2021). More concretely, each client trains M times the same model on the same data but randomly varying the starting point of the gradient descent algorithm, exploited for the training. This approach advantageous as it does not require any modifi- cations to the model architecture, the standard loss function, or the learning process. By adopting this method, it follows that the overall posterior at the server will be mixture of Gaussian distributions and, in general, multimodal. More precisely, starting from Equation (2), the posterior for C clients and M mixtures can then be written as:
p( (D. Dc) = a p(w)C-
The first term of the right-hand side is normalizing con- stant and does not depend explicitly on the model param- eters, while the factor p( ((W)C1 avoids the over-counting of the prior. A proof is given in the Appendix. We em- phasize that this result allows for the exact combination of distributed inference outcomes represented by the client posteriors (D))))/__ without any loss of information. In other words, combining the posteriors via Equation (2) provides the same information about the model as if it had been trained in the usual centralized setting, by combining directly the clients' data, thus naturally enabling Federated Learning in One-Shot.
p(w www/Wil I:M
I:C
11MM )
x
C M
II>
M M
N(W my A%)
(4)
p(w)
If we are interested in finding the maximum of the global posterior, i.e., to perform a Maximum Posteriori (MAP) estimate of the parameters for the global model, we can rely on the 1eg-oostrror and drop the terms that do not depend on w, resulting in the following objective function:
where: w is the m-th maximum-a-posteriori estimate of the c-th client, found through standard gradient-based op- timizers, and similarly A is the precision matrix. The precision matrix involved the LA is derived from the Hessian of the negative log-likelihood evaluated at w To ensure positive definiteness, it is typically approximated using the generalized Gauss-Newton method (Daxberger et al., 2021a). The previous expression is used by the server as a proxy for the global posterior. Specifically, to identify different modes of this posterior, used in the second phase of our approach, the server computes the po-ppostrrio,, thus obtaining the following objective function:
[(w)= Σ log(p(w(D.)) (C- ()g((((())
Remark. It is worth mentioning that under the assumption of an (improper) flat prior the obtained result reduces to what has been presented in (Al-Shedivat et al., 2021), and considered as starting point in (Jhunjhunwala et al., 2024). While those works focus solely on inferring the maximum of the posterior, providing a unimodal approximation of the local loss, we propose a multimodal method that properly accounts for the prior by capturing its different modes.
M
[(w/WIC
LC ALCC I.M
log Σ
1 M
1
N(ww m. ((.
c=1
m=1
(C 1)log(p(w))
(5)
Remark. In the previous equations, we implicitly assumed that each mixture component has equal weight. This as- sumption aligns with (Eschenhagen al., 2021), as we do not expect any particular model to be more important than the others, given that the only difference between the mixture components is their random initialization.
4. Proposed Algorithm: FedBEns
In this Section we discuss in detail our approach, which operates in two stages. In the first stage, the global posterior estimated, followed by a second phase where its different modes are identified. The details of our method are given in the pseudocode presented in Algorithm 1.
4.2. Model predictions
4.1. Global posterior estimation
Given the posterior distribution, fully-Bayesian approach makes predictions by relying on marginalization over the posterior of the model parameters, rather then using a sin- gle setting of weights (Wilson & Izmailov, 2020). More specifically, by marginalizing, one computes the predictive distribution for a given input x as (Wilson & Izmailov, 2020; MacKay, 1995);
As mentioned, we first aim to construct a tractable and ef- fective approximation of the complex local posteriors that can capture different modes. To achieve this, we employ mixture of Gaussian distributions to approximate the local posterior distribution of each client, (w))D.). More specifi- cally, we exploit a mixture of Laplace approximations (LA) (MacKay, 1992) of independently-trrinn works (Eschenhagen et al., 2021). The rationale behind
Deep Neural Net-
ply|x, D) - / w)(w\/Dd
(6)

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
As mentioned earlier, the posterior is generally multimodal, indicating that it is possible find different parameter set- tings that can fit the data well. Therefore, it is advantageous to consider a mixture of these solutions and use all of them to make predictions, in line with a Bayesian Model Aver- aging perspective, represented by Equation (6) (Wilson & Izmailov, 2020).
posterior (Lakshminarayanan et al., 2017), initializing from the median of the different MAP solutions obtained from the clients' mixtures. This ensemble represents the final model that the server sends back to the clients.
4.3. Implementation considerations
The Laplace approximation involves computing the Hes- sian at the MAP estimate. However, constructing the full Hessian matrix is quadratic in the number of parameters, making it impractical for modern Neural Networks. Typ- ical approaches for applying the Laplace approximation include diagonal or Kronecker factorization of the Hessian (Grosse & Martens, 2016; Ritter al., 2018), or applying the Laplace approximation to subset of the model (e.g., to the last layers) (Daxberger et al., 2021a). Here, we consider two approaches for approximating the Hessian:
Algorithm 1 FedBEns: One-Shot FL through Federated Bayesian Ensemble
Require: Data DI:c (data of each client); number of mix- tures M; client and server learning rates Te, 7s: and number of iterations Ker K,. respectively.
1: Server randomly initializes M model weights wi.M and broadcast them to C' selected clients
2: for each client e in C' in parallel do
3: [cc)
WI::: A ":M < C1 i ng( I:M: Der K2)
(ee)
Diagonal+Full. This approach approximates the Hessian simply with its diagonal elements, with the exception of the last layer for which it exploits the full Hessian. This resem- bles the subnetwork LA (Daxberger et al., 2021 Kristiadi et al., 2020) that treats only subset of the model parameters probabilistically, typically focusing on the last linear layer while ignoring the uncertainty associated with the feature extractor layers. Last-layer Laplace approximations proved to be effective in practice, the usual centralized setting (Daxberger al., 2021a).
4: Client e sends updated mixture weights and precision
matrices w (e) A 1:M to the server
5: end for
6: Server aggregates client updates to infer global optima:
7: "1:M
w-olohal + i onw I.M AE&
1:C"
8: Clients receive global and make predictions by ensem- bling:
9:
pluix, Dr.)) EmEr(by p(y|x, w
10: Function:
i.M, De, Ko)
11: for each mixture model m do
Kronecker. More expressive alternatives to the simple diag- onal approximation are block-diagonal factorizations, such as Kronecker-factoree factored approximate curvature (Martens & Grosse, 2015), which factorizes the Hessian of each layer as a Kronecker product of two smaller matrices. The di- mensions of these matrices are determined by the number of input-output neurons. Let di and d2 denote the number of input and output neurons, respectively. Consequently, the computational cost for storing and communicating the Hessian scales as Ode + dZ), instead of ((d) )2
12: for k = 0..... Ke - do
13:
w t w nc E, V((w mi X yi)
14: end for
15: Set Wm w K.-1
16: Compute precision matrix of Laplace approx. Am
17: end for
18: return WI,M, AlArr
19: Function: ServerAsgregattion I:C
l:y A I.C
1-M- K.)
20: Compute global loss [(w|w]:C 1:M: (:C) via Eq. (5)
21: Compute starting points: wi.s = median, WIC
Temperature Scaling. In Bayesian deep learning it is com- mon to consider a tempered posterior by raising the likeli- hood to the power 1/T: Ptemp (D) oc p( w)) 1/T p(w) (Wilson & Izmailov, 2020). Smaller values for the temper- ature (i.e., T < 1) lead to more concentrated posteriors. A low temperature is beneficial since it can improve our approximation of the true Laplace by reducing the over- estimation of variance in certain directions, which arises when covariances between layers are ignored (Ritter et al., 2018). Additionally, it enhances the LA itself by preventing from placing too much probability mass in low-probability regions of the true posterior (Ritter et al., 2018).
22: for each mixture model m do
23: for k = O..... K₁ do
1--1 24: wh-11
w 7/ V ((wh))
25: end for
26: Set win
global
w wK.-- -1
27:
end for
elohal
28:
return I:M
A convenient way to approximate Equation (6) is again through ensembling so that the predictive distribution is ap- proximated as p(g) x, D) 37 EM= p(y]x, W..), where wn, represent a mode of the posterior (Wilson & Izmailov, 2020; Murphy, 2023). To find different modes of the global posterior, the server adopts the same approach as each client, by running SGD M times on the approximate global log-
Choice of the Prior. We exploited the commonly-used zero- mean Gaussian diagonal prior, p(w) = [N o² where d² is the variance (Wilson& & Izmailov, 2020).
5

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
5. Experimental Setup
symmetric Dirichlet sampling procedure with parameter (Hsu et al., 2019), where smaller value of a results in more heterogeneous data split. In addition, data are normal- ized before splitting, and small fraction of the training data (i.e., 500 samples) is kept at the server as validation data for the hyperparameter tuning of the various approaches.
5.1. Benchmark methods
We compared our method with well-cstablished state-of- the-art approaches for one-shot federated learning, ensuring that we included representative algorithms from each cate- gory described in Sec. 2.2, with special focus on methods that do not explicitly require dataset to be available at the server or involve multiple server-client communications. FedFisher (Jhunjhunwala et al., 2024). As we discussed in Sec. 2.2, this strategy is based on an unimodal approxima- tion of the loss as a quadratic function. The Fisher matrix is approximated through a diagonal matrix or through a Kro- necker factorization. Here, consider for comparison the method based on Kronecker approximation, referred to as FedFisher K-FAC, which achieves better results.
Regarding the models employed for the classification task, we used two CNNs as done in the literature (Jhunjhunwala et al., 2024): LeNet (LeCun et al., 1998) for FNShionMNIST, and a larger CNN proposed by (Wang et al., 2020) for the other two datasets.
5.3. Training details
We kept the basic training procedure fixed for all the ap- proaches and experiments to compare them fairly. SGD is utilized as local optimizer, with each client training for 20 epochs on LeNet and 50 epochs on the more complex CNNs for SVHN and CIFARIO 10. Moreover, we set the batch size to 64, the learning rate Te 0.01, and the momentum to 0.9, in line with (Jhunjhunwala et al., 2024). For the various parameters of each baseline, we used their default settings and the official implementations provided by the authors.
RegMean (Jin et al., 2023). It is model fusion approach that merges models in their parameter space, guided by weights that minimize prediction differences between the merged model and the individual models.
DENSE (Zhang et al., 2022). It is strategy that relies on data-free knowledge distillation, by exploiting GANs to artificially generate data for distillation at the server. OTFusion (Singh & Jaggi, 2020). It is a popular neuron- matching method that utilizes optimal transport to align neurons across the models by averaging their associated parameters, rather than directly averaging the weights.
Concerning FedBEns, we used the standard PyTorch weights initializer for the random initialization of the lo- cal ensemble. To perform Laplace approximation of the local posteriors, we exploited the open-source laplace package (Daxberger et al., 2021a), easy-to-use software library for Py offering access to the most common LA methods, discussed in Section 4.3. The server performs 300 steps to minimize the global loss using the Adam op- timizer with its standard default hyperparameters. During each server's optimizer run, conducted separately for each ensemble model, the validation performance is evaluated ev- ery 30 steps and the parameters configuration that achieves the best validation performance is selected as the final com- ponent of the ensemble. Last, for all the experiments we employed a (cold) tempered posterior for each client, with T = 0.1, and a diagonal Gaussian as prior with variance o² 0.1. An ablation study on these two hyperparameters can be found the Appendix.
Fisher Merge (Matena & Raffel, 2022). It is a model fu- sion baseline that, similarly to (Jhunjhunwala et al., 2024), employs diagonal approximation to the Hessian of the log-postenn posterior.
5.2. Datasets and models
To assess the performance of the proposed strategy, and to compare it with other baselines, we conducted several ex- periments on standard benchmarks and settings common in the literature (see for example (Jhunjhunwala et al., 2024)). We considered the following datasets:
FashionMNIST (Xiao et al., 2017). It is dataset of Za- lando's 's article images, consisting of a training set of 60, 000 samples and test set of 10, 000 instances. Each sample is a 28 x 28 grayscale image, associated with a label from 10 classes. Here we refer to this dataset as FMNIST. SVHN (Netzer et al., 2011). It is a real-world image dataset obtained from house numbers in Google Street View images. consists of 600, 000 digit images associated with a label from 10 classes.
All the experiments were performed on a machine equipped with an Intel Xeon 4114 CPU and NVIDIA Titan XP GPU with 12Gb of RAM. Our code is available at: https: git com/ 3 a copot 9 /FedBEn
6. Numerical Results
CIFAR10 (Krizhevsky et al., 2010). It consists of 60, 000 32 xxx33 32 color images in 10 different classes (e.g. airplane, dog, truck, etc.), with 6, 000 images per class. There are 50, 000 training images and 10, 000 test images.
6.1. Varying heterogeneity
In this Section, we compare the different methods in sce- nario where the number of clients is fixed (C' 5) while varying the heterogeneity parameter We used the same
To simulate data heterogeneity among C' client datasets, we partitioned the original image dataset into C subsets using a
6

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Table Test performance metries, 5 clients, varying heterogeneity parameter
Mean Accuracy & Standard Deviation (5 seeds)
EndBEns
FedBEns
FedFisher K-FAC
Fisher Merge
Dataset
RegMean
DENSE
OTFusion
M 5 Kron
M = Diag+Full
a 0.05
FMNIST SVHN
59.89 3.61
55.58 + 2.96
53.92
1.89
48
38.48
5.56
37.26
1.01
43.11
6.118
71.55 = 1.65
53.97 + 1.87
56.46
3.70
65.57
+ 1.83
56.89
= 0.75
39.17
0.42
43.53
2.54
CIFARIO
51.43 + 0.79
33.69
+ 3.38
36.97 + 2.01
33.77 + 0.52
30.61
= 1.95
29.59
+ 1.77
28.44
144 3.37
Average
60.90
47.75
49.58
49.13
41.99
35.33
38.36
ar 0.1
FMNIST
75.91 3.21
70.04 + 1.13
68.06
2.64
56.40
+ 5.54
47.16
1.73
2.44
57.38
5.66
SVHN
80.67 0.48
63.97 + 1.51
65.52
1.42
71.01
+ 1.29
62.65
0.91
51.94
0.67
52, 66
4.04
CIFAR1O
57.62 = 1.24
46.68 + 2.64
46.92 + 1.79
40.86
+ 0.77
38.04 + 1.78
40.41
+ 0.73
36.56 + 1.51
Average
71.40
60.23
60.17
58.13
49.32
45.30
48.87
a 0.2
FMNIST
79.03 + 2.71
77.24 + 1.88
72.89
3.19
71.43
+ 1.61
69.07
1.23
57.06
2.19
66.89
+ 3.50
SVHN
82.91 + 0.26
77.71 + 0.31
71.61
5.14
78.21
+ 0.81
77.13
3.39
72.19
1.00
68.72
0.68
CIFAR10
59.78 + 0.92
49.12 + 1.15
19.95
2.51
43.42 + 1.54
44.94 + 2.50
40.64 + 2.07
40.54 + 3.11
Average
73.91
68.02
63.72
64.36
64.81
56.63
58.72
Q 0.4
FMNIST
84.42 0.77
81.81 + 0.26
74.76
2.33
74.45
1: 1.89
77.91
1: 1.87
67.29
1
4.78
78.16
1 0.62
SVHN
85.58
82.15 1 0.71
78.31
0.32
80.11
1.03
78.40
2.49
73.79
3.25
73.34
1.77
CIFARIO
61.49 1 0.72
48.96 1 0.16
51.89
1.99
41.82
1 3.70
46.47
1 2.07
45.84
0.30
41.78
3.98
Average
77.16
70.92
68.32
65.46
67.59
62.31
64.43
FashionMNIST
30
Accuracy (1)
Rcculracy (16
Accutres (16
2 0.05
d 0.05
d= 0.05
0 0.1
22 0.1
a 0.1
0.2
0.2
0.2
0.4
0.4
Numbar of Mixtures
Number of Mixtures
Number of Mixtures
Figure Test accuracy as a function of the number mixtures for FedBEns Kronecker factorization. each dataset, mean accuracy is represented by solid line, while the shaded band indicates the standard deviation, both computed over 5 seeds, for various heterogeneity parameters.
parameters across all settings, such the number of local epochs, to ensure fair comparison. Specifically, consid- ered settings ranging from high to moderate heterogeneity (a {0.005,0 0.1, 0.2, 0.4}) (Jhunjhunwala et al., 2024) to analyze their impact on the different approaches.
deviation for our method with fixed number of mixtures, M 5, and compare it with the other considered baselines. Results were obtained over 5 experiments with different random seeds. Additionally, we report the average accuracy across datasets for each model and experimental setting to facilitate a clearer comparison. Bold numbers indicate the highest accuracy, while underlined numbers represent the
Table shows the average accuracy along with the standard

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Table 2. Test performance metrics, o 0.3, varying number of clients
Mean Accuracy & Standard Deviation (5 seeds)
EndBEns M= = Kron
FndBEns
FedFisher K-FAC
Fisher Merge
Dataset
RegMean
DENSE
OTFusion
M = 5 Diag+Full
10
FMNIST
80.57 0.86
78.66 1.73
73.85
3.91
73.97 + 3.18
74.33
3.95
64.34
66,84 4.54
SVHN
83.75 13.75-20 + 0.22
75.98 + 1.85
74.94 + 0.69
78.64
0.95
64.99
= 5.75
67.76
07.76
3.39
63.99 + 0.39
CIFARIO
58.14 + 1.19
48.18 + 1.52
52. 14
+ 1.73
44.48 +
1.74
45.39
= 3.18
41.39
+ 1.28 38.10 + 2.58
Average
74.15
67.61
66.98
65.65
61.57
57.83
56.31
20
FMNIST
80.83 + 0.54
76.16 1.04
72.83
3.21
72.91
1.56
74.26
= 3.89
63.46
6.44
64.11
+ 4.49
SVHN
81.21 + 0.85
74.33 + 3.72
78.27 = 0.48
79.04
0.89
58.93
= 3.81
59.25
- 4.66
62.23
3.18 3.18
CIFARIO
56.41 + 1.81
43.39 + 2.58
48.85 = 0.84
42.81 + 2.37
43.75
= 1.91
38.98 + 1.69
33.69 + 2.22
Average
72.77
64.63
66.65
64.79
58.98
53.90
53.34
40
FMNIST
77.27 0.84
72.69
+ 1.55
72.06 2.73
72.09 + 1.05
71.81
= 1.23
61.4844.1
64.29 + 4.33
SVHN CI
80.06 + 0.56
65.96 + 2.97
76.85 + 0.98
76.94
0.86
59.84
+ 3.76
50.91
+ 4.48
52.56
5.18
54.87 + 1.72
42.38 + 1.45
47.76 + 0.69
44.35 + 0.97
41.35
+ 1.83
34.53
+ 3.02
33.81 + 2.56
Average
70.73
59.68
65.56
64.46
57.67
49.97
50.22
second-best performance. From this table, it is possible to conclude that FedBEns with Kronecker factorization, systematically outperforms the baselines. This highlights the effectiveness of FedBEns One-Shot FL algorithm, especially in scenarios with high data heterogeneity.
this computational-accuracy tradeoff are provided in the Appendix. Additionally, we compare our approach with FedFisher when run over multiple communication rounds, ensuring comparable total communication costs between the two approaches. We also conduct separate comparison where we equalize client computation costs by rumning Fed- BEns on smaller datasets. Even in these settings, FedBEns demonstrates competitive performance (see the Appendix for more details).
To further examine the impact of the number of mixture components M, Figure shows the accuracy of FedBEns with Kronecker factorization of the Hessian as function of the ensemble cardinality. The results demonstrate an almost monotonic improvement in accuracy as the number of components increases until a plateau is reached, with significant boost observed even with just two three com- ponents. Notably, our proposed approach achieves compet- itive performance with respect to state-of-te oatt--of-the-att methods while only using a few components, and often surpasses them with just two components. fact, the accuracy of FedBEns with single mixture (M 1) is comparable to, or slightly better than, the baselines across all datasets and heterogeneity levels (compare Figure 2 and Table 1), and FedBEns outperforms the baselines already for MM = 2.
6.2. Varying number of clients
Table 2 reports the results for varying the number of clients (C E {10, 20, 40)), with the heterogeneity parameter fixed at a = 0.3, to provide insights into the behavior of the proposed approach as the number of clients increases.
It is worth noting that, in this scenario, each client dataset is smaller, which increases the risk of local models overfitting and makes the aggregation more challenging. From this analysis emerges that FedBEns with Kronecker factorization outperforms the other baselines.
We also conducted an ablation study on Hessian approxima- tions, detailed in the Appendix, highlighting the advantages of Kronecker factorization over other methods.
Overall, our approach demonstrates highly competitive per- formance, outperforming state-of-the-art models, by up to 10 percentage points in some cases, across different hetero- geneity settings and client numbers.
The possibility of choosing the number of mixtures adds degree of flexibility control the trade-off between com- putational/coomuniccttoo costs (higher when the number of components is higher) and accuracy, making our ap- proach suitable to real-world applications with widely dif- ferent requirements. Further analyses and discussions of
7. Conclusion and Future Work
In this paper, we proposed a novel approach for One-Shot Federated Learning based on Bayesian inference. Our
8

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
method utilizes a mixture of Laplace approximations to model client posteriors, which are used to estimate the global posterior. Our proposed approach outperforms other
tional Conference on Learning Representations (ICLR), 2021.
state-of-the-an art models, with respect to classification accu- racy, across different scenarios.
A Allouah, Y., Dhasade, A. B., Guerraoui, R., Gupta, N., Ker- marrec, A.-M., Pinot, R., Pereira Pires, R., and Sharma, R. Revisiting ensembling in one-shot federated learning. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Many interesting research directions remain open for explo- ration. One of the main limitations of the proposed approach its computational cost, since it based on a mixture of M models trained in parallel instead of just one. However, it is worth mentioning that there are approaches that miti- gate this problem (e.g. (Havasi et al., 2021)) by mimicking ensemble prediction within a single model. Moreover, we plan to study the privacy risks associated with the proposed approach and try to enhance it, e.g., through differential privacy techniques. Additionally, we would like study the adaptability of FedBEns to dynamic environments where, for instance, new clients join the Federated I Learning scheme and/or clients collect new types of data and experience a distribution shift.
Cao, L.., Chen, H.,, Fan, X., Gama, J., Ong, Y. and Kumar, V. Bayesian federated learning: A survey. In International Joint Conference on Artificial Intelligence (IJCAI), 2023.
Chen, H.-Y. and Chao, W.-L. Fedbe: Making bayesian model ensemble applicable to federated learning. In Inter- national Conference on Learning Representations (ICLR), 2021.
Cho, Y. J.. Manoel, A., Joshi, G., Sim, R., and Dimitriadis, D. Heterogeneous ensemble knowledge transfer for train- ing large models in federated learning. In International Joint Conference on Artificial Intelligence (IJCAI), 2022.
Acknowledgement
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, Laplace redux- effortless bayesian deep learning- In Advances in Neural Infor- mation Processing Systems (NeurlPS), 2021a.
The research leading these results has been partially funded by the Italian Ministry of University and Research (MUR) under the PRIN 2022 PNRR framework (EU Contri- bution- - NextGenerationEU - M. 4,C. I. 1.1), SHIELDED project, ID P2022ZWS82; by the French government, through the 3IA Côte d' Azur Investments in the Future project managed by the National Research Agency (ANR) with the reference number ANR- HIR-19-P31A-00022 by the Euro- pean Network of Excellence dAIEDGE under Grant Agree- ment Nr. 101 120726; by the EU HORIZON MSCA 2023 DN project FINALITY (G.A. 101 191166166 16); by the Groupe La Poste, sponsor of the Inria Foundation, in the framework of the FedMalin Inria Challenge.
Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., and Hernández-Lobato, J. M. Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning (ICML., 2021b.
Eschenhagen, R., Daxberger, E., Hennig, P., and Kristiadi, A. Mixtures of laplace approximations for improved post-hoc uncertainty in deep learning. arXiv preprint arXiv:2111.03577, 2021.
Fort, S., Hu, H., and Lakshminarayanan, B. Deep en- sembles: A loss landscape perspective. arXiv preprint arXhv:191..2277, 2019.
Impact Statement
This work advances the field of machine learning by intro- ducing an algorithm that enables Federated Learning with a single round of communication between a central server and clients. We believe this work mainly aligns with the increasing demand for user privacy in widely different tech- nological domains, which one of the main advantages of the Federated Learning paradigm. This has many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Gong, X., Sharma, A., Karanam, S., Wu, Z.,, Chen, T., Doer- mann, D.,, and Innanje, A. Ensemble attention distillation for
federated learning. In IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
Gong, X., Sharma, A., Karanam, S., Wu, Z.,, Chen, T., Doer- mann, D., and Innanje, A. Preserving privacy in federated learning with ensemble cross-domain knowledge distil- lation. In AAAI Conference on Artificial Intelligence. 2022.
References
AI-Shedivat, M., Gillenwater, J., Xing, E., and Ros- tamizadeh, A. Federated learning via posterior averaging: A new perspective and practical algorithms. In Interna-
Grosse, R. and Martens, J. A kronecker-factored approxi- mate fisher matrix for convolution layers. In International Conference on Machine Learning (ICML), 2016.

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Guha, N., Talwalkar, A., and Smith, V. One-shot federated learning. arXiv preprint
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z.,, Cormode, G., Cummings, R., et al. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1-2):1-210, 2021.
2019.
Guo, H., Greengard, P., Wang, H., Gelman, A., Kim, Y., and Xing, E. Federated learning as variational inference: A scalable expectation propagation approach. In Interna- tional Conference on Learning Representations (ICLR), 2023.
Kristiadi, A., Hein, M., and Hennig, P. Being bayesian, even just a bit, fixes overconfidence in relu networks. In International Conference on Machine Learning (ICML), 2020.
Hamer, J., Mohri, M., and Suresh, A. T. Fedboost: A communication--efiiiee algorithm for federated learn- ing. In International Conference on Machine Learning (ICML), 2020.
Krizhevsky, A., Nair, V., and Hinton, G. Cifar-10 10 (canadian institute for advanced research). URL http://www. cs. toronto. edu/kriz/cifar. hmml, 5(4):1, 2010.
Hasan, M., Zhang, G., Guo, K., Chen, X., and Poupart, P. Calibrated one round federated learning with bayesian inference in the predictive space. In AAAI Conference on Artificial Intelligence, 2024.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Siimle and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Process- ing Systems (NeurIPS), 2017.
Havasi, M., Jenatton, R., Fort, S., Liu, J. Z., Snoek, J., Lakshminarayanan, B., Dai, A. M., and Tran. D. Train- ing independent subnetworks for robust prediction. In International Conference on Learning Representations (ICLR), 2021.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 1):2278 2324, 1998.
Lin, T., Kong, L., Stich, S. U., and Jaggi, M. Ensemble distillation for robust model fusion in federated learning. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2016.
Heinbaugh, C. E., Luz-Ricca, E., and Shao, H. Data-free one-shot federated learning under very high statistical heterogeneity. In International Conference on Learning Representations (ICLR), 2023.
Liu, C., Lou, C., Wang, R., Xi, A. Y., Shen, L., and Yan, J. Deep neural network fusion via graph matching with applications to model ensemble and federated learning. In International Conference on Machine Learning (ICML), 2022.
Hsu, T.-M. H., Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual clas- sification. arXiv preprint
Liu, X., Liu, L.,, Ye, F., Shen, Y., Li, X., Jiang, L., and Li, J. Fedlpa: One-shot federated learning with fayer-wise posterior aggregation. arXiv preprint arXiv:23 10.00333, 2024.
2019.
Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. G. What are bayesian neural network posteriors really like? In International Conference on Machine Learning (ICML)), 2021.
MacKay, D. A practical bayesian framework for backprop- agation networks. Neural Computation, ((3):448-477 1992.
Jhanjhunwala, D., Wang. S., and Joshi, G. Fedfisher: Lever- aging fisher information for one-shot federated learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
MacKay, D. Probable networks and plausible predictions- review of practical bayesian methods for supervised neural networks. Network: Computation in Neural Sys- tems, 6(3):469, 1995.
Jin, X., Ren, X., Preotiuc- Pietro, D., and Cheng, Data- less knowledge fusion by merging weights of language models. In International Conference on Learning Repre- sentations (ICLR), 2023.
MacKay, D. Information theory, inference and learning algorithms. Cambridge University Press, 2003.
Jordan, K., Sedghi, H., Saukh, O., Entezari, R., and Neyshabur, B. Repair: Renormalizing permuted activa- tions for interpolation repair. In International Conference on Learning Representations (ICIR), 2023.
Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and
Wilson, A. G. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems (NeurIPS), 2019.
10

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Martens, J. and Grosse, R. Optimizing neural networks with kronecker-ffctoned approximate curvature. In Interna- tional Conference on Machine learning (ICML), 2015.
Zhou, Y., Pu, G., Ma, X., Li, X., and Wu, D. Distilled one- shot federated learning. arXiv preprint arXiv:000900999, 2020.
Matena, M. S. and Raffel, C. A. Merging models with fisher- weighted averaging. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. Communication-efien efficcient learning of deep networks from decentralized data. In International Con- ference on Artificial Intelligence and Statistics (AISTATS), 2017.
Murphy, K. P. Probabilistic machine learning: advanced topics. MIT Press, 2023.
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng.
A. Y., et al. Reading digits natural images with unsu- pervised feature learning. In NIPS Workshop Deep Learning and Unsupervised Feature Learning, 2011.
Ritter, H., Botev, A., and Barber, D.
A scalable laplace approximation for neural networks. In International Con- ference on Learning Representations (ICLR), 2018.
Singh, S. P. and Jaggi, M. Model fusion via optimal trans- port. In Advances Neural Information Processing Sys- tems (NeurIPS), 2020.
Wang, Liu, Q., Liang, H., Joshi, G., and Poor, H. V. Tackling the objective inconsistency problem in hetero- geneous federated optimization. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Wilson, A. G. and Izmailov, P. Bayesian deep learning and a probabilistic perspective of generalization. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1788:0777 2017.
Yadav, P., Tam, D., Choshen, I., Raffel, C., and Bansal, M. Resolving interference when merging models. arXiv preprint ar arXiv:2366.07708 1, 2023.
Zhang, J., Chen, C., Li, B., Lyu, L.,, Wu, S., Ding, S., Shen, C., and Wu, C. Dense: Data-free one-shot federated learning, In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Zhang, Y., Jia, R., Pei, H., Wang. W., Li, B., and Song, D. The secret revealer: Generative model-inversion attacks against deep neural networks. JEFE/CVF Conference on Computer Vision and Pattern Recognition (ICCV), 2020.

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
A. Proof: Global Posterior
Proposition A.1. Given datasets and the posteriors p(w (Da) of the same model f(w), under the assumptions that (i) the datasets are conditionally independent given the model, and (ii) the prior is the same across clients, the global posterior can be written as:
p(w DI. Dc) II I(w()
p(w)
Proof:
The starting point is the global likelihood function p(D)w), w). Under the assumption that clients conditionally independent, it can be decomposed as:
p(D = p(D1... Delw)
II p(D./w) c=1
(8)
To obtain the posterior, it possible to rely on Bayes' theorem p(w/2) p( w))()
p(D)
I P(De w) p(w) p(D1, .., De)
p(w ()) p(w(DI,... Dc)
(9)
Equivalently, instead of combining the likelihoods, it is possible rewrite the previous Equation terms of each client posterior, exploiting Bayes' theorem again:
p(w)
p(w(Di....Do
II ()
(10)
p(D1, Dc)
p(w)
C p(Do)p(w(Da
(11)
P(D ...,, Dc)
p(w)
1 P(D,)
I Ie (.)
(12)
p(D Dc)
C-w
20
Remark 1. This expression can be generalized to different local priors, denoted (w(w,, as:
p(w (D.)
pi.,,,,.., Dc) Pelobbll (w) II
(13)
(w(w)
where Pelobal (w) represent the global prior, provided by the server.
Remark The conditional independence assumption plays crucial role. Otherwise, one has to model the joint likelihood explicitly, since p(D), .., Pe/w) * Ie p(Da(w). A possible approach consider an autoregressive-like likelihood. For instance, for clients, correlated log-likel likelikhkhod can be derived from Gaussian- approximation of joint likelihood, as:
(p(pp w)) = log(p(D) w)) + log(p(D2 w)) + p. (D,, DD,,,2, w)
(14)
where: p E [-1, correlation coefficient and g is function that captures the functional dependencies among clients' data.
12

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
The global posterior in Equation (2) factorizes independently across clients. Consequently, incorporating new clients into the federation forgetting clients that leave the federation is straightforward: it is sufficient simply add remove their corresponding mixture component from the global posterior.
Remark 3.
B. Additional Related Work
In this section, we discuss additional related work on FL, with a focus on papers based on an ensemble of models. An early contribution that leverages ensemble and Bayesian inference in the standard multi-rounds FL setting FedBE (Chen & Chao, 2021), where the aggregation at the server is based on the local posterior approximated as a Gaussian Dirichlet distribution, relying on the stochastic weight average- Gaussian (SWAG) method (Maddox et al., 2019). Models are then sampled from the global (unimodal) distribution and the knowledge of this ensemble is then distilled in single model, under the assumption that a certain amount of unlabelled data is available at the server.
average-Gaussian
Regarding other approaches that exploit ensembles in a Federated Learning scenario, it is worth mentioning FENS (Allouah et al., 2024), FedKD (Gong et al., 2022), Fed-ET (Cho et al., 2022), FedBoost (Hamer al., 2020), and FedeDF (Lin et al., 2020). FENS One-Shot FL method that operates in two stages: first, clients train models locally and send them the server; second, clients collaboratively train a lightweight prediction aggregator model using standard multi-rounds FL. Other approaches (c.g, FedKD and Fed-ET) are designed for standard multi-rounds FL settings where knowledge distillation is exploited to reduce communications costs (Gong et al., 2022) and/or improve generalization by introducing a data-aware weighted consensus from the ensemble of models, with a feedback loop to transfer the server model's knowledge to the client models (Cho et al., 2022).
Our approach differs from these works does not require data on the server side or additional server-clients communication rounds. Moreover, the underlying theoretical motivation is also different: our method leverages ensembles to achieve tractable and effective approximation of the posterior and its Bayesian marginalization, whereas other approaches employ ensembles guided principles of knowledge distillation and/or stacked generalization.
C. Ablation and Further Studies
C.1. Hessian matrix approximations
One key design feature of the proposed approach is the number of mixtures, M. Increasing the ensemble size improves predictive performance but raises server-client communication costs. Additionally, also the choice of the approximation structure for the Hessian matrix estimation impacts performance and computation. Here we consider the following approximations:
Diagonal. It is the simplest and most scalable approach that ignores off off-diagonal terms of the Hessian. Denoting the total number of parameters of the model as d, diagonal approximation requires additional d parameters for the Hessian matrix.
Diagonal+Full. It adds expressivity to the previous diagonal approximation by considering a full Hessian for the last layer. In this case, the computational cost raises w.r.t the diagonal approximation by an amount that scales with the last-layer number of neurons, Tlaasts and the number of classes, Relasses, as O (Tiaass MPclasses) )
Kronecker. As we discussed in Section 4.3, it is one of the most common approximations of the Hessian since it is a good trade-off between computational cost and expressivity (Daxberger et al., 2021a).
In general, the diagonal approximation is less expensive, while the comparison between Kronecker and Diagonal+last Full depends on the model architecture. For instance, for the considered CNN employed on SVHN and CIFARIO and denoting, as said, the total number of parameters d, Kronecker factorization requires 3.9d parameters, while Diagonal+ Full approximation ~ 2.7d. However, for other architectures (e.g. those with a substantially higher number of classes), Kronecker may become more cost-effective.
To assess the impact of each approximation on predictive performance, we conducted an ablation study examining three types of approximations across different numbers of mixtures. Our analysis focused on two different heterogeneity scenarios, ie., a 0.05 (highly heterogeneous), and a = 0.4 (medium/low heterogeneity). Figure 3 shows the results. As expected,
13

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
KRON
DIAG FULL
DIAG
FMNIST a 0.05
SVHN a = 0.05
CIFAR10 a 0.05
70
65
Accuracy (96
Accuracy (9)
Accuracy (%)
45
50
55
40
50
45
35
40
40
30
FMNIST 0.4
SVHN a= 0.4
CIFAR10 0.4
84
86
62.5
82
84
Accuracy (%)
Accuracy (%6
Accuracy (9)
57.5
55.0
78
80
52.5
76
78
74
76
47.5
45.0
2
Number of Mixtures
Number of Mixtures
Number of Mixtures
Figure Test accuracy as function the number of mixtures with different Hessian approximations. Curves represent mean accuracy computed over 3 seeds.
the diagonal approximation yields the worst performance, while the Kronecker factorization consistently achieves the highest accuracy, highlighting the importance of capturing correlations among model parameters for each layer.
C.2. Temperature parameter
In Section 4.3 we introduced temperature scaling and, particularly, we asserted that cold temperature T < ensures to obtain a more concentrated posterior that prevents Laplace approximation from giving too much mass to certain regions (e.g. due to the approximation of the Hessian). Figure 4 presents the results of an ablation study using our approach with mixtures and Kronecker Hessian approximations. while low temperatures tend to yield better results, also demonstrating certain robustness with respect the temperature choice. It is worth noting that, in principle, the performance of our approach could be further improved by tuning the temperature parameter for each experiment (e.g. through a grid search on the server if validation dataset is available), while the results we presented are with a fixed temperature of 0.1.
general, high temperatures (i.e., T 1) result poor performance,
C.3. Prior variance
We present an ablation study on the variance of the isotropic Gaussian prior p(w) = o(00022), Figure reports the results for EndBEns with Kronecker factorization, 5 mixtures, and 7' = 0.1 on 5 clients. The predictive performance remains robust across a reasonably wide range of prior variance values.
C.4. On the role of the validation set
As stated in Section 5.3, our method uses small validation set during the server's gradient descent runs, to find the parameter configuration that achieves the best validation performance. In the following, we report an ablation study where no validation dataset is used and the number of server epochs is fixed to 100. The results are reported in Table 3. Remarkably, FedBEns still outperforms the considered baselines, which continue to leverage the validation dataset on the server side
14

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
FashionMNIST
SVHN
CIFAR10
60
Accuracyy (86)
70-
55-
Accuracy
Accuracy
50
60
50-
Temperature
Temporature
Temperature
Figure 4. Test accuracy for FedBEns with Kronecker factorization and mixtures, as a function of the temperature parameter. Curves represent the mean accuracy computed over seeds.
-
0.05
0.4
FashionMNIST
CIFAR10
60
90
75
55
Accuracy (86)
Accuracy (86)
70
Accuracy (96)
65
60
65
40
35
10
10.
10"2
10
10P
10
10
10
10
10-
300
10'
10
10-7
10
10"
10'
10'
Figure 5. Test accuracy for EndBEns with Kronecker factorization and 5 mixtures, function of the prior variance o², Curves represent the mean accuracy computed over seeds.
tune certain hyperparameters depending on the approach, with only minor drop in accuracy, up to few p.p. (see Table I in the paper and Table 3 below). This makes FedBEns applicable even in scenarios where such a small amount data not available at server side.
C.5. Computational and communication costs
In this Section, we comment on the computational and communication costs associated with the different approaches. We observe that the FndBEns per-client computational and communication costs are both linear in the number mixtures M. the contrary, the computation time at the server seales as M², since there the need to find M modes of global log-posterior through M1 SGD runs, and the cost of each evaluation of the global log-posterior scales linearly with M (see Equation (4)).
In Table 4 we provide the communication cost and a wall-clock time comparison with the other baselines computed on CIFARIO with 5 clients on the hardware specified in Sec. 5.3. The reported values are the average over runs. Overall, our approach increases computation and communication costs depending on the number of mixture components, but also improves predictive performance. Whether this additional cost is justified depends on the practical context. the following paragraphs, we provide further insights into the trade-off between predictive performance and computational cost.
15

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Table 3. Test performance metries of EndBEns with no validation
Mean Accuracy & Standard Deviation (3 seeds)
FedBEns
Dataset
M = 5 Kron
a: 0.05
FMNIST
59.15 3.87
SVHN
70.26 + 1.71
CIFARIO
48.81 + 0.98
a = 0.4
FMNIST
82.16 + 0.83
SVHN
84.69 1 0.27
CIFARIO
60.00 1 0.75
Table ConputaioonalConmmiitttioo costs
Method
Client runtime (s)
Server runtime (8)
Communication cost (MB)
FedBEns Kron
30.9 M
7.2 - M
17.8 - M
FedFisher K- K-FAC
30.9
13.4
17.8
RegMean
26.0
5.1
7.3
DENSE
25.9
170.2
3.65
OTFusion
25.9
1.1
3.65
FisherMerge
26.2
0.9
7.31
C.6. Mixture in One-Shot vs. multiple rounds
In this Section, we compare the proposed approach with its main competitor, FedFisher, while ensuring the same server-client communication cost. In fact, FedFisher has been extended to support multiple rounds of server-client communication, as detailed in (Jhunjhunwala et al., 2024). Specifically, we compare FedFisher K-FAC with rounds of communication to our approach using mixtures in the One-Shot setting. In Table 5 we report the accuracy for the described methods in two scenarios, i.e., high (a 0.05) and medium/low (a = 0.4) heterogeneity.
Although the fairness of this comparison (one round of communication vs. multiple rounds) may be questionable, is remarkable that FedBEns gives competitive results, especially in more balanced scenario. Note that even though the communication cost is the same, sending all the needed information in single round of communication, if the client-server communication channel has sufficient capacity, may be preferable in specific cases, e.g. when the client-server connection is not continuously available or possible. Moreover, while FedBEns with M mixtures and FedFisher with M rounds entail similar per-client communication/coopuattio costs, the total time a client must remain active during training can differ significantly due to server-to-client transmissions and the straggler effect.
First, consider server-to-client communication. Assuming equal server/clients uplink and downlink capacities, the communi- cation time for M-round FedFisher scales as 2.M, since the model must be sent from client server and back in each round. In contrast, FedBEns communication time is halved since mixtures are sent once and then aggregated. If communication is the system bottleneck (as is often the case in FL), training with M-round FedFisher would require each client to remain available for roughly twice as long as FedBEns with M mixtures. Notably, the final model in FedBEns remains unchanged even if a client disconnects after submitting its mixtures.
Second, consider the straggler effect. The reasoning above assumes homogeneous clients and synchronization. However, real-world systems affected by stragglers, clients whose updates arrive significantly later than those of others (Kairouz et al., 2021). Stragglers slow down the entire system, as the server must wait for their inputs before generating a new model to distribute. This delay, in turn, forces other clients to remain idle, further extending their total participation time. One-shot algorithms mitigate this issue, as clients can disconnect immediately after submitting their contributions.
Finally, It is also important to highlight the value of minimizing client participation time. cross-device setting, clients often exhibit volatile participation patterns. Therefore, minimizing the required client's ss participation time is critical: may
16

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Table 5. Test set performance metrics, multiple mixtures multiple rounds
Mean Accuracy & Standard Deviation (3 seeds)
EndBEEs M = 5 Kron
FedFisher K-FAC
Dataset
5 communication rounds
a 0.05
FMNIST SVHN
60.08 3.60
66.49 4.89
71.39
1.63
67.27 + 2.05
CIFARIO
51.38 + 0.79
58.28 + 1.03
a = 0.4
FMNIST
84.39 0.78
82.3 0.94
SVHN
85.57
0.15
83.06 1 1.53
CIFARIO
61.56 1 0.73
68.14 1 2.03
Table 6. Test set performance metries, multiple mixtures smaller datasets
Mean Accuracy & Standard Deviation (3 seeds)
FedBEns
Dataset
M = 5 Kron
FedFisher K-FAC
Reduced Datasets
0.05
FMNIST
50.87 4.46
53.92 1 1.89
SVHN
59.12 + 1.42
58.46 + 3.70
CIFARIO
45.38
0.62
36.97 = 2.01
0.4
FMNIST
77.95
0.31
74.76 1 2.33
SVHN
81.15 + 0.25
78.31 + 0.32
CIFARIO
53.82 1.16
51.89 + 1.99
determine whether client contributes to the training drops out before completion, thereby preventing the federation from leveraging its dataset and degrading the final model's quality (Kairouz et al., 2021). Finally, we observe that in FndBEns each client can transmit each mixture as soon as it is computed. As a result, even if a client disconnects before completing all M mixtures, the system may benefits from its partial contribution (see also C.8).
C.7. Mixture in One-Shot on smaller datasets vs. single model on larger datasets
In this Section, we present an additional comparison with FedFisher, where we equalize the client computation of our approach by running FedBEns with mixture components on client training datasets that are 5 times smaller. Remarkably, FedBEns performs better than FedFisher in most scenarios, as illustrated by the results reported in Table
C.8. System heterogeneity: customizing the number of mixtures for each client
As previously discussed, FedBEns introduces a computational overhead that depends on the number of mixture components. In practical scenarios, clients may have different computational capabilities, so this additional computation could be an issue for some of them, but not for others. this Section, we explore the possibility of customizing the number mixtures based on each client's computational power. As an illustrative example, we considered a setting with 5 clients: 3 can only employ single Laplace approximation (M = 1) and the remaining 2 are able exploit three mixtures (M 3). We considered two different approaches. In the first, all 5 clients use unimodal approximation of the posterior (M = 1). In the second, the number of mixture components is customized: 3 clients use M = 1, while the remaining use M = 3, obtained by running 3 SGD optimizer in parallel from the same initialization. In both cases, the final model returned to each client consists of a single model.
Table 7 reports the results the described settings, indicating that the additional information coming from the clients with
17

---

FedBEns:
One-Shot Federated Learning based on Bayesian Ensemble
Table 7. Test set performance metrics, customizing the number of mixtures
Mean Accuracy & Standard Deviation (3 seeds)
FedBEns Kron
FedBEns Kron Mcustomized
Dataset
M = 1
a 0.05
FMNIST
57.02 + 3.73
59.57 1.44
SVHN
66.11 + 0.77
68.39 + 1.57
CIFARIO
39.89 + 1.71
42.32 + 2.25
Q = 0.4
FMNIST
77.07 + 0.29
78.31 + 1.76
SVHN CIFARIO
82.86 1 0.21
82.75 1 1.41
53.97 1 0.86
54.68 1 1.88
Table 8, Test set performance metrics, ensemble vs, single model
Mean Accuracy & Standard Deviation (3 seeds)
Dataset
Ensemble
Min
Max
0.05
FMNIST
60.08 + 3.60
53.41 + 1.78
59.98 + 1.68
SVHN
71.39 + 1.63
62.24 + 1.67
68.99 + 1.4
CIFARIO
51.38 1 0.79
36.21 + 4.68
42.77 + 0.32
0.4
FMNIST
84.39 1 0.78
76.69 1 0.95
79.87 1 0.76
SVHN
85.57 + 0.15
81.25 + 0.95
83.38 = 0.11 54.27 = 0.81
CIFARIO
61.56 0.73
49.22 = 1.60
more computational power benefits the entire federation. This advantage allows the server to find a better global model, especially in heterogencous settings.
C.9. Ensemble vs. single component for global model
As discussed in Section 3, our approach leverages a multimodal approximation of the global posterior and exploits an ensemble of final models, in line with Bayesian Model Averaging perspective. In this Section, we analyze how much the ensemble contributes to enhance the predictive performance compared to selecting, for example, the ensemble component with the highest accuracy.
In Table we present the final accuracy of the ensemble, along with the highest (Max) and lowest (Min) accuracies achieved by its individual components. These results highlight the benefits of ensembling, as relying on the solution provided by a single SGD run on the global loss may result in suboptimal performance.
C.10. more challenging setting: ResNet18 on CIFARIO0
We also conducted an experiment on a more challenging task, represented by CIFAR-100 (Krizhevsky al., 2010). As a base model, we used the ResNet18 architecture (He et al., 2016) without batch normalization ensure compatibility with the different One-Shot approaches. Training was conducted for 100 epochs for each client. Given our hardware limitations, were able to run FedBEns with only 2 mixtures using Kronecker Factorization. Table 9 summarizes the performance of 5 clients in two different heterogeneity settings. As reference, the employed model achieves an accuracy of 58.55% in the usual centralized setting. Even in this more complex task, in which training in One-Shot is challenging, FedB Ens shows the highest performance, followed by the DENSE algorithm. In particular, it is worth mentioning that at the cost of doubling the computaton/coommnicttiion cost w.r.t FedFisher, FedBEns achieves a higher gain in accuracy (doubled or even more).
18

---

FedBEns: One-Shot Federated Learning based on Bayesian Ensemble
Table 9. Test set performance metrics, 5 clients, varying heterogeneity parameter
Mean Accuracy & Standard Deviation seeds)
EndBEns M 2 Kron
FedFisher K-FAC
Fisher Merge
Dataset
RegMcan
DENSE
OT'Fusion
0.05
CIFARI00
8.56 0.47 3.14 +: .23 3.01 0.54 5.88 + 1.07 4.79 1. 0.19 2.81 1 0.33
a 0.4
CIFARIO0
17.85
0.67
4.79
0.31
4.85
1.16
13.92 + 0.54
4.74 + 0.56
3.87
0.48
Communication costs (MB)
716
358
93
47
47
94
Local runtime (min)
28.0
14.0
9.2
9.2
10.1
Server runtime (min)
2.3
1.2
0.6
10.4
0.2
0.3
19


Reviews:
[{"review_id": "LgKuWNcvpS", "reviewer": "Reviewer_wHXT", "summary": "The paper introduces FedBEns, a one-shot federated learning (FL) algorithm using Bayesian inference to address multimodal local loss functions. It approximates local posteriors with a mixture of Laplace approximations (GMM) and aggregates them to estimate the global posterior. The server identifies global modes via SGD and performs ensemble predictions. Experiments on FashionMNIST, SVHN, and CIFAR10 show FedBEns outperforms baselines (e.g., FedFisher, RegMean) by up to 10% in accuracy, especially under high data heterogeneity (α=0.05).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are happy the reviewer appreciated the paper.\n \n**1) Runtime/communication costs quantification:**\nWe observe that the FedBEns per-client computational and communication costs are both linear in the number of mixtures M. On the contrary, in our implementation, the computation time at the server scales as M$^2$ (we have to find M modes of the global log-posterior through M SGD runs, and the cost of each evaluation of the global log-posterior scales linearly with M, see Eq. (4) ).\nWe provide experimental results for the CIFAR10 dataset with 5 clients on the hardware specified in Sec. 5.3. Times are averaged over 5 runs.\n- Local training time: for each mixture component, 25.9s for training the NN and 5.0s for computing the Hessian with Kronecker factorization.\n- Communication cost: for each mixture component, 17.8 MB (3.6 MB for the NN and 14.2 for the Hessian with Kronecker factorization).\n- Server aggregation time: ($O(M^2)$). We report the execution time for different values of M: M=1: 7.1s, M=2: 28.5s, M=3: 65.9s, M=4: 113.5s, M=5: 173.5s.\n \nWe provide the communication and computation costs with the other baselines in the following table. For additional comments on this aspect and results comparing FedBEns and FedFisher under similar computational cost, see answer 2) to reviewer dxFK:\n \n|Method|Client runtime (s)|Server runtime (s)|Communication cost (MB)|\n|-|-|-|-|\n|FedBEns(kron)|30.9$\\cdot$M|7.1$\\cdot$M$^2$|17.8$\\cdot$M|\n|FedFisher-KFAC|30.9|13.4|17.8|\n|RegMean|26.0|5.1|7.3|\n|FisherMerge|26.2|0.9|7.31|\n|DENSE|25.9|170.2|3.65|\n|OTFusion|25.9|1.1|3.65|\n \n**2) Validation set usage:**\nWe thank the reviewer for the interesting question. First, we note that FedFisher, RegMean, DENSE, OTFusion rely on a validation set at the server, both in their original papers and in our experiments. Our method uses a validation set during the server's gradient descent runs, to find the parameter configuration that achieves the best validation performance (Sec. 5.3). We have performed an additional experiment where no validation dataset is used and the number of server epochs was fixed to 100. Remarkably, FedBEns still outperforms the considered baselines, which continue to take advantage of the validation dataset, with only a minor drop in accuracy (up to a few p.p.; see Table 1 in the paper and the supplementary table below). We will highlight this important point more clearly in the revised version. \n\n**FedBEns, 5 mixtures Kron, NO VALIDATION**\n|Dataset| Accuracy [%] (Average, 3 seeds)|\n|-|-|\n|**Alpha = 0.05**||\n|FMNIST|59.15|\n|SVHN| 70.26|\n|CIFAR10|48.81|\n|**Alpha = 0.4**||\n|FMNIST|82.16|\n|SVHN|84.69|\n|CIFAR10|60.00|\n \n**3) Feature heterogeneity:** \nWe are not sure how to interpret the expression “feature heterogeneity.” If it refers to structurally different data across clients—e.g., clients having different numbers or types of features—this would require a different model architecture per client. Since our training method assumes a shared model architecture across all clients, it is not directly applicable to such a setting. In contrast, our method explicitly assumes *statistical heterogeneity* across client datasets—that is, data drawn from different underlying distributions. We do not impose strong assumptions on these distributions beyond conditional independence. If the reviewer refers to *domain shifts* after global model training, our method remains applicable in certain scenarios, such as covariate shift. In fact, FedBEns can be particularly beneficial in this context: when new clients join the federation with different data distributions, the approximate posterior produced by FedBEns can serve as a prior. This acts as a regularizer, helping to prevent the model from drifting too far when fine-tuned on data that differs significantly from the original training distribution. We believe that exploring this aspect more thoroughly would be an interesting direction for future work.\n \n**4) Learnable Weights:** \nA principled way to assign the weights would be to set them equal to the model evidence of each Gaussian component. We experimented with this approach and found that performance remained essentially unchanged, as the model’s evidences were comparable to each other, in line with what has been found in [1]. Another possibility would be to let the server weigh each distribution, or all distributions belonging to a given client, based on the predictive performance of the model trained by that client. However, this approach may require the server to have explicit access to a dataset while, in our approach, it is not strictly needed (as discussed also in answer 2). \n \n[1] Immer et al. “Scalable marginal likelihood estimation for model selection in deep learning”, ICML ‘21\n \n**5) Statistical Significance:** all results (Table 1 and 2) are statistically significant, p-value<0.05, computed with a paired t-test between our models and the best-performing competitor.", "claims_and_evidence": "Employing GMMs to model multi-modality of local loss for each client is straightforward and the proposed method is solid. Extensive numerical experiments support its validity.", "methods_and_evaluation_criteria": "The method uses a mixture of Laplace approximations, which aligns with the goal of capturing multimodal posteriors in one-shot FL. Evaluation on standard datasets (FashionMNIST, SVHN, CIFAR10) with Dirichlet sampling for heterogeneity is appropriate.", "theoretical_claims": "Proposition 3.1 (Equation 2) claims the global posterior can be derived by combining local posteriors under conditional independence and same-prior assumptions. The proof (Appendix A) is correct, using Bayes’ theorem and factorization of likelihoods. Extension to unimodal to multimodal posteriors is straght-forward.", "experimental_designs_or_analyses": "Experiments (Tables 1-4) vary heterogeneity and client numbers, with 5 seeds for robustness. The ablation studies on mixture size, Hessian approximations, temperature, and prior variance (Figures 2-5) are sound. especially, an ablative study for the number of mixture components validates the superiority of employing multimodality instead of unimodal gaussian. However, the lack of statistical significance tests (e.g., t-tests) for accuracy comparisons may limit the reliability of superiority claims. In addition, the assumption of having validation data at the server is unrealistic in federated learning (FL), as FL typically assumes no raw training data exists at the server. If such data is available, it should be used for training rather than just validation.", "relation_to_broader_scientific_literature": "FedBEns builds on Bayesian FL by focusing on one-shot learning and multimodal posteriors, unlike FedPA’s multi-round MCMC approach. It extends Laplace approximations (MacKay, 1992) to mixtures, similar to the work in Eschenhagen et al. (2021). It contrasts with unimodal loss-based methods like FedFisher (Jhunjhunwala et al., 2024), emphasizing multimodal loss landscapes.", "essential_references_not_discussed": "Please refer to \"Relation to Broader Scientific Literature\" Section", "other_strengths_and_weaknesses": "The claim of computational efficiency is weakly supported—while a trade-off with mixture size is mentioned, no runtime or communication cost analysis is provided.", "other_comments_or_suggestions": "- Validation Set Usage: Why use a server-side validation set for hyperparameter tuning, given FL’s privacy constraints? If removed, how would performance change?\n\n- Computational Cost: Can you quantify the runtime and communication costs of FedBEns compared to baselines?\n\n- Does the proposed method work under feature hetergeneity such as domain shift?", "questions_for_authors": "- Learnable Mixture Weights: Is it possible to make each mixture component weight learnable?", "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": "", "extra_scores": {}}, {"review_id": "9HHH2hIZ3n", "reviewer": "Reviewer_b5tc", "summary": "This paper focuses on one-shot Federated Learning, where the model is aggregated in a single communication round. The authors provide an analysis through the lens of Bayesian inference and then propose a method to leverage the inherent multimodality of local loss functions to find better global models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1) Additional server-side computation:**\nWe appreciate the reviewer’s observation—this is indeed an important point. Our method does entail increased server-side computation as the number of mixture components grows. However, we would also like to clarify that some of the baselines considered can be similarly demanding in terms of server-side resources. For instance, DENSE involves both synthetic data generation and knowledge distillation at the server, which results in a training time comparable to that of FedBEns with 5 mixture components.\nFurthermore, we would like to highlight that FedBEns outperforms DENSE even when using fewer mixture components (compare the results for FedBEns in Figure 2 with those in Table 1 in the manuscript) while requiring less server-side training time. To illustrate this more concretely, we provide below a table—also to be included in the paper—that reports the server training time for the CIFAR-10 experiment with 5 clients. The number of mixture components used in FedBEns is denoted as M. \n \n|Method|Server average execution time (seconds)|\n|-|-|\n|FedBEns (Kron)|M=5 173.5, M=4: 113.5, M=3: 65.9, M=2: 28.5, M=1: 7.1|\n|DENSE|170.2|\n|FedFisher|13.4|\n|Regmean|5.1|\n|OTFusion|1.1|\n|FisherMerge|0.9|\n \nOur approach increases server computation but improves predictive performance. Whether this additional computational cost is worth it depends on the practical context, and it can be justified if the server has enough computational power.\n\nFor an additional discussion on the computational/communication cost, please refer to answer 1) to reviewer wHXT.   \n\n**2) Figure 1, clarification:** \nThe two subfigures in Figure 1 illustrate different functions, as also indicated by the respective titles at the top of each subfigure. The left plot shows the local loss functions for client 1 (orange curves) and client 2 (blue curves), each with its own optimum (represented by orange and blue points, respectively). The right plot depicts the global loss function (green curve), obtained using Eq. (2) by combining the losses of the two clients from the left plot. The key takeaway is that, when estimating the global loss, capturing the secondary optimum is more crucial than the primary one, emphasizing the need for a multi-modal approach. We will modify the figure caption to make it clearer.\n \n**3) Faster convergence when trained on the server?:** \nIf we have correctly understood the reviewer’s point, our method can indeed also be used in a centralized setting—that is, where all training data is available on the server. However, in this case, the server’s role in merging different posteriors becomes redundant, as there is no need to combine client-specific posteriors. Naturally, centralized training also eliminates communication overhead, providing an additional efficiency advantage. On the other hand, the federated approach could benefit from the inherent parallelism of distributed data processing. Overall, the relative training speed between centralized and federated setups depends on the interplay of these three factors, as well as potential differences in the hardware available in each case.\nAs an illustrative example, in a centralized setting and using the same hardware described in Sec. 5.3, on CIFAR10, computing one mixture component over the entire dataset would require 129.5 seconds for training the Neural Network, plus 25 seconds for computing the Hessian. In a federated setting, the time to compute a client’s mixture component would decrease proportionally with the number of clients, due to parallelism. However, aggregating the results incurs a cost that scales linearly with the number of clients, as computing the global log-posterior in Eq. (4) requires summing over the clients' posteriors. Assuming identical hardware at the server and at the clients and ignoring communication delays, the federated approach on CIFAR10 and 5 mixtures is more efficient up to $\\sim 20$ clients. This threshold decreases when communication overhead is included or when the centralized server has superior computational capabilities. \n\nLast, also in a centralized setting, accuracy increases with mixtures. As an example we report the average accuracy (3 seeds) for CIFAR10 as a function of the number of mixtures M, with the same training details reported in the paper: M=1: 79.48%, M=2: 82.11%, M=3: 83.28%, M=4: 84.99%, M=5: 85.46%.", "claims_and_evidence": "na", "methods_and_evaluation_criteria": "na", "theoretical_claims": "na", "experimental_designs_or_analyses": "na", "relation_to_broader_scientific_literature": "na", "essential_references_not_discussed": "na", "other_strengths_and_weaknesses": "Overall, the results look promising. However, it would be beneficial to include discussions on the following aspects:\n\n- computation: The algorithm involves additional server-side training. It would be fair to compare its computational cost with other methods, as some model fusion methods don't require additional training. \n\n-  Figure 1 is not very clear or easy to understand. For example, where is the global loss in the left figure? Are both figures plotted in the same way? Clarifying these points would improve readability.\n\n-  It would be interesting to see whether this method converges faster when trained on the server.", "other_comments_or_suggestions": "na", "questions_for_authors": "na", "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": "", "extra_scores": {}}, {"review_id": "t1Z39kGXVq", "reviewer": "Reviewer_dxFK", "summary": "FedBEns proposes a one-shot federated learning utilizing a Bayesian ensemble approach. Unlike standard FL methods that simply rely on averaging, FedBEns combine client models using a mixture of Laplace approximations to model multimodal local posteriors. Empirically, FEDBEns demonstrates superior performance on benchmark datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**1) Computational/Communication overhead quantification:** See answer 1) to reviewer wHXT.\n \n**2) Computational/Communication trade-offs and practical applicability:**\nThe proposed method is inherently more computationally intensive, as it is based on an ensemble of M models. Its applicability depends on the specific application setting and the base model (e.g., NN architecture) used for training in the federated setting. Regarding the trade-off between computational costs and prediction performance, it is worth mentioning that our proposal substantially improves over existing methods with just a couple of mixtures (see Table 1 and Figure 2 of the manuscript).\n \nThe reviewer suggested additional exploration of the computational and communication trade-offs. We provide the following insights in response.\n\n***2.1) Comparison with FedFisher:***\nFedFisher is our direct competitor. In Appendix B.4, we extend FedFisher to a multi-round setting. By letting FedFisher train over 5 rounds, its total communication and computation costs equal those of FedBEns with 5 mixture components. Our method shows to be competitive despite relying on a single communication round. We performed an additional experiment in which we equalized the client computation with FedFisher by running FedBEns with 5 mixture components on clients’ training datasets 5 times smaller. Remarkably, FedBEns performs better than FedFisher in most scenarios (see the table below and FedFisher K-FAC results in Table 1 of the manuscript for comparison).\n\n**FedBEns, 5 mixtures Kron, REDUCED DATASETS**\n| Dataset  | Accuracy [%] (Avg ± STD, 3 seeds) |\n|-|-|\n| **Alpha = 0.05**  |  |\n| FMNIST   | 50.87 ±  4.46  |\n| SVHN   | 59,12 ± 1.42   |\n| CIFAR10 | 45.38  ± 0.62 |\n| **Alpha = 0.4**   |  |\n| FMNIST | 77.95 ± 0.31  |\n| SVHN   | 81.15  ±  0.25  |\n| CIFAR10 | 53.82  ±  1.16 |\n\n***2.2) Other tasks:***\nTo further explore the accuracy-costs tradeoff for practical applications, we conducted a more challenging experiment with a ResNet-18 on CIFAR100. We followed the experimental setting of the paper; however, we limited FedBEns to only 2 mixtures and Kronecker factorization. Local training lasted for 100 epochs. The table below shows that at the cost of doubling the computation/communications cost w.r.t. FedFisher, FedBEns achieves a higher gain in accuracy (doubled or even more).\n| Metric  | FedBEns M=2 | FedFisher | RegMean | DENSE | OTFusion | Fisher-Merge |\n|:-|-:|-:|-:|-:|-:|-:|\n| **Accuracy**, CIFAR100 (alpha=0.05) | 8.56±0.47 | 3.14±0.23 | 3.01±0.54 | 5.88±1.07 | 4.79±0.19 | 2.81±0.33 |\n| **Accuracy**, CIFAR100 (alpha=0.4)| 16.95±0.67 | 4.79±0.31 | 4.85±1.16 | 13.92±0.54 | 4.74±0.56 | 3.87±0.48 |\n| **Communication Costs (MB)** | 716 | 358 | 93 | 47 | 47 | 94 |\n| **Local runtime (min)** | 28.0 | 14.0 | 9.3 | 9.2 | 9.2 | 10.1 |\n| **Server runtime (min)** | 2.3 | 1.2 | 0.6 | 10.4 | 0.2 | 0.3 |\n\n***2.3) System heterogeneity:***\nAnother comment regarding the practical deployment of our approach is that the number of mixtures can be customized on the basis of each client's computational power. As an illustrative example, we considered CIFAR-10 with 5 clients: 3 employ a single Laplace approximation (M=1) and 2 are able to exploit three mixtures (M=3). In both cases, the final model returned to the clients is composed of a single model. The additional information coming from the clients with more computational power benefits the entire federation and leads the server to find a better global model, especially in heterogeneous settings. The table below shows the mean accuracy over five seeds for FedBEns in two settings: all clients use a single component (M=1 ALL) versus the customized setup described above.\n \n|Dataset|Accuracy [%] (Avg ± STD, 3 seeds)||\n|-|-|-|\n| | **M=1 ALL**| **M CUSTOMIZED**|\n|**Alpha = 0.05**| | |\n|FMNIST| 57.02 ± 4.73| 59.97 ± 1.44|\n|SVHN| 66.11 ± 0.77| 68.39 ± 1.57|\n|CIFAR10| 39.89 ± 1.71| 42.32 ± 2.25|\n|**Alpha = 0.4**| | |\n|FMNIST| 77.07 ± 0.29|78.31 ± 1.76|\n|SVHN| 82.86 ± 0.21|82.75 ± 1.41|\n|CIFAR10| 53.97 ± 0.86|54.68 ± 1.88|\n\n\n In the case of acceptance, we will integrate these additional results into the paper.\n\n**3) Quite limited novelty:** \nWe agree with the reviewer that our proposal relies on concepts commonly used in the Bayesian learning landscape. However, we would like to emphasize that our key contribution lies in recognizing how these ideas can significantly improve FL, specifically in a One-Shot setting. In particular, we highlight the importance of explicitly capturing different modes of the posterior rather than framing the FL task purely as an optimization problem. Additionally, we believe that, by building upon a well-studied theoretical framework, our method, also thanks to its inherent simplicity, is both principled and robust. Moreover it brings, as pointed out by the reviewer, very strong empirical results.", "claims_and_evidence": "The claims are generally supported by empirical evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria (datasets like FMNIST, SVHN, and CIFAR10) are appropriate and standard in federated learning literature. Experiments varying the heterogeneity parameter and client counts effectively illustrate the method’s robustness in diverse scenarios.", "theoretical_claims": "This paper does not provide any new theoretical results.", "experimental_designs_or_analyses": "Experimental designs are sound, with clear baselines chosen from different categories of one-shot federated learning (FedFisher, RegMean, DENSE, OTFusion, FisherMerge). The analysis across varying levels of data heterogeneity and client numbers is comprehensive and methodologically robust.", "relation_to_broader_scientific_literature": "FedBEns extends the literature on Bayesian federated learning methods, particularly enhancing one-shot federated learning. It directly builds upon prior works (FedBE, FedPA, FedPop, FedHB), explicitly modeling multimodal local posteriors through Laplace mixture approximations. However, the approach is quite incremental, essentially applying known Bayesian ensemble ideas and aggregation explicitly to the one-shot federated learning setting. While simple, no other work has explicitly pursued this direction in one-shot FL, to the best of my knowledge.", "essential_references_not_discussed": "No critical references appear to be missing that would significantly alter the context or understanding of the proposed contributions.", "other_strengths_and_weaknesses": "**Strengths:**\n\n- Clear justification for the multimodal Bayesian approach.\n\n- Empirical results demonstrating significant accuracy improvements over baselines.\n\n- Good evaluation across multiple datasets and varying heterogeneity levels.\n\n\n**Weaknesses:** \n\n- Computationally intensive due to multiple local models required per client, significantly impacting scalability, particularly in large-scale or resource-constrained settings.\n- Quite limited novelty: The paper’s contribution is essentially an explicit application of existing Bayesian ensemble and aggregation ideas to one-shot FL without substantial conceptual innovation. If empirical results were not as strong, this lack of novelty would likely not warrant a paper.\n- Lack of detailed analysis on practical constraints, such as communication overhead and latency, limiting the clear demonstration of practical applicability in real-world scenarios.", "other_comments_or_suggestions": "- Readability and clarity of the algorithm and experimental setups are very clear.\n\n- Further exploring and clearly quantifying computational and communication trade-offs in practical deployments would strengthen the paper's applicability and appeal.", "questions_for_authors": "- Could you further clarify the scalability and communication overhead compared to simpler baseline methods?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper uses a mixture of (Laplace) Gaussians in one-shot federated learning. Reviewers agree that the multimodal nature of local posteriors is well-motivated, that the method itself is straightforward, and that the experiments are performed well (good benchmarks and ablations). Concerns such as computation / communication cost, and some additional experiments, were addressed by the authors during rebuttal. The authors have committed to adding these results and discussions into the paper. Overall, this is a solid contribution to ICML that takes well-known methods and then applies them to a new setting, with good motivation and experiments. 

Reviewer dxFK asked about the benefit of one-shot federated learning (but M times the cost) vs M shots. I agree with the authors that this is still an interesting and important setting.
Decision: Accept (poster)