Personalized Federated Learning via Variational Bayesian Inference

Xu Zhang; Yinchuan Li; Wenpeng Li; Kaiyang Guo; Yunfeng Shao

Source-linked AI summary

Personalized Federated Learning via Variational Bayesian Inference

Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, Yunfeng Shao

arXiv:2206.07977v1Jun 16, 2022cs.LG

TL;DR

Client data are often non-i.i.d. and limited; personalized methods can overfit with limited data, while Bayesian federated-learning methods can degrade under client diversity. pFedBayes represents client and server networks with parameter distributions and balances each client’s private-data fit against KL divergence from the global distribution. On small, medium, and large MNIST datasets, pFedBayes’s personalized model outperforms other SOTA comparison algorithms by 1.25%, 1.78%, and 0.52%, respectively.

Problem
Client data are often non-i.i.d. and limited; personalized methods can overfit with limited data, while Bayesian federated-learning methods can degrade under client diversity.
Source: §1 · p. 1
Method
pFedBayes represents client and server networks with parameter distributions and balances each client’s private-data fit against KL divergence from the global distribution.
Source: §2 · pp. 3–4
Results
On small, medium, and large MNIST datasets, pFedBayes’s personalized model outperforms other SOTA comparison algorithms by 1.25%, 1.78%, and 0.52%, respectively.
Source: §5.3 · p. 7
Takeaways & Limitations
The method provides personalized models and output-uncertainty quantification, with reported performance advantages especially when data are limited.
Source: §6 · p. 9
For large datasets, Bayesian neural networks need other techniques to match non-Bayesian performance; using the same techniques as for small datasets results in performance degradation when aggregating models.
Source: §5.3 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

The introduction positions pFedBayes as a variational-Bayesian personalized federated-learning method for addressing client heterogeneity and limited-data overfitting together. It presents the method, its theoretical guarantee, and its experimental performance alongside limitations of earlier approaches.

1. Introduction

Client data are often non-i.i.d. and limited, motivating an approach to client heterogeneity and overfitting together.Source: §1 · p. 1

1.1. Main Contributions

The method uses a global distribution as the prior for each client’s local distribution, which balances private-data fit against divergence from the global distribution and quantifies output uncertainty.The paper describes Bayesian model averaging as the means of quantifying output uncertainty.Source: §1.1 · pp. 1–2, §1 · p. 1
The averaged generalization-error bound has estimation error O(1/n), and a suitable parameter count yields minimax-optimal convergence up to a logarithmic term.Source: §1.1 · p. 2
pFedBayes attains the highest personalized-model accuracy on MNIST, FMNIST, and CIFAR-10 across small, medium, and large data volumes, with competitive global models on small and medium volumes.The experiments compare pFedBayes with other state-of-the-art personalization algorithms.Source: §1.1 · p. 2

1.2. Related Works

Earlier global federated-learning methods perform poorly on non-i.i.d. client data, while personalized methods address heterogeneity but can overfit when client data are limited.Source: §1.2 · p. 2
Bayesian neural networks and Bayesian federated-learning methods address overfitting, but existing approaches can perform poorly on non-i.i.d. data.Source: §1.2 · pp. 2–3

2. Personalized Bayesian Federated Learning Model with Gaussian Distribution

The model represents client and server parameters as distributions and uses the global distribution to guide personalized client distributions. Clients balance private-data fit with divergence from that global distribution, under shared architecture and data-setting assumptions.

At each iteration, clients upload their parameter distributions and download the aggregated global distribution, which serves as the prior for personalized models.Source: §2 · p. 3
The analysis assumes clients share a fully connected network architecture with bounded parameters, as well as equal data counts and equal noise variances.Source: §2 · p. 3
The local objective balances reconstruction error on private data against KL divergence from the global distribution, which replaces an assumed prior.The server minimizes the sum of KL divergences between the global distribution and local distributions; the replacement avoids requiring a prior compatible with collected data.Source: §2 · p. 4
Each client and the server use Bayesian neural networks with Gaussian parameter distributions, under a mean-field decomposition assumption.Source: §2 · pp. 3–4

3. Theoretical Analysis

The analysis bounds averaged generalization error and establishes a minimax-optimal convergence rate up to a logarithmic term under stated assumptions. It also identifies a trade-off between the bound and personalization.

3. Theoretical Analysis

The analysis assumes equal-width networks and a 1-Lipschitz activation function, alongside additional conditions on network and sample parameters.The stated assumptions also require s0, n, M, and L to be sufficiently large.Source: §3 · p. 4
The upper bound contains estimation and approximation error terms, which respond oppositely to increases in network parameter count T.Estimation error increases with T but decreases with sample size n; approximation error decreases as T increases.Source: §3 · p. 5
For the typical function class considered, the approximation-error analysis assumes β-Hölder-smooth functions and intrinsic data dimension d.The convergence-rate argument then uses a choice of network parameter count based on n, d, and β.Source: §3 · p. 5
pFedBayes’s generalization-error convergence rate is minimax optimal up to a logarithmic term for bounded functions.The result follows by combining upper and lower bounds.Source: §3 · p. 6
Choosing ζ involves a trade-off: increasing it decreases the approximation error and upper bound but also decreases personalization.The paper recommends selecting ζ to balance the degree of personalization against the global upper bound.Source: §3 · p. 6

4. Algorithm

The algorithm represents network parameters with variational means and standard deviations, then updates personalized client models and a localized global model using stochastic gradients. The server aggregates updates from a random subset of available clients.

4. Algorithm

The algorithm reparameterizes each weight distribution using mean μ and standard deviation σ_m = log(1 + exp(ρ_m)), with ρ ensuring nonnegative standard deviations.It uses vectors μ and ρ in place of θ in the optimization problem.Source: §4 · p. 6
Clients estimate the non-KL part of their objective with Monte Carlo samples and use minibatch gradient descent, while the KL terms have closed-form results.The stochastic estimator uses minibatch size b and Monte Carlo sample size K.Source: §4 · p. 6
At each iteration, clients update personalized and localized global models; after R iterations, the server averages updates from a random subset of S available clients.The server uses an additional parameter β in its global-model update.Source: §4 · pp. 6–7

5. Experiments

Experiments compare pFedBayes with other methods on non-i.i.d. MNIST, FMNIST, and CIFAR-10 under varying data sizes. Personalized pFedBayes reports gains over SOTA comparison algorithms, with results varying by dataset and data size.

5.1. Experimental Setting

Each client has data from 5 of 10 labels; experiments use 10 clients for MNIST/FMNIST and 20 for CIFAR-10, with three training-set sizes per class.Accuracy is evaluated using the highest result over 700–800 communication rounds.Source: §5.1 · p. 7

5.2. Experimental Hyperparameter Settings

For pFedBayes, experiments set ρ = −2.5, ζ = 10, and personalized and global learning rates η1 = η2 = 0.001.The experiments tune learning rates and algorithm-specific parameters around recommended settings for comparison.Source: §5.2 · p. 7

5.3. Performance Comparison Results

On MNIST, pFedBayes is stable after 50 iterations and is reported clearly ahead of other algorithms on the small dataset.Figure 2 compares convergence rates across small, medium, and large MNIST datasets.Source: §5.3 · p. 7, Figure 2 · p. 8
Personalized pFedBayes beats SOTA by 1.25%, 1.78%, and 0.52% on MNIST; 0.42%, 0.63%, and 0.79% on FMNIST; and 11.71%, 7.19%, and 6.33% on CIFAR-10, across small, medium, and large datasets.The table covers results on all three datasets; the prose reports these gains in small-to-large order.Source: §5.3 · pp. 7–8, Table 1 · p. 8
For large datasets, Bayesian networks may need additional techniques to match non-Bayesian performance; using the same techniques as for small datasets degraded aggregation performance.The authors describe this as a fairness choice in their comparisons.Source: §5.3 · p. 8
pFedBayes also provides model-uncertainty estimates; Figure 3 illustrates prediction uncertainty for personalized models.The reported uncertainty figure uses a Gaussian assumption.Source: §5.3 · p. 8

6. Conclusions

The conclusion presents pFedBayes as a personalized federated-learning model based on variational Bayesian inference, with a client-level data/KL trade-off and a minimax-rate generalization guarantee. Experiments report better performance than many advanced personalization methods, especially with limited data.

pFedBayes updates each client's personal distribution by balancing private-data construction error against KL divergence from the aggregated global distribution.The paper also describes an alternating update algorithm for global and personal models.Source: §6 · p. 9
The averaged generalization error across clients is shown to achieve minimax-rate optimality up to a logarithmic factor.The conclusion says theoretical analysis provides an upper bound on this error.Source: §6 · p. 9
Experiments report better performance than many advanced personalization methods in most cases, especially when data are limited.Source: §6 · p. 9

A.1. Proof of Lemma 1

The supplied appendix passages outline proofs for Lemma 1, Lemma 2, and Lemma A.2, including a variational solution and derivative conditions.

A.1. Proof of Lemma 1

Lemma 1's proof invokes the Donsker–Varadhan representation of KL divergence and applies Jensen's inequality.The passages then refer to averaging over clients.Source: §A.1 · p. 12

A.2. Proof of Lemma 2

Lemma 2's proof introduces an optimal global variational solution for given client distributions and derives an upper bound under stated distributions and assumptions.Source: §A.2 · pp. 12–14

A.3. Proof of Lemma A.2

For Lemma A.2, the proof sets partial derivatives of the global objective with respect to its mean and scale parameters to zero.It then combines the resulting equations.Source: §A.3 · p. 15

B.1. Effect of Hyperparameters

The appendix reports hyperparameter choices for pFedBayes and comparison methods, and includes a figure illustrating personalized-model prediction uncertainty across clients and training rounds.

B.1. Effect of Hyperparameters

On the medium MNIST setting, the reported pFedBayes choices are η1 = η2 = 0.001, ζ = 10, and ρ = −2.5.The passage says ζ adjusts personalization: increasing it can improve global test accuracy while weakening personalized-model performance.Source: §B.1 · p. 15
The comparison setup uses recommended or tuned settings, including α = 0.5 for HeurFedAMP despite noting that it is not a fair comparison to α = 0.1.The text says higher α means less interaction with the global model.Source: §B.1 · pp. 15–16
Figure 3 presents prediction uncertainty for five clients at training rounds 0, 1, and 10 under Gaussian assumptions.Source: Figure 3 · p. 16

B.2. pFedBayes Uncertainty Estimation

pFedBayes estimates uncertainty in personalized-model predictions; estimates become more confident during training and can inform client aggregation.

B.2. pFedBayes Uncertainty Estimation

Before training, random model parameters do not provide accurate uncertainty estimates; confidence in classification results increases as training progresses.The paper describes this progression for the personalized model's uncertainty estimates.Source: §B.2 · pp. 16, 18
Client-model uncertainty can help assess model quality and inform whether to aggregate a client's model.Source: §B.2 · p. 18