Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation
Authors: Ting Wei
Biao Mei
Junliang Lyu
Renquan Zhang
Feng Zhou
Yifan Sun
Venue: NeurIPS 2025 poster ()
Keywords: Personalized Bayesian Federated Learning
Stein Variational Gradient Descent
Wasserstein Barycenter Aggregation
Primary area: general_machine_learning

Abstract: Personalized Bayesian federated learning (PBFL) handles non-i.i.d. client data and quantifies uncertainty by combining personalization with Bayesian inference. However, current PBFL methods face two main limitations: posterior inference on clients often assumes restrictive parametric forms, and server-side posterior aggregation typically relies on naive parameter averaging. To overcome these issues, we propose FedWBA, a novel PBFL method that enhances both local inference and global aggregation. At the client level, we use particle-based variational inference for nonparametric posterior representation. At the server level, we introduce particle-based Wasserstein barycenter aggregation, offering a more geometrically meaningful approach. Theoretically, we provide local and global convergence guarantees for FedWBA. Locally, we prove a KL divergence decrease lower bound per iteration for variational inference convergence. Globally, we show that the Wasserstein barycenter converges to the true parameter as the client data size increases. Empirically, experiments show that FedWBA outperforms baselines in prediction accuracy, uncertainty calibration, and convergence rate, with ablation studies confirming its robustness.

Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation
Ting Wei School of Statistics, Renmin University of China Beijing, China veiting1006êruc. edu. cn
Biao Mei School of Statistics, Renmin University of China Beijing, China 2022103705@ruc. edu
Junliang Lyu Guanghua School of Management,
Renquan Zhang
School of Mathematics Science, Dalian University of Technology Dalian, China zhangrenquan@dlut. edu, cn
Peking University Beijing, China #yjunliang02110rus edu. on
Feng Zhou'
Yifan Sun Center for Applied Statistics and School of Statistics, Renmin University of China
Center for Applied Statistics and School of Statistics,
Renmin University of China
Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing Beijing, China
Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing Beijing, China sunyifangruc. edu. cn
feng. monggucc wew.uuc..ed..oom edu.cn
Abstract
Personalized Bayesian federated learning (PBFL) handles non-i.idd client data and quantifies uncertainty by combining personalization with Bayesian inference. However, existing PBFL methods face two limitations: restrictive parametric as- sumptions in client posterior inference and naive parameter averaging for server aggregation. To overcome these issues, we propose FedWBA, a novel PBFL method that enhances both local inference and global aggregation. At the client level, we use particle-based variational inference for nonparametric posterior repre- sentation. At the server level, introduce particle-based Wasserstein barycenter aggregation, offering a more geometrically meaningful approach. Theoretically, provide local and global convergence guarantees for FedWBA. Locally, we prove a KL. divergence decrease lower bound per iteration for variational inference convergence. Globally, true parameter as the client data size increases. Empirically, experiments show that FedWBA outperforms baselines in prediction accuracy, uncertainty calibration, and convergence rate, with ablation studies confirming its robustness.
show that the Wasserstein barycenter converges to the
Introduction
Federated learning (FL) enables clients 443, making particularly advantageous privacy-sensitive domains such as finance 33)
collaborative model training across decentralized
"Corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

healthcare [35] and IoT [32. Conventional FL. however, faces two key challenges from non-i.i.d client data: convergence degradation due to client drift and (2) suboptimal global model perfor- mance 226
40 These limitations impede deployment in aafety-critical fields. Personalized FI. (PFL.)
addresses these limitations through client-specific adaptation techniques including transfer learning
13 12, meta-learning [20] [4], and parameter decoupling [O]. Nevertheless, conventional PFL
frameworks relying on frequentist approaches remain vulnerable to overfitting and lack uncertainty quantification crucial shortcomings for safety-critical applications. This motivates personalized
Bayesian FL (PBFL) [T] which integrates Bayesian inference with PFL via two key mechanisms:
prior regularization against overfitting and posterior inference for uncertainty quantification [45] 28
Although PBFL has many advantages, existing methods face two main challenges: (1) Performing
posterior inference on clients is challenging due to the lack of analytical solutions for model parameter
posteriors, requiring approximation methods such as variational inference (VT) [9] To simplify inference, many studies assume parameterized variational distribution, such as Gaussian, to obtain
a tractable evidence lower bound (ELBO) 45 33 46). However, this parameterization introduces
S (2) Performing aggregation on the central server is challenging because aggregating the posteriors from clients is not straightforward. To simplify the aggregation, many works upload the parameters of variational distributions from the clients and then average these parameters on the server to obtain the updated global prior [45] 77. However, this is problematic since information geometry [ indicates that distributions lie on manifold, where "averaging" differs from that in parameter space.
error,, as the true posterior is often complex and non- Gan--aussian
To address the aforementioned issues, we propose a novel PBFL method, federated learning with Wasserstein barycenter aggregation (Fe) which enhances both local posterior inference and global aggregation. At the local level, we avoid parameterizing the variational distribution and instead employ particle-based VI. It represents the posterior with a set of particles, offering more nonparametric flexibility than parameterized VI. At the global level, we propose particle-based Wasserstein barycenter aggregation method. It finds the barycenter of local posteriors on a manifold induced by the Wasserstein distance, combining client particles into a new set for the global prior. This method has a clearer geometric interpretation than parameter averaging. Theoretically, we guarantee the convergence of the proposed method. Our main contributions are follows:
(1) We propose a novel PBFL method called FedWBA. At the local level, we use particle-based for greater nonparametric flexibility, while at the global level, we introduce a particle-based Wasserstein barycenter aggregation, which more geometrically meaningful.
(2) We provide theoretical convergence guarantees at both the local and global levels. Locally, we prove lower bound on the Kullback-Leibler (KL) divergence decrease per iteration, ensuring the convergence of variational inference. Globally, we show that as client data size approaches infinity, the Wasserstein barycenter converges the true parameter.
(3) Comprehensive experiments demonstrate the superiority of FedWBA in prediction accuracy, uncertainty calibration, and convergence rate compared to baselines. Additionally, ablation studies evaluate the robustness of our approach w.r.t. different components.
2 Related Works
Personalized Federated Learning can be taxonomized into two principal paradigms according to [88] The first personalizes global models through client-specific fine-tuning. Regularization- based methods, such as FedProx (25] add a proximal term to the loss function to measure model discrepancics, providing an intuitive and effective solution. Meta-lcarning approaches, such as Per-FedAvg [14) learn model initializations for rapid client adaptation, later enhanced by pFedMc 37 through Moreau envelope optimization. The second strategy focuses on personalizing the models by modifying the FL aggregation process, aligning with our PFL goal. Parameter-decoupliin methods, exemplified by FedPer [5), decompose deep neural networks into shared base layers for feature extraction and client-specifie heads for task adaptation. Multi-task learning methods, such as FedAMP [8] employs attention- based similarity weights, though sensitive to data quality. Our method can be considered as a form of the meta-learning approach.
Bayesian Federated Learning focuses on the inference and aggregation of posterior distributions over model parameters across clients. Depending on how the posterior is represented, BFL methods

---

can be broadly categorized into parametric and nonparametric approaches. Parametric methods assume specific form for the posterior-typically Gaussian for tractability. Representative works include pFedBayes 45 FOLA 27,, PFedVEM 46, FedPA
3 FedEP [77. and BA-BFL. [19] Although BA-BFI. introduces geometric aggregation perspectives, still relies on restrictive Gaussian posterior assumptions. FedHB 23 extends this to Gaussian mixtures, yet remains constrained by parametric forms. contrast, nonparametric methods avoid strong assumptions on the posterior, enabling more flexible representations. For instance, FedPPD 7 employs MCMC to characterize local posteriors, while distributed SVGD (DSVGD) [22] adopts a particle-based approach. However, DSVGD requires three SVGD updates per communication round, which significantly increases computational cost. Our method advances this line of work by enabling single-round particle transport while preserving flexible posterior representations free from Gaussian constraints.
Preliminaries
In this section, we provide an overview of Stein variational gradient descent and Wasserstein barycen- ter aggregation.
3.1 Stein Variational Gradient Descent
SVGD approximates target distribution p(x) by iteratively transforming particles from initial distri- bution (xx through t(x) the KL. divergence reduction rate. The objective is to find a that minimizes (I{q (x))((x)) Computing the derivative of the KL w.r.t. e at 0 yields a closed-form solution:
cιb(x), where is the step size and Φb-) R - RM maximizes
(xKL(gr)())]l(()) == [rtcc
where (,) Vx (())(() + V (xx.. To maximize the rate of decrease in KL divergence, we aim to select & such that By(x) [trace A (x))) is as large as possible.
The constrained optimization is resolved via reproducing kernel Hilbert spaces (RKHS) 29 Let k(,,-) be positive-definit kernel defining an RKHS H. with HDD = 7 x ... x H denoting its D- dimensional product space for vector-valued functions f (fi,.... fD) where f. H. Constraining Φ e HD with || < 1, the steepest descent direction admits the following analytic expression:
(¹(-) = (b(1)/11bll (00)=((()//1111111
(+) = Ex (Apk() -]]-
we approximate the distribution qxx) using a finite set of particles located at {xi} the iterative algorithm at the 1-th iteration can be formulated as follows:
(111)
k(x)0) x V logp(x) (V<((xxx")))
The first term in the update formula attracts particles towards high-probabiliin regions, while the second term acts as a repulsive force, preventing particle clustering and ensuring a more uniform exploration of the distribution's support.
3.2 Wasserstein Barycenter Aggregation
The Wasserstein distance from optimal transport theory provides geometric distribution comparison through minimal transport cost, with the Wasserstein case defined as:
(2(p(x), qxx))
min Bxxxx))
x'112).
(1)
*ell
where I (xx x') represents a joint distribution with marginals p and q respectively, and II denotes the set of all such joint distributions (39. The Wasserstein distance has a clear physical interpretation. The norm term ||x - x') x'I represents the cost of transporting a unit mass from x to x', and <xx, x') is a transport plan that specifies the amount of mass move from x to x'. Therefore, the Wasserstein distance is the minimal transportation cost, and the optimal n² (x, x') the plan that achieves this
minimal cost among all plans.
3

---

Server
(0.
(6r
(a.)
(9
Σ (o)
10
{0
(oud
Client
Client 2
Client K
wrg max R (o: (Pua)
D2
Dg
(L( (()
Figure 1: Overview of FedWBA. Left: System diagram. Clients upload local posterior particles to server for aggregation, server updates global prior particles and redistributes them to clients. Right: Local posterior particles from maximizing local ELBO, global prior particles as Wasserstein barycenter of K local posteriors.
The Wasserstein barycenter provides geometric mean of distributions by minimizing the weighted sum of Wasserstein distances to given distributions [2]. For set of probability distributions (p: i=1 i=l wi 1, 2 0, the Wasserstein barycenter is defined as:
K
with weights [w;} satisfying > K
K
P arg min
wa W2
pEP
K
where P the of probability distributions. The Wasserstein barycenter has the advantage of clear geometric interpretation, but a drawback the challenging computation of the Wasserstein distance. This difficulty arises because identifying the optimal transport plan between continuous distributions is complex and analytically tractable only in special cases [2]. Nevertheless, it is worth noting that when the given distributions are discrete, the Wasserstein distance can be estimated via linear programming, making the computation of the barycenter tractable [6 [1]
Methodology
In this section, we introduce our federated learning with Wasserstein barycenter aggregation (Fed- WBA) approach. As illustrated in Figure I the training process is iterative: clients begin by down- loading the global prior from the server. Using this global prior, each client updates its local posterior, represented as a set of particles, via SVGD. The updated particles are then uploaded to the server for aggregation, performed using Wasserstein barycenter aggregation. This iterative process continues until convergence. Each client adapts the global prior to its local data, ensuring personalization, while the global prior, formed by aggregating local posteriors, reflects overall patterns.
4.1 Problem Definition
We consider distributed system consisting single server and clients. Assume all clients share the same model with parameters 0 R. and each of the K clients has its dataset [D We represent the global prior the server using set of particles {0 and similarly, use (0x, /==1 to represent the local posterior on the k-th client. Distributions without hat denote continuous distributions, while those with hat represent the corresponding empirical (discrete) distributions expressed through particles.
Local Posterior via Stein Variational Gradient Descent
In each communication round, on the k-th client, global prior over the model parameters 7(0) downloaded from the server. Our goal is to adapt this prior to the local data using Bayes' rule, resulting in the posterior distribution of the model parameters:
p D&) & p(Px ())
where (0) is the prior, (DDs 0) the categorical likelihood, and p(0 (Dx) the posterior. Since the likelihood is usually parameterized by neural network, is typically non-conjugate to the prior.

---

As result, the posterior generally lacks an analytical expression. To address this, approximate inference methods are employed approximate the posterior. VI among the most widely used approaches for this purpose. Specifically, in VI, the posterior p(O (D&) is approximated by variational distribution q (0). The optimal variational distribution obtained by minimizing the KI. divergence between them equivalently, by maximizing the ELBO:
F 0)) = (o) p(Dx - (L( P(0)).
commonly impose parametric assumptions (e.g., Gaussian) on qu (())) for ELBO tractability, yet suffer from approximation mismatch when the true posterior deviates from assumed forms [8. Inspired by [22]. we advocate for using the SVGD method, which employs set of particles to flexibly represent the variational distribution. This approach eliminates the need for specific parametric assumptions. Specifically,
Prior methods 45
33 46
assume that the variational distribution qu (0) is represented by a set of partícles {0x.)) j=1. which are iteratively updated according to SVGD. The particle update rule in the (1)-th iteration given by:
N
((+1) R. k,
@
Σ
Vo logp(0 Dz) +ook(0,0))) 0
(2)
N
where V lologp(0 P&)) Vo ( + Vo P(D according to Bayes' rule, k(:;)) is positive definite kernel corresponding to a RKHS.
As the number of particles N increases, the empirical distribution of { asymptotically converges to the optimal variational distribution in the corresponding RKHS; when we use only a single particle N 1, the maximum posteriori (MAP) estimate be obtained [29]
4.3 Global Prior via Wasserstein Barycenter Aggregation
In each communication round, the server needs to aggregate the posterior distributions uploaded by multiple clients into global prior, which challenging procedure. Previous works 45 7 used naive approach by averaging the parameters of the uploaded posteriors and treating the resulting distribution as the aggregated one. We argue that this aggregation method is problematic because, according to information geometry [44], distributions lie on a manifold, where the "average" on the manifold differs from the "average" in the parameter space. visual illustration of this difference can be seen in Figure (Appendix A..
To address this issue, propose using Wasserstein barycenter aggregation, leveraging the geometric properties of the manifold where the local posteriors reside. Unlike previous works, our approach replaces parameter averaging with Wasserstein distance-based averaging, offering a clearer geometric interpretation. Specifically, after applying SVGD, the variational distribution yı (0) on the k-th client approximated as an empirical measure (x(()) using particles J. similarly. the global prior p(0) is also approximated as an empirical measure p()) using particles {0,)N
N 8e. (0).
(0)
S0.. (0), (0), 70)
N
N
where 50(-) the Dirac measure, To define the Wasserstein distance between p( ()) and (0), we define distance matrix M, NxN between the global prior particles and local posterior particles:
MR.., = 100 1x
where i.j represent the entry in the i-th row and j-th column. We then define the transport plan as matrix Т₃ & X which is discrete version the joint distribution (..-) in Equation (I). Then the Wasserstein distance between p(0) and (f) can be written as:
1 W (x) min (M., Tt) s.t. T.1 1 T -INII. Te 0. T,ET
where T is the set of all possible T₁. and (::)) is the Frobenius inner product between two matrices. The constraints of Tr correspond to the marginal properties of .(.)). Let Tद denote the optimal transportation plan in Equation (3). It is easy that this is a linear programming problem, which
5

---

can be solved using general optimization packages. However, more specialized algorithms,, such as Orlin's algorithm, also exist 334 15].
After obtaining T we can update the global prior as the Wasserstein barycenter of K local posteriors:
(0) arg min 10
(Ms. TT)
1 I N be the stack of global prior particles and O₁ I N to be the stack of local posterior particles on the k-th client, it is easy to prove that the objective function in Equation is quadratic w.r.t. θ so the optimal O has an analytical solution (see proof in Appendix
Denoting [
[NN]] 10........ 0
= K
Σ k=1
Oiiig((NN
(5)
After obtaining the optimal global prior particles {0) Ji=1: a continuous global prior is required on the clients for SVGD updates, as the gradient of the oog-prior must be computed in Equation (2). To achieve this, we employ kernel density estimation (KDE) to construct continuous global prior:
p p0(0)
Σ 10,00)
(6)
where i(.,)) is normalized kernel in KDE, which is different from the kernel k(.;)) in SVGD.
4.4 Algorithm
In summary, at the client level, each client receives the global prior particles and reconstructs continuous prior using Equation (6), then updates the local posterior particles via Equation (2). At the server level, the global prior particles are updated by aggregating the uploaded local posterior particles using Equation (5). We term this method FedWBA, with its pseudocode provided in Appendix
Theoretical Analysis
In this section, we aim to establish the convergence properties of our proposed method. Locally, we demonstrate that the ELBO increases with each iteration. By providing lower bound for the difference in ELBO between consecutive iterations, we ensure that as the iterative process of SVGD proceeds, the variational distribution represented by particles getting closer to the true posterior. Globally, we show that as the data size on each client tends to infinity, the Wasserstein barycenter converges to a delta measure centered at the true parameter. This convergence implies that the aggregated global prior effectively captures the overall patterns. The proofs are presented in Appendices D and E. respectively.
Assumption 5.1. Vet(O) is positive definite matrix, where the eigenvalues e, 0 for 1....
Assumption 5.2. is small enough s.t. ||e (eVoo(()) 1
Theorem 5.3. Under Assumptions 5.1 and 5.2 given SVGD iteration I, with client k scheduled, the satisfies the inequality:
increase in the ELBO from iteration |
F( (0) 2 cD(a (0),p(0 Dw)), (( (0)) ((++)())))
(7)
where e is the step size in SVGD, 5(0 DD) denotes the unnormalized posterior distribution and D(g.p) stands for the kernelized Stein discrepancy:
D(g.p)
max (E,(o) (race(A,,(())) s.t. | I 11. $EHD
Theorem 5.3 shows that the ELBO increases with each iteration, which equivalently ensures that the KI. divergence between the variational distribution and the true posterior decreases, with the lower bound of its decrease determined by the kernelized Stein discrepancy. This implies that ga (0) converges to p(0 D&) the iterations progress.

---

Assumption 5.4. E is compact space in metric. and dg is an interior point of E. All clients have equal data size, i.e., St S/ for ....., K, where Sg the data size of client &t and the total data size.
Assumption 5.3. For any 0, E and 1,.... K, there exist positive constants a and CL, s.t. the following inequality holds:
2 (0, 0' > CLP (0,000
Assumption 5.6. There exist constants Cj 0, 0 C2 2 a function (u, > that is non-increasing in a E R+ and non-decreasing in r IR For all k I..... K, any u, r > and all sufficiently large s, the generalized bracketed entropy Hj satisfies
HD (n(u, (n (D(0)) E E, (0, hr), (u,r) for all k: = =1,..... -,...,, K;
and
Gr
v (u, du C2ver².
1212/212
Assumption 5.7. There exist positive constants and or s.t., uniformly over k 1...., K,
II 1
p(Dxs - S log 2 $ p(D
E Poo Σ 1=
exp
k log
2 exp( K log² s),
where log* max(log x, 0) for O.
Assumption 5.8. The metric satisfies the following property: for any N EI,, UN. 0' E and nonnegative weights Σ W/s =
Σ
w 0'
w.p(O 0').
P
Theorem 5.9. Assuming Assumptions 5.4 5.8 hold for all client posteriors {p(0 Dc)) let $ [Du] denote the data size on each client. Then, as > 00, we have:
1og2/a
6e. (0) = =Op
8 8//aa
where dø, denotes the Dirac delta measure centered at the true parameter Oo. the Op notation is w.E.. the probability measure 0.
p(S) and a a positive constant introduced in Assumption 5.5
Theorem 5.9 shows that the Wasserstein barycenter converges the true parameter Og at rate of up to logarithmic factors, which demonstrates that our global aggregation mechanism effectively combines client information to approximate the true parameter.
6 Experiments
In this section, we utilize four real-world datasets to evaluate the performance of FedWBA in terms of prediction accuracy, uncertainty calibration, and convergence rate. We perform all experiments using a server with GPU (NVIDIA GeForce RTX 4090).
6.1 Experimental Setup
Baselines:
We conduct comprehensive comparisons against state-of-the-art FLJ/FI, methods im-
plemented in the standardized framework [44]. Specifically, we compare our approach against the
following FL methods: (1) FedAvg (31], (2) FedProx (25), (3) Scaffold [21]: PFL methods: (4)
FedPer [5) (5) PerFedAvg [1], (6) pFedME (37; and BFL methods: (7) pFedBayes [45), (8)
VEM 46 (9) PFedGP M and (10) DSVGD 22
"Our code is publicly available at

---

Table 1: Test accuracy (% - SEM) over 50, 100, 200 clients on MNIST. FMNIST, CIFAR-I0 and CIFAR- 1-00. Best and second-best results are bolded and underlined respectively, Dashes denote
DSVGD results unavailable within 10 hours on CIEAR-10/CIFAR-100)
DATASET
METROD
50 CLIENTS
100 CLIENTS
200 CLIENTS
CLIENTS
100 CLIENTS
200 CLIENTS
DATASET
FRnAVG
91.98 0.07
91.76 0.08
90.94
08:24 0.70
55.96 0.21
51.78
PEDPROx
92.12 + 0.08
92.04 :: 0.11
90.82 à: 0.16
60:26 0.42
56.87 + 0.00
58.86 0.00
SCARFOLD
92.90 0.07
92.
90.85 0.11
02.90 0.38
61.42 0.51
60.51
PEDPER
96.33 + 0.02
95.98 0.05
94.19 0.04
0.07 T3.53
68-95 0.21
65.41 0.11
PERFEDAVG
91655
90.61
67.20
62.17
61.20
MNIST
PFEDME
96,70 +: 0.02
95.00 1: 0.00
93.82 4: 0.02
72,62 0.17
71:33 4: 0.15
69.72 1: 0.22
CIFAR-10
PPEDBAYES
95.83
92.77
72.83
62
66.
PEEDVE
97.90 +: 0..5
97.11 01.0
96.42 0.10
73.20 0.20
71.90 4: 0.1
70.10 #: 0.30
PFEDGP
97.69
90.48 0.21
72.61
71.87
DSVGD
96.41 + 0.02
96, =: 0100
95.97 + 0.01
OURS
97.99 0.03
97.36 0.01
96.95 0.01
73:40
72.08 0.14
TEDAYG
84.60 1: 0.22
84.06 1: 01.09
60.14 :: 0.05
25.09 a: 0.14
24.49 a: 0.24
21.53 1: 0.19
PEDPROX
84.65
$3.51
92.95 +
25.34
34.56
34.56
SCAFFOLD
85.49 d: 0.19
63.57 1: 0.10
$2.56 à: 0.16
47.45 0.67
41.62 0.16
41.48 à: 0.50
FEDPER
91.65
90.01
25
52.88
68
PERFEDAVG
90.03 + 0.10
87.40 0.00
85.03 + 0.03
51.42
0.12 56.40
49.02 0.20
35.84 + 0.19
FMNIST
PFEDME
90.32 0.14
89.63
89.19
57.66
CIFAR-100
PFEDBAYES
91.70 1: 0.03
91.53 1: 01.66
90.50 2: 0.06
61:45 a: 0.10
60.58 a: 0.17
55.92 1: 0.31
PEEDVEM
91.80
91.40
90.70 4:
57.57
55.32
40.99
PFEDGP
91.83 1: 0.05
91.51 4: 0.07
90.58 4: 0.15
63.30 4 0.
61:30 a: 0.20
59. 0 0.22
DSVGD
92.41
OURS
92.50 1: 0.06
91.65 d: 0.02 90.97 +: 0.08
64.21
0.02
61.73
0.16
59.22 :: 0.12
Datasets:
To benchmark under realistic non-ii.d. conditions with label skew, evaluate four vision datasets: MNIST, FashionMNIST (FMNIST) (42), CIFAR-10, and CIFAR-100 (24). We adopt the setup from 45 46 I in which each client receives 5 unique labels for MNIST, FMNIST, and CIFAR-10, and 10 labels sampled from distinct superclasses for CIFAR-100.
Setup: We conduct all experiments with 100 communication rounds and 20% client participation
per round, sufficient for algorithm convergence. We evaluate performance with client counts K e (50, 100, 200}, considering that more clients disperse the training data. Following established architectures in prior work, we implement lightweight models: for MNIST and FMNIST, adopting the lungle-hidden-layee MLP from 45 46). and for CIFAR-10/100 deploying the LeNet-style CNN in [T] to accommodate resource constraints.
Hyperparameter: Similar to 29 in all SVGD experiments, we employ the radial basis function (RBF) kernel (00 0o) = exp( |0 0o1l.). 10-0.112
The bandwidth set to h med? / log N, with med being the median of pairwise particle distances in the current iteration. The bandwidth of the Gaussian kemel k(.,-) used for KDE is set to 0.55. We use AdaGrad with momentum to set the learning rate e. Considering communication constraints from uploaded data size, we set the number of particles to 10, balancing computational accuracy and communication overhead in SVGD, mitigating communication bottlenecks without sacrificing much variational inference performance.
6.2 Performance of Prediction
As shown in Table [] the algorithm performance reported under optimal hyperparameter config- urations. Given that the global model acts prior regularizer, we primarily focus on client-side performance; additionally, we designed experiments to verify that the global model's performance is comparable to FL. baselines, as detailed in Appendix F.T Due to DSVGD failing to produce results within 10 hours on CIFAR- 10 and CIFAR-100, we only report its performance on MNIST and FMNIST. As is evident in Table I the proposed method outperforms the existing approaches in nearly all scenarios. This demonstrates FedWBA's strong adaptability to diverse real datasets, delivering consistent performance irrespective of data type or scale. Partial results are copied from 46)..
The performance superiority of FedWBA stems from its unique nonparametric design. In our approach, particles flexibly represent both the local posterior and the global prior, eliminating the need for parametric constraints. This nonparametric representation gives the model strong fitting ability, allowing to better capture data patterns, thus having an edge over baseline methods with parametric assumptions, Also, as Bayesian algorithm, FedWBA excels in shot scenarios. When
8

---

PerFedAvg
PFedBayee
pFedGP
FedWBA
BCB 1ib--uu 0J077
BCB
000002
BCE -0/297
BCE 0.0.9
Acomncy
0.75
0.75
0.75
0.75
0.50
0.50
0.25
0.25
0.25
0.25
0.00
Confilence
Confidlacee
Confislenee
Confilence
Figure 2: Reliability diagrams of top four methods CIFAR-100. The perfect calibration is plotted as a red diagonal, and the actual results are presented as bar charts. The gap between the top of each bar and the red line represents the calibration enror. The ECE is calculated and placed in the corner of the figure. FedWBA BA demonstrates the best calibration performance, ranking first in terms ECE.
HART
Cenan
CINARIN
20
40
10
110
30
10
10
Figure 3: Comparison of convergence rates of different methods on MNIST, FMNIST, CIFAR- 10, and CIFAR-100 with 50 clients. FedWBA exhibits the fastest convergence, with rapid growth in the first 10 communication rounds followed by steady improvement.
the number of clients increases and per client data volume is small, FedWBA features to learn and infer from limited data, resulting in better performance. Moreover, FedWBA has relatively small accuracy fluctuations. This stability is crucial in practical applications, especially in distributed learning with data heterogeneity and limited data, providing more reliable results.
its Bayesian
6.3 Performance of Uncertainty Calibration
Benefiting from Bayesian principles, FedWBA provides superior predictive uncertainty quantification. Under identical experimental conditions as detailed in Section 6.1] we evaluate calibration perfor- mance through reliability diagrams in Figure These diagrams compare model confidence (bars) against perfect calibration (diagonal), with Fed showing closest alignment. We also calculate the expected calibration error (ECE), which, [16] suggests, measures the weighted average of empirical accuracy and model confidence. For the CIFAR-100 dataset, we present only the top-four performing methods; details of the remaining algorithms are provided in Appendix F.2 Results for additional datasets are also detailed in Appendix [22, further corroborating the consistent superiority of our approach across diverse visual recognition tasks.
FedWBA attains a lower ECE by leveraging Bayesian principles that integrate prior knowledge and update beliefs based observed data, which is essential for precise uncertainty quantification. Empowered by SVGD, it adapts well to data characteristics, enabling more accurate uncertainty quantification. In contrast, baseline models usually lack this adaptability, leading subpar calibration and less accurate uncertainty quantification.
6.4 Convergence Rate
We compare the convergence rate of FedWBA with baseline models, updating local parameters 10 times per client before each server upload. The experiment involves 50 clients and runs for 100 communication rounds across MNIST, FMNIST, CIFAR-10, and CIFAR-100, with all other settings identical to those in Section 6.I Test accuracy convergence curves for these datasets with 50 clients shown in Figure with additional details in Appendix F
Clearly, FedWBA shows an excellent convergence rate. It improves rapidly in the first 10 rounds and then rises steadily until convergence. In contrast, some baseline models have slow convergence rates, and others like FedAvg have performance fluctuations due to data heterogeneity. The better convergence performance of FedWBA is attributed to the use of SVGD at the client side. SVGD enables our method to better adapt to local data, efficiently capture the local posterior distribution, and thereby improve the convergence efficiency in federated learning.

---

6.5 Ablation Study
We perform ablation studies to assess key components of the model under the setting of 100 clients, with all other configurations identical to those in Section 6. to deepen our understanding of the model's behavior. We consider (1) Iteration number in SVGD, (2) Kernel selection in SVGD, and (3) Kernel bandwidth in SVGD. Additionally, (4) Number of labels per client, (5) Client scheduling ratio per communication round, (6) AdaGrad parameters for SVGD learning rate and (7) Kernel bandwidth in KDE to estimate the global prior are also examined (Appendix (F4). Extended analysis in Appendix F5 quantifies the
ECC-communicttion efficiency trade-ofts,
revealing practical implementation insights.
Iteration number in SVGD: shows that accuracy improves with more SVGD iterations on MNIST and FMNIST but plateaus after about 30 and 50 iterations, respectively. This indicates that large number of local iterations are unnec- essary, avoiding excessive local computational pressure.
Table 2: Ablation studies were conducted to inves-
tigate the impact of the number of SVGD iterations, the choice of SVGD kernel, and kernel bandwidth on prediction performance. The experiments performed using 100 clients on the MNIST and FMNIST datasets.
MNIST
FMNIST
Kernel selection in SVGD: As shown in Ta-
Number SVGD
we systematically evaluate Laplacian, sig- moid, and polynomial kernels under identical hyperparameters. Laplacian and RBF both em- ployed the median heuristic for bandwidth se- lection, while the sigmoid kernel used scaling factor a = 1 and bias mial kernel adopted exponent results highlight the importance of kernel se- lection, with the RBF kernel yielding the best performance empirically.
ble 2
lerrtion Number
Acc(%)
Number
leration
Aco(%)
30
97.02
0.03
50
91.62 1: 0,01
40
97.23 ::: 0.01
60
91.65 =: 0.02
SVGD kernels
0, and the polyno-
kemnel
Acc(%)
kernel
Aco(%)
2. These
sigmoid
93.90 d: 0.07
sigmoid
87.10 d: 0.02
Laplacian
97.12 0.02
Laplacian
91.43 1: 0.02
polynomial
polynomial
91.52
RBF
97.23 d: 0.01
RBF
91.65 :: 0.02
Kernel Bandwidth SVGD
bandwichh
Acc(%)
bandwidth
Ace(%)
Kernel Bandwidth in SVGD: As shown in Ta- ble (2] the Gaussian kernel bandwidth signifi- cantly affects SVGD performance. Small band- widths yield high precision but cause unstable, fluctuating approximations and oversensitivity local noise. Large bandwidths favor global explo- ration yet obscure local details, slowing convergence. Using the median heuristic effectively balances these aspects, supporting both exploration and local refinement for stable, efficient convergence.
97.34 1: 0.05
91.71 =: 0.10
med
97.26 0,00
med
91.65 + 0.02
97.10 0.03
91.22 0.03
The findings guide configuring FedWBA's Bayesian components, showing deliberate selection of SVGD iterations, kernel type, and bandwidth critical for optimal accuracy FL.
Conclusions
In summary, we propose FedWBA, novel PBFL method that simultaneously enhances local infer- ence and global aggregation. At the client side, we employ particle-bar basel variational inference for nonparametric posterior representation. This allows clients to capture complex local posterior distri- butions by leveraging the flexibility of particle-based methods, better approximating the posterior of model parameters. At the server side, we introduce particle-base bassd Wasserstein barycenter aggregation, offering more geometrically meaningful way to aggregate local client updates. Theoretically, we establish convergence guarantees for FedWBA. Comprehensive experiments reveal that FedWBA outperforms baselines in prediction accuracy, uncertainty calibration, and convergence rate.
Acknowledgments
This work was supported by the NSFC Projects (No. 12171479, No.62576346, 12375511), the
MOE Project of Key Research Institute of Humanities and Social Sciences (22JJD11000) 10001), the fundamental research funds for the central universities, and the research funds of Renmin University of China (24XNKJ13), and Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing.
10

---

References
[1] Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes, Advances in Neural Information Processing Systems, 34:8392-8406, 2021.
[2] Martial Aguch and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-924, 2011.
[3] Maruan Al- I--hedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Representations, 2021.
[4] Shun-ichi Amari. Information geometry and its applications, volume 194. Springer, 2016.
[5] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint Stt.::I912.0055 2019.
[6] Gennaro Auricchio, Federico Bassetti, Stefano Gualandi, and Marco Veneroni. Computing wasserstein barycenters via linear programming. In International Conference on Integration of Constraint Programming, Artificial Intelligence, and Operations Research, pages 355-363. Springer, 2019.
[7] Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Federated learning with uncertainty via distilled predictive distributions. In Asian Conference on Machine Learning, pages 153-168. PMLR, 2024.
[8] Christopher M Bishop. Pattern Recognition and Machine Learning. springer, 2006.
[9] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859-877 2017.
[10] Due Bui, Kshitiz Malik, Jack Goetz, Honglei Liu, Seungwhan Moon, Anuj Kumar, and Kang G Shin. Federated user representation learning. arXiv preprint arXiv: /2535, 2019.
[11] Longbing Cao, Hui Chen, Xuhui Fan, Joao Gama, Yew-Soon Ong, and Vipin Kumar. Bayesian federated learning: A survey. arXiv preprint arXiv:2304. 13267,, 2023.
[12) Yiqiang Chen, Xin Qin, Jindong Wang, Chaohui Yu, and Wen Gao. Fedhealth: A federated transfer learning framework for wearable healthcare. IEFE Intelligent Systems, 35(4):33-93 2020.
[13] Li Daliang and Wang Junpu. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint 1r0.v:191.003381,,2 2, 2019.
[14] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in neural information processing systems, 33:3557-3568, 2020.
[15] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[16] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration modern neural networks. In International conference machine learning, pages 1321-1330. PMLR, 2017.
[17] Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning as variational inference: A scalable expectation propagation approach. In The Fleventh International Conference on Learning Representations, 2023.
[18] Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang. Jiangehuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning non-iid data. In Proceedings of the AAAI conference on artificial intelligence, pages 7865-7873, 2021.
[19] Nour Jamoussi, Giuseppe Serra, Photios A Stavrou, and Marios Kountouris. Ba-bfl: Barycentric aggregation for bayesian federated learning. arXiv preprint 1.Xi::24/2.116444 2024.
[20] Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv: 00.009//4440 /2488, 2019.
[21] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference machine learning, pages 5132-5143. PMLR, 2020.
11

---

[22] Rahif Kassab and Osvaldo Simeone. Federated generalized bayesian learning via distributed stein variational gradient descent. IEEE Transactions on Signal Processing. 2022.
[23] M. Kim and T. Hospedales. Fedhb: Hierarchical bayesian federated learning. arXiv preprint arX1v:2305.04979, 2023.
[24] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Handbook of Systemic Autoimmune Diseases, ((4), 2009.
[25] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 1:429-4500, 2020.
[26] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang. and Zhihua Zhang. On the convergence of fedavg on non-iid data. arXiv preprint arXiv:1907.02 189, 2019.
[27] Liangxi Liu, Xi Jiang, Feng Zheng, Hong Chen, Guo-Jun Qi, Heng Huang, and Ling Shao. bayesian federated learning framework with online laplace approximation. IEEE Transactions on Puttern Analysis and Machine Intelligence, 46(1):1-16, 2023.
[28] Liangxi Liu, Xi Jiang, Feng Zheng. Hong Chen, Guo-Jun Qi. Heng Huang, and Ling Shao. A bayesian federated learning framework with online laplace approximation. IEEE Transactions on Puttern Analysis and Machine Intelligence, 2023.
[29] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
[30] Guodong Long, Yue Tan, Jing Jiang. and Chengqi Zhang. Federated learning for open banking. In Federated learning: privacy and incentive, pages 240-254. Springer, 2020.
[31] Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fficien learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[32] Dinh C Nguyen, Ming Ding. Pubudu N Pathirana, Aruna Seneviratne, Jun Li, and H Vincent Poor. Federated learning for internet of things: A comprehensive survey, IEEE Communications Surveys & Tutorials, 23(3): 2021.
[33] Atahan Ozer, Kadir Burak Buldu, Abdullah Akgül, and Gozde Unal. How to combine variational bayesian networks in federated learning. arXiv preprint arXiv:2206. /0897, 2022.
[34] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(--6):3556677 2019.
[35] Nicola Rieke, Jonny Hancox, Wenqi Li, Fausto Milletari, Holger R Roth, Shadi Albarqouni, Spyridon Bakas, Mathieu N Galtier, Bennett A Landman, Klaus Maier-Hein, et al. The future of digital health with federated learning. NPJ digital medicine, 3(1):1-7, 2020.
[36] Sanvesh Srivastava, Cheng Li, and David B Dunson. Scalable bayes via barycenter in wasserstein space. Journal of Machine Learning Research, 19(8):1-35, 2018.
[37] Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in neural information processing systems, 33:21394 21 405, 2020.
[38] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[39] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2008
[40] Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623. 2020.
[41] Wing Hung Wong and Xiaotong Shen. Probability inequalities for likelihood ratios and convergence rates of sieve mles. The Annals of Statistics, pages 339-362, 1995.
[42] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:/708.00 777..07787777777 2017.
[43] Qiang Yang. Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
12

---

[44] Jianqing Zhang, Yang Liu, Yang Hua, Hao Wang, Tao Song. Zhengui Xue, Ruhui Ma. and Jian Cao. Pfilib: Personalized federated learning algorithm library. arXiv preprint 2xXx:2332..04992.
2023.
[45] Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized federated learning via variational bayesian inference. In International Conference on Machine Learning. pages 26293-26310, PMLR, 2022.
[46] Junyi Zhu, Xingehen Ma, and Matthew B Blaschko. Confidence-aware personalized federated learning via variational expectation maximization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
2023.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope. See the Abstract and Introduction sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitation of the work is discussed in the Appendix H.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

For all theoretical results, the paper provides the corresponding proofs in the appendix. Specifically, all assumptions are detailed Appendix D, while Appendices E and F present the proofs for the theorems, respectively.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper fully discloses all the information needed to reproduce the main experimental results in the paper. See the Experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide the data and code the supplemental material to reproduce the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ ctaaps..c/public/guue
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify all the training and test details necessary to understand the results. See the Experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the statistical significance of the experiments. In the tables presenting the experimental results, provide the standard deviations across multiple runs, with all results averaged over five independent trials to ensure reliability and consistency in the reported metrics.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All experiments conducted the server whose details provided in the paper. See the Experiments section.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't i make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httptp/ss/euurrps..ccpppll////// ctw.ccddeeeicdiieeeeeieenne
Answer: [Yes]
Justification:
The research conducted in the paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction)
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Guidelines:
17

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All code, models, and datasets mentioned in the text are appropriately cited with their original papers.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
18

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thood cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New assets introduced in the paper, such as code, are well documented. The documentation is provided alongside the assets the supplementary material.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMS as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurips: /att::///earrrpsssc//oooeerreee/ cc////ooeeeeeecs//22/////.. for what should or should not be described.
20

---

A Comparison of Aggregation Methods
Figure presents three different aggregation methods for three Gaussian distributions: parameter averaging, mixture, and Wasserstein barycenter. Notably, both parameter averaging and the Wasser- stein barycenter result Gaussian distributions, while the mixture does not. The choice of Gaussian distributions is motivated by the Gaussian assumption imposed by pFedBayes on the variational parameter family, which serves representative baseline in our comparative analysis of Bayesian federated learning approaches.
Comparison of Aggregation Methods
0.55
NI31E)
Dicealy
0.3
0.15
N:4A
0.19
-10
Figure 4: Comparison of Three Aggregation Methods: Wasserstein Barycenter (WB), Parameter Averaging (Avg), and Arithmetic Mean (Mixture).
B Proof of Equation (5)
Define O def diag( -) and O def diag @ Ok). Then, let M, denote the squared Euclidean distance matrix between the pairwise elements of θ and Ø which can be written as:
M @1 1 200
(8)
Following this form of Mg, the Frobenius inner product between M, and T. is given by:
(M., = -r Ir. @1 + (T 1N (Tr
(9)
1 N
tr ( (r 2T T
(10)
By minimizing Equation (9. and taking into account that O₁ is a constant, obtain:
1
T arg min (M,, Tx)
arg min - (r( 2 (0, T,Op).
(1I)
T-ET N
TEET
After obtain Te. obtain the global prior the Wasserstein barycenter of K local posteriors:
K
(0)
arg min
MM. Ti F
K
k=1
arg min
(Odiag(NV/2) TioN O diag( 11
K
(12)
Consequently, the optimal @* has an analytical solution
K
θ =
T (Ndiag(N-))
(13)
K k=
21

---

C Pseudocode for Algorithm
Algorithm FedWBA
Input: Size Z, kernels ((;;)) .) and k((,,))
) Initialize: local posterior particles ( k= global prior particles (0:)
for (communication round)
t- Server randomly samples subset of clients of size Z;
Server broadcasts {0.);=1 to each client k K;;
[Local updates]
for client k in K. do
Rebuild the continuous prior by Equation 6
Update local posterior particles by Equation ( end for
Each client k E K. upload {0k. 1 the server; [Server aggregates]
Update global prior particles by Equation
D Proof of Theorem 5.3
Proof. The ELBO (() can be expressed as the negative KL divergence: F(qx(0)) (II((x(0)) p(0) Dz)). where p(0 Dz) p(Dx)0)p(()) is the unnormalized posterior distribu-
tion.
For iteration on client k, the change in ELBO from iteration / to + 1 is given by:
F(gk (1+1)
(0)) (0))
F( (0)) = KL
(
(0)
(D)))
(T
(1+1) (0) D2))
KL
( Da)) KL 9k
(1
(0) (0) (p(0) Da))
log (Dx) log t (p(0) [D))
E %
(0)
(14)
n
E
logg(0 (
+ log (t(()) (Dx) + log del (Vot(0))
(i) is obtained from the SVGD based particle transformation ((++1) t(0l)) = /²)) + eΦ E (0))
and the definition of KL divergence. (ii) follows from the change of variable formula for densities t=(()()) = p(t(0)) det (Vot(0))
Applying a first-order Taylor expansion for A:
1 2 V log (t)) Dxx) R logí( +x)++VV log Dx) (e(
" 2
(15)
According to Assumption
5.1
we have
log det (Vot(0))
Σ log legle:1>
ir ((--"))-((- (Vot(0))=1)
(16)
1=1
Given that Vot(O) = I eV and under Assumption 5.2 applying the Neumann expansion, we get the approximation of (Vot(0))
(Vet(0)) (cV (eVo)(0))²
(17)
Then Equation
16 simplifies to:
log det (Vot(0)) 2 ir -VVoo)())
(18)
22

---

Substituting Equations
13
and (8) into Equation
we obtain:
- F( (0) 2 cB
Vo p D&) (0) [o02
E
Ir
(Voo
v log
Da (2)
2
(0). p DE))
tr (Voo
2 v log p Dz) (2)
(0)
(19)
Given that the learning rate is sufficiently small, the second term becomes negligible. Therefore, the lower bound of the ELBO increase per iteration is:
F(kk
( (((0)) (1-1) (0) F(ak 1 (0)) 2 <D( 12 (0), Dz))
(20)
E Proof of Theorem 5.9
E.1 Preliminaries
We begin by recalling the definition of Pseudo Hellinger distance and then state key lemmas that essential for the proof of Theorem 5.9 Definition E.1. (Pseudo Hellinger distance) The pseudo Hellinger distance between probability measures Pe. Pe is
2 0')
h {p(Des O),p(D{< }0)}
(21)
1/2
where represents the data index and h(p1.P2) = -
v vraww)
2 dy is the Hellinger
distance between generic densities p1.
Lemma E.2. (Generalization of [41) Theorem 1) Assume Assumption 3.6 holds. Then for any
8 > 0, there exist positive constants q1. Q2 that depend on C2 C2. such that for all subsets D3 with l..... K and all sufficiently large s,
pl D 0
p(S) 9o
2 exp((=9160) S 44ep(-92200)
sup (0.00
DRe 0n)
Lemma E.3. Assume Assumption 5.6 holds. Then for any 0, there exist positive constants T1, T2
that depend on Cx. such that for every subset DE (k I,... K) for any 2 1222
p( Disl0) I1(d0) S exp(-r,,t) O(D
p(
exp(-rrsst.
Lemma E.4. Let v denote the W2 barycenter of N measures VI...., in P2(E). Then for any Og E E, the following inequality holds:
W2(v, do.) S N W2(uy, 50.)
This lemma establishes relationship between the barycenter and the distances from each individual measure to a fixed point 0g. The proofs Lemmas 2 to E.4 can be found in [36].
Lemma E.5. Suppose that Assumption 5.4 Assumption 5.8 hold for the k-th subset posterior p(0 Dx)
with = 1,... K. Then there exists a constant Cg that depends on Cr, Cr. O2, K, Cg and does not depend on k such that as + 0o,
log²
Ep,, w3 W3 (L) 60, < SCa
23

---

Proof. Let = (ogg
2u Due to the compactness of in assumption Assumption 5.4 there exists a large finite constant Mo such that p(O. 0n) I My. We begin with decomposition of the W2
distance from the k-th subset posterior p(0 to the delta measure at the true parameter 0g:
D2)
EP W (p(0) (x). (.) -BPo p (0, 0u)p(d0 D2)
Eroo S 110:(0..0)$C....
(0, (0.00))(d010) Dx)) Eroo [ 110000.0.>0... (0,,0)) D&))
s (Cana)² + My Pe (0 0o) > Can.lDa)
(22)
will choose the constant C4 as C4 (2r Iez where Ct. Qt, ri are the constants in Assump-
tion 5.4] Assumption Lemma F.2 and Lemma Using Assumption 5.3 we can further replace
the e metric by the pseudo Hellinger distance:
p(0 E :: :00,00)) C.na S p(0 E: (a Poo.k) > )CL(CuM)) Db))
II
p(D (0)
Il(d0) (23)
P(Dr. 0a)
10 (0.00)> nP
pD 0) p(De.
I[de)
For the denominator Equation 22 by Assumption and Lemma
when is sufficiently
large, with probability at least exp( - 20)
p Dk 0)
ll(d0) > exp(-ri
(24)
Das 0o)
and Lemma E.2 setting /
/2r
For the numerator in Equation (23 by Assumption 5.6
11. we
get that with probability at least 4exp
2r192 Sm 41
pl Dks ()) p(Dxr 0.))
sup
$ exp((22rsSpp
(25)
(0.0:h.
Therefore, based on Equations (23) to 25 with probability at least 1 4 4exp exp(-r2sm2")
22192 Sm1
p(0 E E: p(0,0.) Canss D2) exp(-2r, +r1Sn120 S eexp(-r_S1)"
(26)
Let An, be the event {0 f0 = p(0 = p(0, 0n) Canss D&) S exp(-ri exp((ri Then can bound the second term in Equation (22 as follows:
E Ero, 0n) > Cans Die)
s Eroo
T( )
(1 (p > Cans D1s)
[(AC )p p(0,00)) 0o) > CAn Dr
+ Po
< exp(--r,s120 20 4exp
2r1 Smp
(27)
91
+ +epp(
< 6 exp((CC5s1)) 29)
for Cs2 min(r1, 72,
2rig2 ), clearly the second term is dominating the other two given m S
91

---

Therefore, for Equation 22
since = (s/logg s)
zur + co, explicit bound will be
log s
EPgg W (00D))),0(())
- C7
+ GMG cxp( -- log 8)
log 1
(28)
log
as becomes sufficiently large, where the constant C' depends C4, Cs. which further depends
on Qı 42, rı, r2, Ct. Since Qı Q2 in Lemma E,2 and rı, Y2 Lemma E.3 depend on C₁, C2 K, Cr it follows that C3 depends CL. C1. Cz. K. Cr.
E.2
Main Proof
Now give the full proof of Theorem 5.9
Proof. For notational simplicity in the subsequent derivations, we abbreviate the Wasserstein barycen- ter distribution ((0) D&) p(0). Following Lemma E.5 we have:
p(
Wa( (0, (1) >
log
c
(29)
1 K
2 log a
<
(S) do
Wa(p(01D2), (-))
K
Next, applying Markov's 's inequality and using the relation between 1 and l2 norms, we obtain:
>
log
2
p(s)
(Walo(0)Da),oo.()) 6o.)) >
EPa
Wa(p(0)D)),
log
2/a
K
8//0
(30)
81/a
Er (p) (-))
CK log"
2/a
kml
(31)
8
log 8 8
KC
(32)
CK log
Ci
(33)
C
where the second step inequality follows from the fact that for nonnegative real-valued random vari-
ables X1. Xn. ( Σi= X,) & Σ X?, and the third step inequality uses the previously
obtained upper bound En,, W3 (p(0 5n, ()) S C1 log
Finally, combing these results, we conclude:
log² 2
W₂ (0o)) Or =0r(1
s
(34)
L
F More Experimental Results
In this section, we provide more experimental results to demonstrate the performance of our method compared to others.
25

---

F11 Performance of Global Model
To validate the performance of the global model. we compare with FedAvg, FedProx, and SCAFFOLD on MNIST, using settings identical to Section 6.1] The results show highly comparable performance, demonstrating the validity of the global model as a prior regularizer.
Table 3: Global model performance with best results bolded.
DATASET
METHOD
50 CLIENTS
100 CLIENTS
200 CLIENTS
FEDAVG
91.98 + 0.07
91.76 + 0.08
90.94 + 0.06
FEDPROX
92.12 + 0.08
92.04 + 0.11
90.82 + 0.16
MNIST
SCAFFOLD
92.90 + 0.07
92.14 0.08
90.85 0.11
OURS
92.02 + 0.02
91.95 + 0.03
89.90 + 0.03
F.2
Performance of Uncertainty Quantification
In this section, we present the performance of the remaining six algorithms in uncertainty quantifica- tion, as shown in Figure
PFcdVEM
DFcdME
FedAvg
1.00
1.00
1.00
0.2517
ECE 0.2114
ECE 0.1852
0.75
0.75
0.75
Accuracy
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.0
0.5
1.0
1.0
0.0
0.5
Confidence
Confidence
Confidence
FedProx
FedPer
SCAFFOLD
1.00
1.00
1.00
0.1806
ECE 0.1595
BCB 0.1070
0.75
0.75
0.75
Accurscy
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
0.0
0.0
Confidence
Confidence
Confidence
Figure 5: Reliability diagrams of the six methods on CIFAR- 100. The perfect calibration is plotted as a red diagonal line, and the actual results are presented as bar charts. The gap between the top of each bar and the red line represents the calibration error. The ECE calculated and placed in the top-left corner of the figure. Among them, the method with the highest ECE value has the worst calibration performance.
Our uncertainty quantification experiments across MNIST, FMNIST, and CIFAR- 10 demonstrate that our method achieves state-of-the-as performance in ECE, outperforming existing baselines on all benchmarks, as presented in Table 4 This consistent superiority highlights enhanced generalization and calibration capabilities under varying data distributions and task complexities.
Convergence Rate
This section presents a comparison of algorithm convergence across MNIST, FMNIST, CIFAR- 10, and CIFAR-100 with 100 and 200 clients (as depicted in Figure for 100 clients and Figure 7 for 200 clients). Similar to observations prior client configurations, FedWBA exhibits rapid test accuracy
26

---

Table 4: ECE of different methods on MNIST, FMNIST and CIFAR-10 with 50 clients. Optimal
results are bolded.
Dataset
method
MNIST
FMNIST
CIFAR-10 10
FedAvg
0.0449
0,0211
0,0394
FedProx
0.0147
0,0308
0.0264
SCAFFOLD
0.0424
0.0343
0.1002
FedPer
0.0061
0,0082
0.0284
perFedAvg
0.0213
0,0258
0,0663
PFedME
0.0097
0.0424
0.0549
PFedBayes
0.01558
0.0132
0.0420
PFedVEM
0.0130
0.0107
0,0476
pFedGP 0.0217
0.0193
0.0245
Ours 0.0014
0.0078
0.0132
growth within the first communication rounds across all datasets, followed by gradual refinement and eventual convergence. This consistent pattern across diverse datasets and client scales highlights the approach's is effectiveness and stability in federated learning, demonstrating robust adaptability to varying data distribution complexities.
Aeseago Auruunry
PNNIST
CHARIE
CIARRIO
1910
Figure 6: Comparison convergence rates of different methods on MNIST, FMNIST, CIFAR- 10, and CIFAR- 1100R-1001A with 100 clients.
Avrrage Accursy 67288 Clee de NINIST
ENINST
CHARIS
CHARII
19
Communicatoon
Figure 7: Comparison of convergence rates of different methods on MNIST, FMNIST, CIFAR-10, 10, and CIFAR-100 with 200 clients.
F.4 Ablation Study
In this section, we conduct ablation studies on four key components: (1) Number of labels per client, (2) Client scheduling ratio per communication round, (3) two AdaGrad parameters for SVGD learning rate, and (4) Kernel bandwidth in KDE to estimate the global prior. In the optimization algorithm, the AdaGrad update rule is:
@(+1) (
17
Gg' +>
VG
where G¹ is the accumulated sum of squared gradients up to the 1-th iteration, g' the gradient vector the 1-th iteration, 7 is the global learning rate determining step-size, and À is smoothing
27

---

term that prevents the denominator from being zero and enhances algorithm stability, which is crucial for SVGD learning rate adjustment.
To validate the method's generality across varying client problem complexities, evaluate MNIST and FMNIST with 2, and 10 labels per client-specifically, construct client heterogeneity in terms of label distribution through this varying number of labels per client, simulating real-world scenarios where clients only hold partial and unequal label information. Smaller label counts indicate simpler client tasks. As shown in Table S accuracy decreases with increasing labels, aligning with expected difficulty trends. Despite performance variation with task complexity, FedWBA maintains stable convergence and high accuracy at each label setting, demonstrating robust adaptability diverse client-side problem scales.
Number of labels per client:
Client scheduling ratio per communication round: We conduct experiments on MNIST and FMNIST with client scheduling ratios of 0.1, 0.2, and 0.5 per communication round. Results show that a higher scheduling ratio correlates with faster convergence of the global prior and higher client- side accuracy-this is attributed to more client updates contributing to the global model optimization at each round. However, scheduling ratio of 0.5 is impractical in real-world scenarios, as it incurs excessive communication bandwidth consumption and computational burdens on the server. Thus, adopt a scheduling ratio of 0.2 for experiments in Section 6. to balance performance and practicality,
Impact of and A: A small global learning rate 7) leads to extremely slow convergence, requiring significantly more iterations to reach or approach the optimal solution. Conversely, a large 1j may cause the model skip the optimal region in the parameter space due overly large step sizes. A small À can result in training fluctuations, while large X weakens gradient accumulation, potentially impeding convergence.
Bandwidth of kernel KDE: The influence the bandwidth KDE is relatively minor. In the context of our experiments, across different datasets and model configurations, varying the KDE bandwidth within reasonable range not lead to substantial changes model's performance metrics.
Table 5: Ablation studies the impact of the number of labels per client, client scheduling ratio per communication round, 7), À and KDE bandwidths on prediction performance, conducted with 100 clients MNIST and FMNIST.
MNIST
FMNIST
NUMBER
LABELS
ACC(%)
LABELS
Acc(%)
97.23 1 0.01
91.65 1 0.02
95.37
CLIENT SCHEDULING RATIO PER COMMUNICATION ROUND
SCHEDULING RATIO
Acc(%)
SCHEDULING RATIO
Acc(%)
96.78
90.77 0.04
0.2
97.57 + 0.02
0.2
91 91.65 + 0.02
0.5
97.71 1: 0.03
D.5
91.73 1 0.01
GLOBAL LEARNING RATE OF ADAGRAD IN SVGD (i)
Acc(%)
Acc(%)
0.01
97.23 + 0.01
0.002
91.55 + 0.02
0.03
95.59 1 0.08
0.004
91.65 +: 0.02
FOR SVGD
Acc(%)
Acc(%)
10-7
97.22 + 0.01
10-7
91.64 = 0.01
0.02
10-9
97.24 1 0.01
10-9
91.66 0.01
97.23
10-10
91.65
BANDWIDTH KDE
BANDWIDTH
Acc(%)
BANDWIDTH
Acc(%)
0.30
97.25 0.02
0.30
91.67
0.55
97.26 4 0.02
0.55
91.67 + 0.01
0.70
97.24 0.01
0.70
91.66 0.01
28

---

F,5 Communication Cost
First, we examine the relationship between the number of particles and communication cost. As stated in Section 4.2 with the increase in the number of particles N, the empirical distribution of particles asymptotically converges to the optimal variational distribution in the corresponding RKHS, In Table [ we report the relationships among the number of particles, the size of uploaded data, and the prediction accuracy. Here, the unit of data size megabyte (MB). Increasing the number of particles leads to slight improvement in accuracy but imposes communication burden. Therefore, 10 particles can strike a balance between communication volume and performance.
Table Ablation study on the impact of particle count on communication overhead and prediction accuracy, conducted with 100 clients MNIST and CIFAR-10.
NUMBER OF PARTICLES
DATASET
COMM.(M)
ACC(S))
1.52
96.63 0.02
3,03
96.71
MNIST
6.07
06.82
15.17 MB
96.86 0.01 0.01
2.31 MB
70.71 0.10
10
4.62 MB
72.00 0.14
CIFAR-10
20
9.25 MB
72.14
50
23.11 MB
72.15 0.01
While using SVGD approximate the posterior introduces moderate communication overhead due to particle updates, we also consider replacing Bayesian layers with standard frequentist layers to reduce costs. With 5 labels per client, we validate this trade-off using a 5-layer CNN for CIFAR-10 prediction, analyzing how the number of Bayesian layers affects communication cost, test accuracy, and ECE. When the number of Bayesian layers reduced to the method effectively reduces to FedAvg common non-Bayesian federated learning baseline-providinn direct comparison point.
Table 7: Impact of number of Bayesian layers communication cost, test Accuracy, and ECE.
THE NUMBER OF BAYESIAN LAYERS
COMM.(M)
Acc(%))
ECE(L)
5
23.99 MB
76.85 + 0.02
0.0376
16.08 MB
76.12 + 0.02
0.0582
2.40 MB
68.84
0.15
0.0853
2.38 MB
62.53 0.44
0. 0.1256
Model with Beyesian Layers
Model with Bayesian Layers
Model with Bayesian: Layers
Model Bayesian Layers
1.0
0.1256
BCE 0.0855
0.0582
BCB 0.0376
0.-
0.5
0.5
0.4
0.4
0.4
0.6
0.2
0.2
0.0
0.
1.00
0.00
0.25
0.50
0.75
0.00
0.50
1,00
0.00
0.25
0.50
Confidence
Confidtecc
Confidaace
Confidmace
Increasing Bayesian Layers
Figure Reliability diagrams for different numbers of Bayesian layers.
As shown in Table 7 and Figure 8 decreasing the number of Bayesian layers lowers communication volume but leads to a noticeable drop in accuracy and an increase in ECE. This indicates that
although Bayesian layers impose higher communication costs, they are critical for maintaining precise uncertainty quantification and prediction reliability. Specifically, our method's low ECE-even at higher communication costs costs-hightt highligtss its suitability for high-risk scenarios where accurate
29

---

uncertainty assessment is paramount, justifying the additional overhead. This trade-off underscores the value of our approach in applications requiring both performance and rigorous uncertainty quantification.
F.6 Additional Experiments
To further verify the scalability of our method on complex datasets and large-scale networks, also conduct experiments on Tiny-ImageNN (a reduced version of ImageNet) using the ResNet18 8 model. Specifically, the experimental setup includes 50 clients (each with 20 distinct classes to simulate non-iid data distribution) and runs for 100 communication rounds. These supplementary experiments further demonstrate the superiority of our method in such complex data scenarios and advanced network setups.
Table 8: Test accuracy (% SEM) over 50 clients Tiny-ImageNet. Best results bolded.
DATASET
METHOP
Acc(%?)
ECE(J)
FEDAVG
27.14 0.12
0.3170
FEDPER
35.43 0.19
0.2890
FEDPROX
24.97 0.25
0.2515
TINY-IMAGENET
SCAFFOLD
31.37
0.0698
PFEDME
36.24
0.0631
PERAVG
1.26
0.1774
OURS
37.91 + 0.07
0.0421
G Limitations and Future Work
FedWBA's key limitation stems from the computational complexity SVGD for nonparametric posterior approximation at clients. While fixed particle count balances efficiency, complexity rises significantly with higher-dimensional models. To address this, future work could explore lightweight Bayesian inference techniques, such as sparse particle approximations, reduce computational overhead while retaining nonparametric flexibility.
30


Reviews:
[{"review_id": "vjaagwW8Q7", "reviewer": "Reviewer_Jqa5", "summary": "This paper proposes a bayesian personalized federated learning approach, that goes beyond previous such approaches by enabling non-parametric posterior distributions.\nThis is made possible through 1) non-parametric local particle based variationak inference via Stein GD, and then 2) aggregation through Wasserstein barycenters.\nThe method is first analyzed theoretically through non asymptotic convergence guarantees, and then in real world experiments.", "questions": "None", "limitations": "None", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe paper is overall very well written and the ideas are clearly exposed. I didn’t read the proof and I am not familiar at all with Bayesian inference, but the results totally make sense. The authors made an effort into making the paper reader friendly.\n\nEnabling non-patzmetrkc modeling is key in modern ML: this seemed to be lacking in the literature here. Parametric assumptions are indeed way too restrictive and can lead to mispecification biases.\n\nThe method seems to be original and new. As I am not an expert at all in VI, I cannot judge the use of variationak Stein GD, although it seems new in this context.\nHowever, I think that exploring new ways to aggregate local weights models or probabilities is generally a promising venue. The use of Wasserstein barycenters here seem to be a natural and well motivated choice.\n\nTheory and experiments seem to validate the method.\n\n**Weaknesses**\n\nI am not an expert (far from it) in Bayesian or variational inference, so it is hard for me to formulate technical weaknesses.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful and encouraging feedback. We especially appreciate your recognition of the clarity of our presentation and the originality of our method. We’re also grateful for your positive remarks on the motivation for non-parametric modeling and the use of Wasserstein barycenters, which are central to our approach."}, {"review_id": "G9H8nED5Ph", "reviewer": "Reviewer_TsB1", "summary": "- The paper aims to tackle the personalized FL using a Bayesian method.\n- Existing methods have an issue of average aggregation in parameter space, which may not be a Euclidean space, but a manifold with information geometry. Motivated from this, they proposed a\nnon-parametric posterior inference and aggregation.\n- Specifically, the client local update is done by particle-based VI using SVGD, while the server aggregation is done by particle-based Wasserstein barycenter aggregation.\n- So, it is claimed that the proposed method has better geometric interpretation than parametric averaging thanks to non-parametric sampling.", "questions": "See \"Strengths And Weaknesses\" section.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "[Strength]\n- The main contribution seems to be the introduction of a non-parametric FL method with particle-based inference and aggregation.\n- Convergence of the method is theoretically provided.\n\n[Questions/Weakness]\n- There are several missing refs which are closely related to the paper. See below.\n- L38-40: \"Problematic averaging posterior params due to info geometry\" -- Any evidence in practice?\n- L34: \"True posterior not Gaussian\" -- There are several methods that deal with non-Gaussians, eg, mixture of gaussians (eg, Ref 3 and 4 below).\n- The experiments do not seem extensible enough to show empirical verification. It is preferred to see results on other FL settings (eg, as a standard FL practice, the problem hparams may need to be varied such as: fraction of participating clients in each round, number of epochs for client local posterior updates, degree of data heterogeneity or non-iid-ness).\n- The experiments are only done with relatively small networks (eg, mlp, lenet). This might be due to the inherent computational issue incurred by non-parametric methods; How scalable is the approach?\n- There are parametric Bayesian methods using non-Gaussian posteriors, eg, Gaussian mixtures. How would the proposed method perform compared to these works? -- I am not asking for an extra set of experiments, but curious about it.\n- Can you provide the convergence analysis results as a function of the number of iterations? Is there any generalization error guarantee? which seems to be done in some previous works.\n- A main drawback of the proposed work is that the particle-based methods do not seem to scale well. What are the practical impacts of the work beyond theoretical interests?\n- As said, WB is tractable for discrete distributions, which is why they used a mixture of dirac measures. However, this might necessitate a large number of samples for high-dim model parameters. So, it looks computationally not scalable.\n- ECE scores in Fig 2 of most methods are sufficiently good (<1%) and statistically less distinguishable?\n\nRefs:\n1, Bayesian Inference Federated Learning for Heart Rate Prediction. Fang et al 2020\n2. Partitioned Variational Inference: A framework for probabilistic federated learning. Ashman et al 2022\n3. FedHB: Hierarchical Bayesian Federated Learning. Kim et al 2023\n4. A Statistical Framework for Personalized Federated Learning and Estimation: Theory, Algorithms, and Privacy. Ozkara et al, ICLR 2023\n5. Fderated Learning as Variational Inference: A Scalable Expectation Propagation Approach. Guo et al 2023", "was_revised": true, "final_justification": "Thanks for the rebuttal. I have upped the score.", "rebuttal": "Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.\n\n> **Q: Several closely related references are missing.**\n\nA: Thank you for pointing out the missing related references. We will make sure to include and discuss these works in the final camera-ready version. \n\n> **Q: Problematic averaging posterior params due to info geometry.**\n\nA: To empirically validate the superiority of our proposed Wasserstein barycenter (WB) aggregation, we **conducted a direct comparison against simple parameter averaging** (i.e., averaging the mean and variance parameters of Gaussian posteriors across clients) on the MNIST dataset. The results, presented in the table below, demonstrate that **WB consistently outperforms parameter averaging across all client scales** (50, 100, and 200 clients) on key metrics such as prediction accuracy and ECE. This confirms that WB, by leveraging optimal transport geometry to aggregate distributions rather than parameters, can effectively improve performance.\n\n|          | Acc ($\\uparrow$) |          |          | ECE ($\\downarrow$) |          |          |\n|----------|------------------|----------|----------|--------------------|----------|----------|\n|          | 50 clients       | 100 clients | 200 clients | 50 clients         | 100 clients | 200 clients |\n| Average  | 97.15 $\\pm$ 0.01 | 97.01 $\\pm$ 0.02 | 96.21 $\\pm$ 0.02 | 0.0021             | 0.0030    | 0.0080    |\n| WB       | **97.99 $\\pm$ 0.04** | **97.36 $\\pm$ 0.01** | **96.95 $\\pm$ 0.01** | **0.0014**         | **0.0021** | **0.0064** |\n\n> **Q: Several methods deal with non-Gaussians, e.g., mixture of gaussians.** \n\nA: We agree that methods like mixture of Gaussians (e.g., FedHB) offer more flexibility than single Gaussian approximations for non-Gaussian posteriors. However, as emphasized by Bishop (2006, Pattern Recognition and Machine Learning, Section 10.1), the accuracy of variational inference is fundamentally limited by the expressiveness of the chosen variational family. When the true posterior lies outside this family—as in non-conjugate models with multi-modal or heavy-tailed distributions—parametric approximations can introduce systematic errors. While mixtures of Gaussians increase flexibility over single Gaussians, they still suffer from **parametric constraints**, such as the need to predefine the number of components based on heuristics. An insufficient number of components may fail to capture the complexity of posterior distributions, which is common in federated settings with highly heterogeneous data. In contrast, our **particle-based nonparametric approach avoids these limitations**, enabling **more flexible posterior approximation** that adapts to the true distributional characteristics of each client.  \n\n> **Q: The experiments lack extensibility for thorough empirical verification; results on other FL settings with varied problem hyperparameters are preferred.**\n\nA: To address the need for extensibility, we have **supplemented our experiments with additional results under varied FL settings and hyperparameters**, as shown in the table below. \n\n**Fraction of participating clients in each round:** We conducted new experiments varying this fraction from 0.1 to 0.5, with clients performing 20 local iterations per round, demonstrating consistent performance gains of our method across all participation rates.\n\n| Dataset | Scheduling ratio | 50 clients       | 100 clients      | 200 clients      |\n|---------|------------------|------------------|------------------|------------------|\n| MNIST   | 0.1              | 96.83 $\\pm$ 0.02 | 96.78 $\\pm$ 0.02 | 96.22 $\\pm$ 0.03 |\n| MNIST   | 0.2              | 97.58 $\\pm$ 0.03 | 97.57 $\\pm$ 0.02 | 97.03 $\\pm$ 0.02 |\n| MNIST   | 0.5              | **97.76 $\\pm$ 0.02** | **97.71 $\\pm$ 0.03** | **97.38 $\\pm$ 0.01** |\n\n**Number of epochs for client local posterior updates:** Since we use SVGD for posterior approximation, this hyperparameter corresponds directly to the Iteration Number in SVGD reported in **Table 2**. Our ablation study there shows **convergence within 10-50 iterations across datasets**.  \n\n**Degree of data heterogeneity:** We systematically controlled heterogeneity by adjusting the number of labels per client, where fewer labels induce higher skew. Results **in Table 5 of the appendix** show our method maintains superiority over baselines even under extreme non-iid conditions (e.g., 2 labels per client).  \n\nThe results consistently validate our method's superiority across diverse configurations, reinforcing its generalizability and practical utility.\n\n> **Q: The experiments are only done with relatively small networks.**\n\nA: Thank you for your suggestion. Due to time constraints, we conduct new experiments on **Tiny-ImageNet**, a reduced version of ImageNet, and supplemented with the updated **ResNet18** model. Specifically, the experimental setup includes 50 clients, each containing 20 distinct classes to simulate non-iid data distribution, and runs for 100 communication rounds. Due to time constraints, we were unable to complete all experiments and compared with the baselines in the table below. These additional experiments further verify the scalability of our method on more complex datasets and more sophisticated network architectures. We believe these results can effectively demonstrate the applicability of our approach in scenarios involving complex data and advanced models.\n\n| Method   | Acc               | ECE     |\n|----------|-------------------|---------|\n| FedAvg   | 27.14 $\\pm$ 0.12  | 0.3170  |\n| FedPer   | 35.43 $\\pm$ 0.19  | 0.2890  |\n| FedProx  | 24.97 $\\pm$ 0.25  | 0.2515  |\n| SCAFFOLD | 31.37 $\\pm$ 0.22  | 0.0698  |\n| pFedME   | 36.24 $\\pm$ 0.14  | 0.0631  |\n| PerAvg   | 34.86 $\\pm$ 0.20  | 0.1774  |\n| Ours     | **37.91 $\\pm$ 0.07** | **0.0421** |\n\n> **Q: Need convergence analysis results as a function of iteration count and generalization error guarantees.**\n\nA: Regarding convergence analysis, Theorem 5.1 in our theoretical analysis rigorously establishes that, with increasing iteration count, the variational distribution converges to the true posterior. Meanwhile, Theorem 5.2 guarantees that as the client data size grows, the Wasserstein barycenter converges to the true parameter. Empirically, Figure 3 validates this by showing the accuracy convergence curves across datasets, confirming stable performance within 50-100 communication rounds.\n\nFor generalization error guarantees, deriving such bounds under the SVGD framework remains a challenging open problem due to its nonparametric nature and reliance on particle-based approximations. While recent works (e.g., [26]) have made progress in specific settings, extending these results to our federated and personalized Bayesian setting is beyond the scope of this paper. We acknowledge this as an important direction for future research and plan to investigate generalization bounds for particle-based FL in subsequent work.\n\n[26] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.\n\n> **Q: The scalability limitations of particle-based methods and the practical implications beyond theoretical contributions require clarification.**\n\nA: While particle-based methods inherently involve higher communication costs compared to traditional parameter-transmission approaches, our experiments (Appendix G.5) confirm that this overhead remains manageable in practice.  \n\nSpecifically, we demonstrate that a **small number of particles (e.g., 10) suffices to achieve strong performance across datasets**. As detailed in Appendix Table 6, the communication cost per client per round is quantitatively modest: for a 100-neuron DNN, the data uploaded per client ranges from 1.52MB (with 5 particles) to 3.05MB (with 10 particles). This magnitude is feasible even for resource-constrained clients (e.g., edge devices in healthcare or IoT), especially when combined with our strategy of reducing Bayesian layers to further lower overhead—all while retaining acceptable accuracy and calibration. \n\n> **Q: WB may require a large number of samples for high-dimensional model parameters, making it computationally unscalable.**\n\nA: We appreciate your feedback, but there appears to be a misunderstanding. In practice, we **do not require a large number of particles for computation**. As shown in Table 6 of the appendix, our experimental results demonstrate that 10 particles are sufficient to achieve robust performance across datasets.  \n\nFurthermore, we have also **proposed a method to reduce the effective number of particles by decreasing the number of Bayesian layers** (detailed in Appendix G.5). This strategy effectively lowers communication and computational overhead while maintaining acceptable prediction accuracy and uncertainty calibration. These findings collectively confirm that WBA remains feasible for large-scale applications.\n\n> **Q: ECE scores in Fig 2 for most methods are sufficiently good with statistically indistinct differences.**\n\nA: You are correct that the ECE scores in Figure 2 for MNIST show minimal differences among top-performing methods, which is attributable to the relative simplicity of MNIST and our focus on presenting the four best-performing methods in the main text. To address this, we have provided **comprehensive ECE results for all datasets and methods in Appendix Table 4**, which reveal more pronounced disparities on more challenging benchmarks like CIFAR-10 and CIFAR-100. In the camera-ready version, we will replace Figure 2 with results from a dataset that better discriminates calibration performance, ensuring clearer differentiation among methods."}, {"review_id": "eP8K9TSC6H", "reviewer": "Reviewer_nESY", "summary": "This paper proposes a novel framework named Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation, which addresses two key limitations in existing PBFL methods: (1) restrictive parametric assumptions in client posterior inference, and (2) naive parameter averaging for global aggregation. FedWBA uses particle-based variational inference for nonparametric client posterior representation and introduces Wasserstein barycenter aggregation to geometrically align local updates on a distribution manifold. Local convergence is guaranteed via a proven KL-divergence decrease bound, while global convergence shows the Barycenter approaches the true parameter as client data scales. Experiments demonstrate state-of-the-art performance in accuracy, uncertainty calibration, and convergence rate.", "questions": "1. How to choose a suitable number of particles?\n2. The paper uses a two-step method to update $\\hat{\\overline{p}}(\\theta)$, is it optimal? I am wondering whether there is a theoretical guarantee for the optimality of $\\hat{\\overline{p}}(\\theta)$.\n3. It would be better to provide the communication comparison with other advanced algorithms.\n4. In Wasserstein barycenter aggregation, how does the computational complexity of the optimal transport plan ($T_k^\\star$) scale with respect to both the number of particles (N) and the number of clients (K)?", "limitations": "The authors states that the main limitations are the computational complexity of SVGD for nonparametric posterior approximation at clients, which was left for future work.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\n1.  FedWBA uniquely combines nonparametric posterior inference with geometric aggregation, addressing two critical flaws in prior PBFL: (i) parametric bias from restrictive variational distributions and (ii) geometric mismatch in naive parameter averaging. This dual improvement enhances both local adaptability and global consistency.\n2. The authors establish rigorous convergence guarantees through dual theoretical frameworks. For ​local convergence, they derive a lower bound on the Kullback-Leibler divergence reduction per iteration of Stein Variational Gradient Descent. For ​global convergence, they demonstrate that the Wasserstein barycenter aggregation asymptotically converges to the true parameter.\n3. Experiments span four datasets with non-IID label skew, 10 baselines, and diverse metrics. FedWBA consistently outperforms competitors, including pFedBayes and pFedGP, which shows robustness across data scales and task complexities.\n\nWeaknesses\n\nThe communication cost is usually higher than the computational cost, but FEdWBA has to send $N$ particles to address the limitations of the paper. The experimental configuration employs N=10 particles, resulting in communication costs that are 5-10 times higher than conventional parameter-averaging approaches (e.g., FedAvg) or Gaussian-variational methods (e.g., pFedBayes). However, this significant communication burden yields diminishing returns, as the accuracy improvements remain marginal compared to these baseline methods.\n\nFurthermore, the paper acknowledges the scalability challenge posed by increasing particle counts, while demonstrating that communication overhead grows linearly with N, it fails to propose mitigation strategies for this critical limitation.", "was_revised": true, "final_justification": "The paper is interesting to introduce particle-based variational inference for nonparametric client posterior representation and Wasserstein barycenter aggregation to geometrically align local updates on a distribution manifold. But the existence of many particles increases the burden of communication. So I keep my rating unchanged.", "rebuttal": "Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.\n\n> **Q: No mitigation strategies proposed for scalability issues from linear communication overhead with increasing particles.**\n\nA: We sincerely appreciate your feedback. Due to space constraints, we initially included the **discussion on reducing Bayesian layers to mitigate communication overhead in the Appendix G.5**. In the camera-ready version, we will elevate this content to the main text for better accessibility. As shown in Appendix Table 7 and Figure 8, reducing Bayesian layers does reduce communication costs, but at the cost of increased ECE. While Bayesian layers introduce higher communication costs, they are indispensable for maintaining precise uncertainty quantification, which are critical for high-stakes applications. This trade-off justifies the additional overhead in scenarios where robust uncertainty estimation is paramount.\n\n> **Q: How to choose a suitable number of particles.**\n\nA: The choice of a suitable number of particles is determined by **balancing performance and efficiency**, as validated by our experimental analysis (see **Appendix G.5**). Specifically, we conducted ablation studies on the impact of particle count $N$ across datasets (MNIST, CIFAR10) by evaluating relationships between $N$, communication overhead, and prediction accuracy (Table 6). Practically, we found $N=10$ to be a robust choice: it achieves near-optimal accuracy while keeping communication overhead manageable, thus striking a balance between performance and efficiency. For scenarios with stricter communication constraints, a smaller $N$ can be used at the cost of slight accuracy degradation.\n\n> **Q: The paper's two-step update of $\\hat{\\bar{p}}(\\theta)$ lacks clarity on optimality and associated theoretical guarantees.**\n\nA:  The global prior $\\hat{\\bar{p}}(\\theta)$ is explicitly defined as the Wasserstein barycenter of local posterior distributions, which, by definition, is the optimal solution to minimizing the weighted sum of Wasserstein distances to all client posteriors (as formalized in Section 4.3). This property is well-established in optimal transportation theory: the Wasserstein barycenter is rigorously proven to be the geometrically meaningful \"mean\" of distributions in the Wasserstein space, ensuring it captures the central tendency of the aggregated posteriors while preserving their manifold structure [31].  \n\nFurthermore, our theoretical analysis (Theorem 2) builds on this foundation by showing that as client data sizes grow, this barycenter converges to the true parameter distribution, reinforcing its optimality in federated settings where local data is heterogeneous. Thus, the construction of $\\hat{\\bar{p}}(\\theta)$ is both theoretically grounded in classic optimal transportation results [31] and extended to our federated learning framework with additional convergence guarantees.\n\n[31] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 11(5-6):355–607, 2019.\n\n> **Q: Need communication comparisons with other advanced algorithms.**\n\nA: We **performed new experiments** in the following table, which compares communication costs across methods under the same MNIST setup as Section 6.1, our Bayesian approach incurs higher overhead than frequentist baselines. However, this enables superior uncertainty calibration, critical for safety-critical applications. \n\n| Dataset | Method    | Communication Cost |\n|---------|-----------|--------------------|\n| MNIST   | FedAvg    | **0.30M**          |\n| MNIST   | SCAFFOLD  | 0.61M              |\n| MNIST   | pFedME    | **0.30M**          |\n| MNIST   | pFedBayes | 0.61M              |\n| MNIST   | Ours      | 0.76M              |\n\n\n> **Q: The computational complexity of the optimal transport plan $T_k^*$ in the Wasserstein barycenter aggregation w.r.t. particle count $N$ and client number $K$.** \n\nA: The computational complexity of computing $ T_k^*$  corresponds to that of solving the linear programming problem in Equation 3.\n\nWith respect to the number of clients $ K $, since $ T_k^* $ must be computed separately for each client, the total complexity scales linearly with the number of clients, i.e., $ O(K) $.\n\nWith respect to the particle count $ N $, the complexity depends on the specific algorithm used to solve the linear program. For example, common solvers such as the simplex method or interior point methods have different computational complexities with respect to $ N $. Therefore, the exact scaling will vary depending on the chosen solver."}, {"review_id": "LhqcSSfHVE", "reviewer": "Reviewer_havj", "summary": "This work proposes FedWBA for federated learning. It improves both local inference and global model aggregation. At the client level, it employs particle-based variational inference for flexible posterior estimation, while at the server level, it uses Wasserstein barycenter aggregation for more meaningful global updates. FedWBA has a strong theoretical foundation for both local and global convergence guarantees.", "questions": "1. Could you include more regular benchmark datasets in the experiment, such as ImageNet and COCO?\n\n2. Could you consider standard model architectures for experiments, such as ResNet and ViT?\n\n3. It lacks the discussion of how assumptions D.1 - D.7 may hold in practice. Please add some discussions on this issue.\n\n4. Could you release the experimental code of this work anonymously?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. This work has a strong theoretical foundation.\n\n2. Sufficient baseline methods have been considered for comparison.\n\n**Weakness:**\n\n1. The experiments could include more regular benchmark datasets, such as ImageNet and COCO.\n\n2. The used models, MLP and LeNet, are too old, and cannot demonstrate the effectiveness of the optimizer in the current era. Please consider standard model architectures such as ResNet and ViT.\n\n3. The assumptions cannot be put in the Appendix, which appears to be intentional. It also reduces the overall readability of this work.\n\n4. It lacks the discussion of how assumptions D.1 - D.7 may hold in practice. Please add some discussion on this issue.\n\n5. Without releasing the experimental code, the reproducibility of this work cannot be fully ensured.", "was_revised": true, "final_justification": "I appreciate the reviewer's response and raise my score to borderline accept.", "rebuttal": "Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.\n\n> **Q: Could you include more regular benchmark datasets in the experiment, such as ImageNet and COCO? \nCould you consider standard model architectures for experiments, such as ResNet and ViT?** \n\nA: Thank you for your suggestion. Due to time constraints, we **conduct new experiments on Tiny-ImageNet**, a reduced version of ImageNet, and supplemented with the updated **ResNet18** model. Specifically, the experimental setup includes 50 clients, each containing 20 distinct classes to simulate non-iid data distribution, and runs for 100 communication rounds. Due to time constraints, we were unable to complete all experiments and compared with the baselines in the table below. These additional experiments further verify the **scalability of our method on more complex datasets and more sophisticated network architectures**. We believe these results can effectively demonstrate the applicability of our approach in scenarios involving complex data and advanced models. \n\n| Method   | Acc               | ECE     |\n|----------|-------------------|---------|\n| FedAvg   | 27.14 $\\pm$ 0.12  | 0.3170  |\n| FedPer   | 35.43 $\\pm$ 0.19  | 0.2890  |\n| FedProx  | 24.97 $\\pm$ 0.25  | 0.2515  |\n| SCAFFOLD | 31.37 $\\pm$ 0.22  | 0.0698  |\n| pFedME   | 36.24 $\\pm$ 0.14  | 0.0631  |\n| PerAvg   | 34.86 $\\pm$ 0.20  | 0.1774  |\n| Ours     | **37.91 $\\pm$ 0.07** | **0.0421** |\n\n> **Q: Assumptions in Appendix reduce readability, seemingly intentional. Lack of discussion on practical validity of assumptions D.1-D.7.** \n\nA: We appreciate your valuable feedback. Due to page constraints in the initial submission, we placed the technical assumptions in the appendix, which may have hindered readability. We will relocate Assumptions D.1–D.7 to the main text in the camera-ready version to enhance clarity.  \n\nIn practice, these assumptions are readily satisfied in typical FL scenarios, with their validity conditions as follows:  \n\n**D1:** This is readily satisfied in SVGD implementations using common kernels like the RBF kernel. The update directions generated by SVGD under typical settings maintain this positive definiteness, ensuring stable transformations.\n\n**D2:** This is a standard requirement for convergence in gradient-based optimization algorithms. In practice, techniques like adaptive learning rate methods (e.g., AdaGrad, used in our experiments) automatically ensure this condition is met for stable updates.\n\n**D3:** In practice, the dimension of model parameters is fixed, and the range of parameter values is constrained by initialization and training strategies (compactness); the true parameter must belong to the parameter space, so this assumption holds.\n\n**D4:** This assumption ensures the statistical distinguishability of parameters based on the data. It holds if the data likelihood $p(D_{ki}|\\theta)$ provides sufficient discriminative power between different $\\theta$, which is generally true for identifiable models and non-degenerate datasets encountered in practical FL tasks. \n\n**D5:** The constraint on the generalized bracketed entropy is used to control the deviation of the posterior distribution. In practice, when the model complexity matches the data scale, the complexity of the function class is controllable and the entropy condition will be naturally satisfied, thereby ensuring that the posterior deviation is within a reasonable range.\n\n**D6:** This assumption requires that the prior $\\Pi$ assigns sufficient mass near the true parameter $\\theta_0$ under a likelihood ratio condition. Using well-specified or weakly informative priors and reasonable initialization strategies (e.g., Kaiming, Xavier) helps satisfy this condition. It guarantees the prior doesn't contradict the true data-generating mechanism too strongly.\n\n**D7:** This property is fundamental to the definition and computation of the Wasserstein barycenter. Crucially, the 2-Wasserstein distance $W_2$ itself satisfies this convexity property on the space of probability measures with finite second moments ($P_2(\\Xi)$). This is a well-established result in optimal transport theory (see references like [Agueh \\& Carlier, 2011] or [Villani, 2008]) and is not an additional restrictive assumption on the model, but rather an inherent geometric property of the space where we perform aggregation.\n\n> **Q: Could you release the experimental code of this work anonymously?** \n\nA: We have uploaded our experimental code in the **supplementary materials**. Additionally, we will include the GitHub repository link in the camera-ready version to facilitate further verification and reuse of our work. Due to this year’s rebuttal policy, we are unable to provide any anonymous external links in the response."}]

Author rebuttal:

Meta-review: In their paper, the authors introduced Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation (PBFL-WB), a novel approach to address data heterogeneity in federated learning. The core contribution is the use of a Wasserstein barycenter to aggregate heterogeneous posterior distributions from different client devices. While the reviewers initially raised several valid concerns, the authors' thoughtful rebuttal successfully addressed these points.

(1) *Novelty and Technical Soundness*: The reviewers unanimously acknowledged the novelty of using a Wasserstein barycenter for aggregating heterogeneous posterior distributions in federated learning. This approach is both theoretically sound and a meaningful contribution to the field. The paper clearly formalizes the problem and provides a new perspective on addressing the non-IID challenges in a Bayesian context.

(2) *Strong Empirical Results*: The paper's empirical evaluation is robust and comprehensive. The experiments clearly show that PBFL-WB outperforms state-of-the-art baselines on various non-IID settings, highlighting the method's practical utility. The inclusion of new results in the rebuttal further solidified the paper's claims and convinced a reviewer to increase their score.

In summary, the paper presents an original solution to an important problem in federated learning. The authors' rigorous response to the reviewers' feedback has demonstrated the paper's technical depth and its readiness for publication. As a consequence, I recommend accepting the paper.
Decision: Accept (poster)