Personalized Federated Learning via Variational Bayesian Inference

Xu ZhangYinchuan LiWenpeng LiKaiyang GuoYunfeng Shao

article2022ICML144 citations

Develops pFedBayes, a personalized federated learning framework that uses variational Bayesian neural networks to prevent client overfitting on non-i.i.d. data while guaranteeing minimax-optimal generalization error bounds.

Listen

Modern machine learning applications increasingly rely on federated learning to train artificial intelligence models across distributed devices while keeping private data localized. However, real-world deployments face two major challenges: private data across individual devices is statistically diverse and non-identical, and local data volumes are often too small to train complex neural networks effectively. These constraints frequently result in model overfitting and severe performance degradation, particularly in high-stakes fields such as finance, healthcare, and autonomous systems.

The article addresses these dual challenges by developing and evaluating a personalized federated learning framework named pFedBayes. The main objective is to demonstrate that integrating Bayesian neural networks—where model weights are treated as probability distributions rather than fixed values—into a two-level optimization process can simultaneously prevent overfitting on limited data and achieve tailored local personalization with quantifiable prediction uncertainty.

To establish this framework, the authors formulated a two-level variational Bayesian inference model where local clients use the aggregated global distribution as an informed prior rather than relying on arbitrary assumptions. The method was rigorously tested through theoretical analysis of generalization error bounds and empirical simulations across standard benchmark image datasets (MNIST, Fashion-MNIST, and CIFAR-10). The experimental setup evaluated performance across 10 to 20 clients under varying data volumes (small, medium, and large) and compared pFedBayes against seven state-of-the-art global and personalized federated learning baselines.

The findings show that pFedBayes consistently outperforms existing baseline methods, especially when local data is scarce. On the complex CIFAR-10 dataset with small sample sizes, pFedBayes surpassed other state-of-the-art algorithms by 11.71% in personalized model accuracy and by 3.47% in global model accuracy. It also demonstrated top performance on MNIST and Fashion-MNIST, outperforming competitors by 1.25% and 0.42% on small subsets, respectively. Additionally, the algorithm exhibited rapid and stable convergence within roughly 50 iterations on small datasets while maintaining the unique ability to quantify predictive uncertainty as training progressed.

These results demonstrate that treating network parameters probabilistically effectively controls overfitting while allowing flexible local adaptation. The ability to measure output uncertainty provides critical decision-support value for safety-critical and regulated applications where understanding model confidence is necessary. While pFedBayes delivered clear superiority in limited and medium data regimes, its global model showed reduced performance advantages on large datasets, reflecting known scaling characteristics of Bayesian neural networks that require specialized aggregation techniques.

Organizations deploying federated learning in data-constrained or heterogeneous environments should consider adopting Bayesian personalized architectures to improve local model accuracy and capture reliable uncertainty metrics. Practitioners must carefully calibrate the regularization parameter balancing personalization and global aggregation, as well as set modest local learning rates to ensure stable convergence. Future work should focus on piloting the approach in operational environments with varied client connectivity and integrating specialized scaling techniques to maintain performance when data volumes expand.

Academy of Mathematics and Systems Science, Chinese Academy of Sciences Huawei bayesian inference federated learning inference personalized federated learning variational Bayesian inference pFedBayes algorithm

arXiv: 2206.07977

Paper: Weight Uncertainty in Neural Network, Charles Blundell et al. (2015). Introduces Bayes by Backprop for variational inference over neural network weights, providing the foundational Bayesian neural network formulation that the source adapts for decentralized federated optimization.
Paper: Personalized Federated Learning with Moreau Envelopes, Canh T. Dinh et al. (2020). Formulates bi-level optimization for personalized federated learning via Moreau envelopes, establishing the core bi-level regularized objective structure that the source extends into a variational Bayesian framework.
Paper: Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach, Alireza Fallah et al. (2020). Pioneers model-agnostic meta-learning for personalized federated learning, establishing key baseline concepts for adapting global shared models to heterogeneous local clients.
Paper: Ditto: Fair and Robust Federated Learning Through Personalization, Tian Li et al. (2020). Develops a bi-level multi-task objective that regularizes local personalized models toward a global model, providing a fundamental baseline for personalized federated learning objectives.
Paper: Federated Learning with Personalization Layers, Manoj Ghuhan Arivazhagan et al. (2019). Introduces layer-based personalization in federated deep neural networks to mitigate statistical data heterogeneity across clients.
Paper: Federated Optimization in Heterogeneous Networks, Tian Li et al. (2018). Introduces proximal regularization to stabilize federated optimization under statistical heterogeneity, forming an essential foundation for regularized client optimization.
Paper: Communication-Efficient Learning of Deep Networks from Decentralized Data, H. B. McMahan et al. (2016). Introduces the standard FederatedAveraging algorithm that serves as the universal foundation for global federated aggregation.
Paper: Federated Multi-Task Learning, Virginia Smith et al. (2017). Presents a multi-task learning formulation for federated systems to address non-identical client distributions, laying the groundwork for personalized federated learning.

Paper: Personalized Federated Learning through Local Memorization, Othmane Marfoq et al. (2022). Extends personalized federated learning under client heterogeneity by leveraging localized spatial memorization within neural feature extractors.
Paper: CD2-pFed: Cyclic Distillation-guided Channel Decoupling for Model Personalization in Federated Learning, Yiqing Shen et al. (2022). Advances personalized federated learning by introducing channel decoupling and cyclic knowledge distillation across layers for heterogeneous client domains.
Paper: FedTGP: Trainable Global Prototypes with Adaptive-Margin-Enhanced Contrastive Learning for Data and Model Heterogeneity in Federated Learning, Jianqing Zhang et al. (2024). Builds on client personalization challenges under statistical heterogeneity by developing trainable global prototype representations via margin-enhanced contrastive learning.

Cover for Personalized Federated Learning via Variational Bayesian Inference

Abstract

Federated learning faces huge challenges from model overfitting due to the lack of data and statistical diversity among clients. To address these challenges, this paper proposes a novel personalized federated learning method via Bayesian variational inference named pFedBayes. To alleviate the overfitting, weight uncertainty is introduced to neural networks for clients and the server. To achieve personalization, each client updates its local distribution parameters by balancing its construction error over private data and its KL divergence with global distribution from the server. Theoretical analysis gives an upper bound of averaged generalization error and illustrates that the convergence rate of the generalization error is minimax optimal up to a logarithmic factor. Experiments show that the proposed method outperforms other advanced personalized methods on personalized models, e.g., pFedBayes respectively outperforms other SOTA algorithms by 1.25%, 0.42% and 11.71% on MNIST, FMNIST and CIFAR-10 under non-i.i.d. limited data.

Topics

bayesian inference federated learning inference neural networks optimization personalizations private data regularization parameters small sample size variational bayesian

Keywords

personalized federated learning variational Bayesian inference pFedBayes algorithm multi-task learning generalization error posterior distribution prior distribution MNIST dataset

1. Introduction
1.1. Main Contributions
1.2. Related Works
2. Personalized Bayesian Federated Learning Model with Gaussian Distribution
3. Theoretical Analysis
4. Algorithm
5. Experiments
5.1. Experimental Setting
5.2. Experimental Hyperparameter Settings
5.3. Performance Comparison Results
6. Conclusions
Acknowledgements
References
A. Proof of Lemmas
A.1. Proof of Lemma 1
A.2. Proof of Lemma 2
A.3. Proof of Lemma A.2
B. Experimental Results on MNIST Dataset
B.1. Effect of Hyperparameters
B.2. pFedBayes Uncertainty Estimation

Knowls

Knowl 1 — pFedBayes Two-Level Variational Bayesian Federated Learning Formulation
model/method
pFedBayes formulates personalized federated learning as a two-level variational Bayesian optimization problem across a central server and $N$ decentralized clients. Each client $i \in \{1, \dots, N\}$ observes a private dataset $D^i = \{(x_j^i, y_j^i)\}_{j=1}^n$ generated according to $y_j^i = f^i(x_j^i) + \varepsilon_j^i$ , where $x_j^i \in \mathbb{R}^{s_0}$ , $y_j^i \in \mathbb{R}^{s_{L+1}}$ , and observation noise $\varepsilon_j^i \sim \mathcal{N}(0, \sigma_\varepsilon^2 I)$ . Instead of deterministic parameter vectors, each client parameterizes its neural network weights $\theta \in \mathbb{R}^T$ as a Bayesian Neural Network with a local variational distribution $q^i(\theta) \in \mathcal{Q}_i$ . The server maintains a global variational distribution $w(\theta) \in \mathcal{Q}_w$ that acts as an empirical prior for all clients.

The global and local objectives are structured as:

$\text{Server subproblem:}\quad \min_{w(\theta) \in \mathcal{Q}_w} \left( F(w) \triangleq \frac{1}{N} \sum_{i=1}^N F_i(w) \right)$

$\text{Client subproblem:}\quad F_i(w) \triangleq \min_{q^i(\theta) \in \mathcal{Q}_i} \left\{ -\mathbb{E}_{q^i(\theta)} \left[ \log p_\theta^i(D^i) \right] + \zeta \operatorname{KL}(q^i(\theta) \parallel w(\theta)) \right\}$

where $p_\theta^i(D^i)$ denotes the local likelihood function, $\operatorname{KL}(\cdot \parallel \cdot)$ is the Kullback-Leibler divergence, and $\zeta \ge 1$ is a regularization tradeoff hyperparameter that balances client personalization (fitting local private data) against global aggregation (anchoring to the shared server distribution).

Under a mean-field Gaussian family decomposition, the local and global weight distributions for each weight component $m \in \{1, \dots, T\}$ are defined by:

$q^i(\theta) = \prod_{m=1}^T \mathcal{N}(\theta_m; \mu_{i,m}, \sigma_{i,m}^2), \quad w(\theta) = \prod_{m=1}^T \mathcal{N}(\theta_m; \mu_{w,m}, \sigma_{w,m}^2)$

which gives the closed-form coordinate-wise KL divergence:

$\operatorname{KL}(q^i(\theta) \parallel w(\theta)) = \frac{1}{2} \sum_{m=1}^T \left[ \log \left( \frac{\sigma_{w,m}^2}{\sigma_{i,m}^2} \right) + \frac{\sigma_{i,m}^2 + (\mu_{i,m} - \mu_{w,m})^2}{\sigma_{w,m}^2} - 1 \right]$
Knowl 2 — pFedBayes Stochastic Gradient Optimization Algorithm
algorithm
The pFedBayes optimization procedure solves the two-level variational inference objective using stochastic gradient descent and the reparameterization trick. To ensure positive standard deviations, the variational parameter vector is defined as $v = (\mu, \rho) \in \mathbb{R}^{2T}$ , where $\sigma_m = \log(1 + \exp(\rho_m))$ for each dimension $m \in \{1, \dots, T\}$ . Weight vectors are sampled via $\theta_m = h(v_m, g_m) = \mu_m + \log(1 + \exp(\rho_m)) \cdot g_m$ with standard normal noise $g_m \sim \mathcal{N}(0, 1)$ .

Each client estimates the expected negative log-likelihood over a mini-batch $\Lambda \subset D^i$ of size $b$ using $K$ Monte Carlo samples:

$\Omega^i(v) \approx -\frac{n}{b} \frac{1}{K} \sum_{j=1}^b \sum_{k=1}^K \log p_{h(v, g_k)}^i(D_j^i) + \zeta \operatorname{KL}(q_v^i(\theta) \parallel w_v(\theta))$

Clients alternate between updating their local personalized variational parameters with $\nabla_v \Omega^i(v)$ and updating their localized global parameters $v_w^i$ using $\nabla_v \Omega_w^i(v_w^i) = \nabla_v \operatorname{KL}(q_v^i(\theta) \parallel w_v^i(\theta))$ for $R$ local iterations before transmitting localized global parameters to the server.
Cloud server executes:
    Input: Communication rounds T_comm, local iterations R, subset size S, sample size K, batch size b, smoothing parameter \beta, initial server parameters v^0 = (\mu^0, \rho^0)
    for t = 0, 1, ..., T_comm - 1 do
        for each client i in {1, ..., N} in parallel do
            v_{w, R}^{t+1, i} <- ClientUpdate(i, v^t)
        S^t <- randomly chosen subset of {1, ..., N} of size S
        v^{t+1} = (1 - \beta) v^t + (\beta / S) * sum_{i in S^t} v_{w, R}^{t+1, i}
ClientUpdate(i, v^t):
    Initialize localized global model: v_{w, 0}^t = v^t
    for r = 0, 1, ..., R - 1 do
        Sample a mini-batch D_\Lambda^i of size b from D^i
        Draw K noise samples g_{i,r,k} ~ N(0, I_T) for k = 1, ..., K
        Compute personalized objective \Omega^i(v_r^t) using mini-batch D_\Lambda^i and noise draws
        Compute gradient \nabla_v \Omega^i(v_r^t) and update local parameters v_{r+1}^t
        Compute localized global objective \Omega_w^i(v_{w, r}^t) = KL(q_{v_r^t}^i(\theta) || w_{v_{w, r}^t}( heta))
        Compute gradient \nabla_v \Omega_w^i(v_{w, r}^t) and update localized global parameters v_{w, r+1}^t
    return v_{w, R}^t to the cloud server
Knowl 3 — Averaged Generalization Error Upper Bound for pFedBayes
theoretical result
Let each client's neural network have depth $L$ , equal hidden layer widths $s_j = M$ for $j \in \{1, \dots, L\}$ , input dimension $s_0$ , 1-Lipschitz continuous activation functions $\sigma(\cdot)$ , bounded weight infinity-norm $\|\theta\|_\infty \le B$ , and total parameter count $T$ . Assume each of the $N$ clients has $n$ private training samples, and let $H = BM$ .

Define the estimation error terms:

$r_n = \frac{(L+1)T}{n} \log M + \frac{T}{n} \log\left(s_0 \sqrt{\frac{n}{T}}\right)$

$\varepsilon_n = n^{-1/2} \sqrt{(L+1)T \log M + T \log\left(s_0 \sqrt{\frac{n}{T}}\right)} \log^\delta(n) \quad (\text{for } \delta > 1)$

and the client function approximation error:

$\xi_n^i = \inf_{\theta \in \Theta(L, s), \|\theta\|_\infty \le B} \|f_\theta^i - f^i\|_\infty^2$

Let $\hat{q}^i(\theta)$ be the optimal variational solution for client $i$ and $w^\star(\theta)$ the optimal global distribution. Under the squared Hellinger distance $d^2(P_\theta^i, P^i) = \mathbb{E}_{X^i} \left[ 1 - \exp\left( -\frac{[f_\theta^i(X^i) - f^i(X^i)]^2}{8\sigma_\varepsilon^2} \right) \right]$ , with dominating probability:

$\frac{1}{N} \sum_{i=1}^N \int_\Theta d^2(P_\theta^i, P^i) \hat{q}^i(\theta) \, d\theta \le C \varepsilon_n^2 + C' r_n + \frac{C''}{N \zeta} \sum_{i=1}^N \xi_n^i$

where $\zeta \ge 1$ is the tradeoff parameter, $C > 0$ is an absolute constant, and $C', C''$ are diverging sequences. The upper bound separates into estimation error on the order of $\mathcal{O}(1/n)$ and approximation error governed by network capacity $T$ .
Knowl 4 — Minimax-Optimal Generalization Convergence Rate of pFedBayes
theoretical result
Suppose the true client mapping functions $\{f^i\}_{i=1}^N$ are $\beta$ -Hölder smooth functions defined on a compact domain with intrinsic data dimension $d$ . The individual approximation error is bounded by $\xi_n^i \le C_0 T^{-2\beta/d}$ for constant $C_0 > 0$ .

By selecting the total number of neural network parameters as a function of sample size $n$ according to $T = C_1 n^{d/(2\beta + d)}$ with constant $C_1 > 0$ , the averaged generalization error under the $L^2$ norm satisfies the upper bound:

$\frac{1}{N} \sum_{i=1}^N \int_\Theta \|f_\theta^i(X^i) - f^i(X^i)\|_{L^2}^2 \hat{q}^i(\theta) \, d\theta \le C_2 n^{-\frac{2\beta}{2\beta + d}} \log^{2\delta'}(n)$

where $\delta' > \delta > 1$ and $C_2 > 0$ is a constant depending on $s_0, \beta, d, L, M, \zeta,$ and $n$ .

Furthermore, for bounded target functions $\|f^i\|_\infty \le F$ and network outputs $\|f_\theta^i\|_\infty \le F$ , this rate matches the minimax lower bound:

$\inf_{\{\|f_\theta^i\|_\infty \le F\}_{i=1}^N} \sup_{\{\|f^i\|_\infty \le F\}_{i=1}^N} \frac{1}{N} \sum_{i=1}^N \int_\Theta \|f_\theta^i(X^i) - f^i(X^i)\|_{L^2}^2 \hat{q}^i(\theta) \, d\theta \ge C_3 n^{-\frac{2\beta}{2\beta + d}}$

where $C_3 > 0$ is a constant. Consequently, the convergence rate of the generalization error in pFedBayes is minimax optimal up to a polylogarithmic factor $\log^{2\delta'}(n)$ .
Knowl 5 — Closed-Form Optimal Global Variational Distribution under Gaussian Mean-Field Assumption
theoretical result
Let the client variational distributions $q^i(\theta) = \prod_{m=1}^T \mathcal{N}(\theta_m; \mu_{i,m}, \sigma_{i,m}^2)$ for $i \in \{1, \dots, N\}$ and the global distribution $w(\theta) = \prod_{m=1}^T \mathcal{N}(\theta_m; \mu_{w,m}, \sigma_{w,m}^2)$ belong to the mean-field Gaussian family. The server optimization problem:

$\min_{w(\theta) \in \mathcal{Q}_w} \frac{1}{N} \sum_{i=1}^N \operatorname{KL}(q^i(\theta) \parallel w(\theta))$

has a unique analytical optimal solution for every parameter index $m \in \{1, \dots, T\}$ :

$\mu_{w,m} = \frac{1}{N} \sum_{i=1}^N \mu_{i,m}$

$\sigma_{w,m}^2 = \frac{1}{N} \sum_{i=1}^N \left[ \sigma_{i,m}^2 + (\mu_{i,m} - \mu_{w,m})^2 \right] = \frac{1}{N} \sum_{i=1}^N \left( \sigma_{i,m}^2 + \mu_{i,m}^2 \right) - \mu_{w,m}^2$

The optimal global mean is the arithmetic mean of client means, while the optimal global variance is the average of local variances augmented by the empirical variance of the client means across the population.

Knowl 6 — Classification Accuracy Across Data Heterogeneity and Sample Size Regimes

data/table

The test accuracy performance of pFedBayes was evaluated against non-personalized methods (FedAvg, Fedprox, BNFed) and personalized federated methods (Per-FedAvg, pFedMe, HeurFedAMP, pFedGP) across MNIST, Fashion-MNIST (FMNIST), and CIFAR-10 datasets under non-I.I.D. partitions (5 of 10 labels per client) across small, medium, and large local data splits.

Dataset	Method	Small (Acc. (%))		Medium (Acc. (%))		Large (Acc. (%))
Dataset	Method	PM	GM	PM	GM	PM	GM
MNIST	FedAvg	-	$87.38\pm0.27$	-	$90.60\pm0.19$	-	$92.39\pm0.24$
	Fedprox	-	$87.65\pm0.30$	-	$90.66\pm0.17$	-	$92.42\pm0.23$
	BNFed	-	$78.70\pm0.69$	-	$80.02\pm0.60$	-	$82.95\pm0.22$
	Per-FedAvg	$89.29\pm0.59$	-	$95.19\pm0.33$	-	$98.27\pm0.08$	-
	pFedMe	$92.88\pm0.04$	$87.35\pm0.08$	$95.31\pm0.17$	$89.67\pm0.34$	$96.42\pm0.08$	$91.25\pm0.14$
	HeurFedAMP	$90.89\pm0.17$	-	$94.74\pm0.07$	-	$96.90\pm0.12$	-
	pFedGP	$85.96\pm2.30$	-	$91.96\pm0.97$	-	$95.66\pm0.43$	-
	Ours (pFedBayes)	$\mathbf{94.13\pm0.27}$	$\mathbf{90.44\pm0.45}$	$\mathbf{97.09\pm0.13}$	$\mathbf{92.33\pm0.76}$	$\mathbf{98.79\pm0.13}$	$\mathbf{94.39\pm0.32}$
FMNIST	FedAvg	-	$81.51\pm0.19$	-	$83.90\pm0.13$	-	$\mathbf{85.42\pm0.14}$
	Fedprox	-	$81.53\pm0.08$	-	$83.92\pm0.21$	-	$85.32\pm0.14$
	BNFed	-	$66.54\pm0.64$	-	$69.68\pm0.39$	-	$70.10\pm0.24$
	Per-FedAvg	$79.79\pm0.83$	-	$84.90\pm0.47$	-	$88.51\pm0.28$	-
	pFedMe	$88.63\pm0.07$	$81.06\pm0.14$	$91.32\pm0.08$	$83.45\pm0.21$	$92.02\pm0.07$	$84.41\pm0.08$
	HeurFedAMP	$86.38\pm0.24$	-	$89.82\pm0.16$	-	$92.17\pm0.12$	-
	pFedGP	$86.99\pm0.41$	-	$90.53\pm0.35$	-	$92.22\pm0.13$	-
	Ours (pFedBayes)	$\mathbf{89.05\pm0.17}$	$80.17\pm0.19$	$\mathbf{91.95\pm0.02}$	$82.33\pm0.37$	$\mathbf{93.01\pm0.10}$	$83.30\pm0.28$
CIFAR-10	FedAvg	-	$44.24\pm3.01$	-	$56.73\pm1.81$	-	$\mathbf{79.05\pm0.44}$
	Fedprox	-	$43.70\pm1.38$	-	$57.35\pm3.11$	-	$77.65\pm1.62$
	BNFed	-	$34.00\pm0.16$	-	$39.52\pm0.56$	-	$44.37\pm0.19$
	Per-FedAvg	$33.96\pm1.12$	-	$52.98\pm1.21$	-	$69.61\pm1.21$	-
	pFedMe	$49.66\pm1.53$	$43.67\pm2.14$	$66.75\pm1.87$	$51.18\pm2.57$	$77.13\pm1.06$	$70.86\pm1.04$
	HeurFedAMP	$46.72\pm0.39$	-	$59.94\pm1.42$	-	$73.24\pm0.80$	-
	pFedGP	$43.66\pm0.32$	-	$58.54\pm0.40$	-	$72.45\pm0.19$	-
	Ours (pFedBayes)	$\mathbf{61.37\pm1.40}$	$\mathbf{47.71\pm1.19}$	$\mathbf{73.94\pm0.97}$	$\mathbf{60.84\pm1.26}$	$\mathbf{83.46\pm0.13}$	$64.40\pm1.22$

PM denotes Personalized Model accuracy and GM denotes Global Model accuracy. On small data regimes where overfitting is severe, pFedBayes achieves substantial gains over the next best algorithm on personalized accuracy: +1.25% on MNIST, +0.42% on FMNIST, and +11.71% on CIFAR-10.

Knowl 7 — Non-I.I.D. and Sample-Size Benchmarking Setup
experimental setup
The empirical validation of pFedBayes uses three image classification benchmarks with partitioned non-I.I.D. client assignments:
- Client Partitioning and Heterogeneity: For MNIST and Fashion-MNIST (FMNIST), data are partitioned among $N = 10$ clients; for CIFAR-10, data are partitioned among $N = 20$ clients. Non-I.I.D. label skew is established by restricting each client to contain samples from exactly 5 of the 10 available classes. The active client subset size selected per communication round is $S = 10$ .
- Sample Size Regimes:
  - MNIST / FMNIST: Small (50 training, 950 test samples per class), Medium (200 training, 800 test samples per class), Large (900 training, 300 test samples per class).
  - CIFAR-10: Small (25 training, 475 test samples per class), Medium (100 training, 400 test samples per class), Large (450 training, 150 test samples per class).
- Network Architectures:
  - MNIST / FMNIST: Fully-connected Deep Neural Network with 3 layers $[784, 100, 10]$ and ReLU activations before a final softmax layer.
  - CIFAR-10: Convolutional VGG network configured as [16, 'M', 32, 'M', 64, 'M', 128, 'M', 128, 'M'] with max-pooling layers.
- Evaluation Metric: Peak test accuracy evaluated over 700 to 800 global communication rounds.
Knowl 8 — Hyperparameter Sensitivity and Personalization Trade-Off in pFedBayes
empirical result
A grid search over hyperparameter configurations on the medium MNIST benchmark demonstrates the impact of key pFedBayes parameters:
- Personalization Weight $\zeta$ : Adjusts the balance between fitting local data and conforming to the global empirical prior. On medium MNIST with $\rho = -2.5$ and $\eta_1 = \eta_2 = 0.001$ , increasing $\zeta \in \{0.5, 1, 5, 10, 20\}$ shifts performance toward the global model: global model accuracy rises from $89.88\%$ (at $\zeta = 0.5$ ) to $93.22\%$ (at $\zeta = 10$ ) while personalized model accuracy remains high ( $97.13\%$ at $\zeta=0.5$ , peak $97.41\%$ at $\zeta=1$ , $97.18\%$ at $\zeta=10$ , and $97.04\%$ at $\zeta=20$ ). $\zeta = 10$ provides the best joint trade-off.
- Weight Variance Parameter Initialization $\rho$ : Evaluated over $\{-1, -1.5, -2, -2.5, -3\}$ . Initializing with $\rho = -2.5$ achieved the optimal balance of personalized model accuracy ( $97.18\%$ ) and global model accuracy ( $93.22\%$ ).
- Learning Rates $\eta_1, \eta_2$ : The personalized model learning rate $\eta_1$ and localized global model learning rate $\eta_2$ yielded the highest convergence and stability when set to $\eta_1 = \eta_2 = 0.001$ .
Knowl 9 — Model Uncertainty Quantification in Client Predictions
empirical result
pFedBayes enables output uncertainty quantification across training communication rounds by leveraging Bayesian model averaging over its posterior weight distributions. At round 0 (prior to training), predicted class probability distributions across distinct clients exhibit high entropy and diffuse uncertainty over output classes. By round 1 and progressing to round 10, predictive probability masses concentrate decisively on true class labels with minimal predictive variance, providing well-calibrated confidence estimates that can be utilized to evaluate client data quality and determine selective model aggregation in federated networks.
Knowl 10 — Global Model Aggregation Degradation on Large Datasets
limitation
While pFedBayes achieves superior personalized model performance across all sample sizes and superior global model performance on small and medium datasets, its aggregated global model degrades relative to deterministic baselines on large datasets. Specifically, on CIFAR-10 with large sample sizes, pFedBayes achieves $83.46\%$ personalized accuracy (outperforming FedAvg's $79.05\%$ and pFedMe's $77.13\%$ ) but its global model achieves only $64.40\%$ accuracy compared to $79.05\%$ for FedAvg and $77.65\%$ for Fedprox. This occurs because mean-field Gaussian Bayesian Neural Networks require additional deep Bayesian sampling or aggregation mechanisms to match deterministic networks when data are abundant, whereas pFedBayes maintains uniform aggregation procedures across small and large data regimes.

Coverage note — No substantial contributed material was omitted; intermediate technical lemmas (Lemmas 1, 2, and A.1) are represented through the main theoretical bounds and exact server formulas.

References

1.Achituve, I., Shamsian, A., Navon, A., Chechik, G., and Fetaya, E. Personalized federated learning with gaussian processes. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
2.Al-Shedivat, M., Gillenwater, J., Xing, E., and Rostamizadeh, A. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Representations, 2021.
3.Alquier, P. and Ridgway, J. Concentration of tempered posteriors and of their variational approximations. The Annals of Statistics, 48(3):1475–1497, 2020.
4.Arivazhagan, M. G., Aggarwal, V., Singh, A. K., and Choudhary, S. Federated learning with personalization layers. arXiv preprint arXiv:1912.00818, 2019.
5.Bai, J., Song, Q., and Cheng, G. Efficient variational inference for sparse deep learning with theoretical guarantee. Advances in Neural Information Processing Systems, 33, 2020.
6.Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859–877, 2017.
7.Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural networks. In Proceedings of the 32Nd International Conference on International Conference on Machine Learning - Volume 37, ICML’15, pp. 1613–1622. JMLR.org, 2015a.
8.Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network. In International Conference on Machine Learning, pp. 1613–1622. PMLR, 2015b.
9.Boucheron, S., Lugosi, G., and Massart, P. Concentration inequalities: A nonasymptotic theory of independence. Oxford university press, 2013.
10.Chen, F., Luo, M., Dong, Z., Li, Z., and He, X. Federated meta-learning with fast convergence and efficient communication. arXiv preprint arXiv:1802.07876, 2018.
11.Chen, H.-Y. and Chao, W.-L. FedBE: Making Bayesian model ensemble applicable to federated learning. In International Conference on Learning Representations, 2021.
12.Chérief-Abdellatif, B.-E. Convergence rates of variational inference in sparse deep learning. In International Conference on Machine Learning, pp. 1831–1842. PMLR, 2020.
13.Chérief-Abdellatif, B.-E. and Alquier, P. Consistency of variational bayes inference for estimation and model selection in mixtures. Electronic Journal of Statistics, 12(2):2995–3035, 2018.
14.Dai, X., Yan, X., Zhou, K., Yang, H., Ng, K. K., Cheng, J., and Fan, Y. Hyper-sphere quantization: Communication-efficient sgd for federated learning. arXiv preprint arXiv:1911.04655, 2019.
15.Fallah, A., Mokhtari, A., and Ozdaglar, A. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557–3568, 2020.
16.Guha, N., Talwalkar, A., and Smith, V. One-shot federated learning. arXiv preprint arXiv:1902.11175, 2019.
17.Hanzely, F. and Richtárik, P. Federated learning of a mixture of global and local models. arXiv preprint arXiv:2002.05516, 2020.
18.Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. Beta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2017.
19.Huang, Y., Chu, L., Zhou, Z., Wang, L., Liu, J., Pei, J., and Zhang, Y. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7865–7873, 2021.
20.Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, L. K. An introduction to variational methods for graphical models. Machine Learning, 37:183–233, 1999.
21.Jospin, L. V., Buntine, W., Boussaid, F., Laga, H., and Bennamoun, M. Hands-on bayesian neural networks-a tutorial for deep learning users. ACM Comput. Surv, 1(1), 2020.
22.Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pp. 5132–5143. PMLR, 2020.
23.Khan, M. E. Deep learning with bayesian principles. Tutorial on Advances in Neural Information Processing Systems, 2019.
24.Krizhevsky, A. Learning multiple layers of features from tiny images. Master’s thesis, University of Tront, 2009.
25.LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
26.LeCun, Y., Cortes, C., and Burges, C. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2, 2010.
27.Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V. Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06127, 2018.
28.Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smithy, V. Feddane: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers, pp. 1227–1231. IEEE, 2019.
29.Li, T., Sahu, A. K., Talwalkar, A., and Smith, V. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50–60, 2020.
30.Li, X.-C., Xu, Y.-C., Song, S., Li, B., Li, Y., Shao, Y., and Zhan, D.-C. Federated learning with position-aware neurons. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10082–10091, 2022a.
31.Li, Y., Liu, X., Zhang, X., Shao, Y., Wang, Q., and Geng, Y. Personalized federated learning via maximizing correlation with sparse and hierarchical extensions. arXiv preprint arXiv:2107.05330, 2021.
32.Li, Z., Lu, J., Luo, S., Zhu, D., Shao, Y., Li, Y., Zhang, Z., and Wu, C. Mining latent relationships among clients: Peer-to-peer federated learning with adaptive neighbor matching. arXiv preprint arXiv:2203.12285, 2022b.
33.Liu, L., Zheng, F., Chen, H., Qi, G.-J., Huang, H., and Shao, L. A bayesian federated learning framework with online laplace approximation. arXiv preprint arXiv:2102.01936, 2021.
34.Liu, X., Li, Y., Shao, Y., and Wang, Q. Sparse federated learning with hierarchical personalization models. arXiv preprint arXiv:2203.13517, 2022.
35.MacKay, D. J. A practical bayesian framework for backpropagation networks. Neural computation, 4(3):448–472, 1992.
36.Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32:13153–13164, 2019.
37.McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pp. 1273–1282. PMLR, 2017.
38.Nakada, R. and Imaizumi, M. Adaptive approximation and generalization of deep neural network with intrinsic dimensionality. Journal of Machine Learning Research, 21(174):1–38, 2020.
39.Neal, R. M. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
40.Osawa, K., Swaroop, S., Khan, M. E. E., Jain, A., Eschenhagen, R., Turner, R. E., and Yokota, R. Practical deep learning with bayesian principles. Advances in neural information processing systems, 32, 2019.
41.Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
42.Pati, D., Bhattacharya, A., and Yang, Y. On statistical optimality of variational bayes. In International Conference on Artificial Intelligence and Statistics, pp. 1579–1588. PMLR, 2018.
43.Polson, N. G. and Ročková, V. Posterior concentration for sparse deep learning. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 938–949, 2018.
44.Reisizadeh, A., Mokhtari, A., Hassani, H., Jadbabaie, A., and Pedarsani, R. Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization. In International Conference on Artificial Intelligence and Statistics, pp. 2021–2031. PMLR, 2020.
45.Rothchild, D., Panda, A., Ullah, E., Ivkin, N., Stoica, I., Braverman, V., Gonzalez, J., and Arora, R. Fetchsgd: Communication-efficient federated learning with sketching. In International Conference on Machine Learning, pp. 8253–8265. PMLR, 2020.
46.Sattler, F., Wiedemann, S., Müller, K.-R., and Samek, W. Robust and communication-efficient federated learning from non-iid data. IEEE transactions on neural networks and learning systems, 31(9):3400–3413, 2019.
47.Sattler, F., Müller, K.-R., and Samek, W. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems, 32(8):3710–3722, 2021. doi: 10.1109/TNNLS.2020.3015958.
48.Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014.
49.Smith, V., Chiang, C.-K., Sanjabi, M., and Talwalkar, A. S. Federated multi-task learning. In Advances in neural information processing systems, pp. 4424–4434, 2017.
50.T Dinh, C., Tran, N., and Nguyen, T. D. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 33, 2020.
51.Thorgeirsson, A. T. and Gauterin, F. Probabilistic predictions with federated learning. Entropy, 23(1):41, 2021.
52.Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
53.Yurochkin, M., Agarwal, M., Ghosh, S., Greenewald, K., Hoang, N., and Khazaeni, Y. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, pp. 7252–7261. PMLR, 2019.
54.Zhang, X., Hong, M., Dhople, S., Yin, W., and Liu, Y. Fedpd: A federated learning framework with optimal rates and adaptivity to non-iid data. arXiv preprint arXiv:2005.11418, 2020.
55.Zong, H., Wang, Q., Liu, X., Li, Y., and Shao, Y. Communication reducing quantization for federated learning with local differential privacy mechanism. In 2021 IEEE/CIC International Conference on Communications in China (ICCC), pp. 75–80. IEEE, 2021.

Citation

MLA

Zhang, X., et al. “Personalized Federated Learning via Variational Bayesian Inference”. International Conference on Machine Learning, vol. 162, 2022, pp. 26293–310, https://proceedings.mlr.press/v162/zhang22o.html.

APA

Zhang, X., Li, Y., Li, W., Guo, K., & Shao, Y. (2022). Personalized Federated Learning via Variational Bayesian Inference. International Conference on Machine Learning, 162, 26293–26310. https://proceedings.mlr.press/v162/zhang22o.html

Chicago

Zhang, X., Y. Li, W. Li, K. Guo, and Y. Shao. 2022. “Personalized Federated Learning via Variational Bayesian Inference”. International Conference on Machine Learning 162: 26293–310. https://proceedings.mlr.press/v162/zhang22o.html.

Harvard

Zhang, X. et al. (2022) “Personalized Federated Learning via Variational Bayesian Inference”, International Conference on Machine Learning. PMLR, pp. 26293–26310. Available at: https://proceedings.mlr.press/v162/zhang22o.html.

Vancouver

1. Zhang X, Li Y, Li W, Guo K, Shao Y (2022) Personalized Federated Learning via Variational Bayesian Inference. In: International Conference on Machine Learning. PMLR, pp 26293–26310

BibTeX

@InProceedings{pmlr-v162-zhang22o,
  title = 	 {Personalized Federated Learning via Variational {B}ayesian Inference},
  author =       {Zhang, Xu and Li, Yinchuan and Li, Wenpeng and Guo, Kaiyang and Shao, Yunfeng},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  pages = 	 {26293--26310},
  year = 	 {2022},
  editor = 	 {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {17--23 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v162/zhang22o/zhang22o.pdf},
  url = 	 {https://proceedings.mlr.press/v162/zhang22o.html},
  abstract = 	 {Federated learning faces huge challenges from model overfitting due to the lack of data and statistical diversity among clients. To address these challenges, this paper proposes a novel personalized federated learning method via Bayesian variational inference named pFedBayes. To alleviate the overfitting, weight uncertainty is introduced to neural networks for clients and the server. To achieve personalization, each client updates its local distribution parameters by balancing its construction error over private data and its KL divergence with global distribution from the server. Theoretical analysis gives an upper bound of averaged generalization error and illustrates that the convergence rate of the generalization error is minimax optimal up to a logarithmic factor. Experiments show that the proposed method outperforms other advanced personalized methods on personalized models, e.g., pFedBayes respectively outperforms other SOTA algorithms by 1.25%, 0.42% and 11.71% on MNIST, FMNIST and CIFAR-10 under non-i.i.d. limited data.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Personalized Federated Learning via Variational Bayesian Inference

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — pFedBayes Two-Level Variational Bayesian Federated Learning Formulation

Knowl 2 — pFedBayes Stochastic Gradient Optimization Algorithm

Knowl 3 — Averaged Generalization Error Upper Bound for pFedBayes

Knowl 4 — Minimax-Optimal Generalization Convergence Rate of pFedBayes

Knowl 5 — Closed-Form Optimal Global Variational Distribution under Gaussian Mean-Field Assumption

Knowl 6 — Classification Accuracy Across Data Heterogeneity and Sample Size Regimes

Knowl 7 — Non-I.I.D. and Sample-Size Benchmarking Setup

Knowl 8 — Hyperparameter Sensitivity and Personalization Trade-Off in pFedBayes

Knowl 9 — Model Uncertainty Quantification in Client Predictions

Knowl 10 — Global Model Aggregation Degradation on Large Datasets

References

Citation

BibTeX

Access the Paper