Title: Personalized Bayesian Federated Learning with Wasserstein Barycenter Aggregation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Personalized Bayesian Federated Learning

Stein Variational Gradient Descent

Wasserstein Barycenter Aggregation
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
**Strengths:**

1. This work has a strong theoretical foundation.

2. Sufficient baseline methods have been considered for comparison.

**Weakness:**

1. The experiments could include more regular benchmark datasets, such as ImageNet and COCO.

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

3. The assumptions cannot be put in the Appendix, which appears to be intentional. It also reduces the overall readability of this work.

4. It lacks the discussion of how assumptions D.1 - D.7 may hold in practice. Please add some discussion on this issue.

5. Without releasing the experimental code, the reproducibility of this work cannot be fully ensured.

QUESTIONS:
1. Could you include more regular benchmark datasets in the experiment, such as ImageNet and COCO?

2. Could you consider standard model architectures for experiments, such as ResNet and ViT?

3. It lacks the discussion of how assumptions D.1 - D.7 may hold in practice. Please add some discussions on this issue.

4. Could you release the experimental code of this work anonymously?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-05-31

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\n1. 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)?\n\nLIMITATIONS:\nThe authors states that the main limitations are the computational complexity of SVGD for nonparametric posterior approximation at clients, which was left for future work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS AND WEAKNESSES:\n[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\n\nQUESTIONS:\nSee \"Strengths And Weaknesses\" section.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: SUMMARY:
This paper proposes a bayesian personalized federated learning approach, that goes beyond previous such approaches by enabling non-parametric posterior distributions.
This is made possible through 1) non-parametric local particle based variationak inference via Stein GD, and then 2) aggregation through Wasserstein barycenters.
The method is first analyzed theoretically through non asymptotic convergence guarantees, and then in real world experiments.

STRENGTHS AND WEAKNESSES:
**Strengths**

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

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

The 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.
However, 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.

Theory and experiments seem to validate the method.

**Weaknesses**

I am not an expert (far from it) in Bayesian or variational inference, so it is hard for me to formulate technical weaknesses.

QUESTIONS:
None

LIMITATIONS:
None
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-04