Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation

Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.

Q: Could you include more regular benchmark datasets in the experiment, such as ImageNet and COCO? Could you consider standard model architectures for experiments, such as ResNet and ViT?

A: 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.

Q: Assumptions in Appendix reduce readability, seemingly intentional. Lack of discussion on practical validity of assumptions D.1-D.7.

A: 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.

In practice, these assumptions are readily satisfied in typical FL scenarios, with their validity conditions as follows:

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.

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.

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.

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.

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.

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.

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.

Q: Could you release the experimental code of this work anonymously?

A: 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.

Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.

Q: No mitigation strategies proposed for scalability issues from linear communication overhead with increasing particles.

A: 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.

Q: How to choose a suitable number of particles.

A: 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.

Q: The paper's two-step update of $\hat{\bar{p}}(\theta)$ lacks clarity on optimality and associated theoretical guarantees.

A: 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].

Furthermore, 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.

[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.

Q: Need communication comparisons with other advanced algorithms.

A: 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.

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$.

A: The computational complexity of computing $T_k^*$ corresponds to that of solving the linear programming problem in Equation 3.

With 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)$.

With 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.

Thank you for your valuable suggestions. We have conducted targeted supplementary experiments and provided corresponding responses to address the points you raised.

Q: Several closely related references are missing.

A: 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.

Q: Problematic averaging posterior params due to info geometry.

A: 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.

Q: Several methods deal with non-Gaussians, e.g., mixture of gaussians.

A: 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.

Q: The experiments lack extensibility for thorough empirical verification; results on other FL settings with varied problem hyperparameters are preferred.

A: 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.

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.

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.

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).

The results consistently validate our method's superiority across diverse configurations, reinforcing its generalizability and practical utility.

Q: The experiments are only done with relatively small networks.

A: 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.

Q: Need convergence analysis results as a function of iteration count and generalization error guarantees.

A: 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.

For 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.

[26] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.

Q: The scalability limitations of particle-based methods and the practical implications beyond theoretical contributions require clarification.

A: 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.

Specifically, 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.

Q: WB may require a large number of samples for high-dimensional model parameters, making it computationally unscalable.

A: 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.

Furthermore, 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.

Q: ECE scores in Fig 2 for most methods are sufficiently good with statistically indistinct differences.

A: 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.

Dear Reviewer TsB1,

Thank you again for your evaluation of our work and the valuable feedback you have provided. We have posted a detailed rebuttal to address the concerns you raised. As the author-reviewer discussion phase will conclude in less than two days, we kindly request your feedback and are keen to know if our clarifications and additional results have effectively addressed your questions.

If any points remain unclear, we are happy to engage in further discussion. We sincerely hope that our rebuttal clarifies the merits of our work, and we would appreciate it if you would consider our response in your final evaluation and revisit your rating.

Best regards, Authors

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.