Federated PAC-Bayesian Learning on Non-IID data

To address privacy concerns in distributed learning, federated learning (FL) has emerged as a viable solution, enabling multiple local clients to collaboratively train a model while retaining their private data and without sharing it [1, 2]. However, in real-world scenarios, data across different devices is not identically and independently distributed (non-IID), which poses challenges in model training and convergence [3].

Significant efforts have been made to improve performance and analyze convergence in non-IID FL [4], but few have provided theoretical guarantees by establishing generalization bounds. Most existing FL generalization analyses rely on the Probably Approximately Correct (PAC) Bayesian theory, first formulated by McAllester [5, 6]. Building on the McAllester’s bound, these analyses typically compute local bounds or apply existing PAC-Bayesian bounds directly, overlooking the non-IID nature of FL. This approach is flawed, as the PAC-Bayesian framework assumes that each data point is IID, ignoring non-IID data and directly employing the PAC-Bayesian theory, which potentially results in inaccurate or overly relaxed bounds. Consequently, techniques developed for the PAC-Bayesian framework are not directly applicable to non-IID FL. Therefore, this work aims to advance the theoretical underpinnings of non-IID FL.

Related works. The PAC-Bayesian framework has been extensively researched in recent years [7, 8, 9], yielding tighter and non-vacuous bounds. However, there has been limited exploration in the context of FL. Some studies have proposed information theoretic-based PAC-Bayesian bounds using Rate-distortion theory to prove generalization bounds [10, 11], providing an information-theoretic perspective on enhancing generalization capacity. Others have followed McAllester’s approach, attempting to directly apply the FL paradigm to the bound. For example, the authors in [12, 13] applied McAllester’s bound in a multi-step FL scenario; Omni-Fedge [14] used the PAC-Bayesian learning framework to construct a weighted sum objective function with a penalty, considering only a local client bound instead of the entire system, which precludes obtaining global information; and FedPAC [15] employed PAC learning to balance utility, privacy, and efficiency in FL. However, these approaches do not account for the non-IIDness of FL.

Our contributions. First, we derive a federated PAC-Bayesian learning bound for non-IID local data, providing a unified perspective on federated learning paradigms. To the best of our knowledge, this is the first non-vacuous bound for a model averaging FL framework. Specifically, due to the non-IID nature of clients, we assume that each client has unique prior knowledge rather than a common one. Additionally, the aggregation weights for non-IID clients vary instead of being uniform. Based on the derived bound, we define an objective function that can be computed by each local client rather than on the server and propose a Gibbs-based algorithm dubbed FedPB for its optimization. This algorithm not only preserves the privacy of each client but also enhances efficiency. Finally, we validate our proposed bounds and algorithm on two real-world datasets, demonstrating the effectiveness of our bounds and algorithm.

In this section, we introduce the federated PAC-Bayesian learning setting. The whole system comprises $K$ clients, each equipped with its own dataset $S_{i}=(x_{i},y_{i})_{i=1}^{n}\subseteq(\mathcal{X},\mathcal{Y})^{n}$ consisting of $n$ IID data points. Here $\mathcal{X}$ denotes the input space and $\mathcal{Y}$ denotes the output space. Each dataset $S_{i}$ is presumed to be drawn from an unknown data generating distribution $D_{k}^{\bigotimes n}$. Moreover, let $\ell:\mathcal{Z}\times\mathcal{W}\rightarrow\mathbb{R}^{+}$ be a given loss function and $h_{k}\in\mathcal{H}$ is a stochastic estimator on client $k$ where $\mathcal{H}$ is the hypothesis class. In the PAC-Bayesian framework, each client holds a tailored prior distribution $P_{k}$. The objective of each client is to furnish a posterior distribution $Q_{k}\in\mathcal{M}$, where $\mathcal{M}$ denotes the set of distributions over $\mathcal{H}$. We then define the population risk:

and the empirical risk:

by averaging over the posterior distribution of each client. In federated learning, each client will upload their posterior distributions to a central server, and then the server will aggregate the transmitted model in a weighted manner:

where $\bar{P}$ and $\bar{Q}$ are the global prior and posterior, respectively, and the averaging weight $p=(p(1),\dots,p(K))$ be a probability distribution on $\{1,\dots,K\}$. For the sake of generality, we can assume that $p(k)\in(0,1)$ and $\sum_{k=1}^{K}p(k)=1$. For intuition of this aggregation, we can see that minimizing the weighted objective function is actually equivalent to maximizing the logarithm of the corresponding posterior: $\min_{h}L(h)=\min_{h}\sum_{k=1}^{K}p(k)L_{k}(h)=\max_{h}\ln\prod_{n=1}^{N}\\ p\left(h\mid\mathcal{D}_{k}\right)^{p(k)}.$ In addition, we denote the Kullback-Leibler (KL) divergence as $D_{KL}(Q\|P)\triangleq\underset{Q}{\mathbb{E}}\left[\log\frac{dQ}{dP}\right]$ if $Q\ll P$ and $D_{KL}(Q\|P)=+\infty$ otherwise.

In this section, we will present our novel bounds on the non-IID FL scenario.

The RHS of Equation 3 comprises two components: the empirical term and the complexity term. Note that our bound eschews the typical smoothness and convexity assumptions on the loss often made by other FL frameworks. Moreover, an intuition can be presented for Equation 3 that the bound will be tighter with the increasing of clients scales, which is further corroborated by the evaluation in Section 5.3.

Equation 5 generally offers a strategy for selecting the value of parameter $\lambda$ so that the complexity term can be minimized.

We denote $\mathcal{L}_{k}=\underset{h_{k}\sim Q_{k}}{\mathbb{E}}\frac{1}{n}\sum_{i=1}^{n}\ell\left(h_{k}\left(x_{k,i}\right),y_{k,i}\right)$. Consider the following local objective function:

In our methodology, we introduce FedPB for a general scenario. It comprises two phases, designed to iteratively optimize the priors and posteriors for each client. Notably, in contrast to previous studies, clients are not required to upload their private prior and posterior distributions to the server, ensuring their privacy.

Phase 1 (Optimize the posterior). Given a fixed parameter $\lambda>0$ and the prior $P_{k}^{t}$ during the training epoch $t+1$, we aim to optimize the posterior $Q_{k}^{t}$ as $\hat{Q}_{k}^{t+1}=\arg\min_{Q_{k}}\mathcal{J}(Q_{k}),$ yielding the solution:

Phase 2 (Optimize the prior). Having derived the optimal posterior $\hat{Q}_{k}^{t+1}$, the prior can be updated by $\hat{P}_{k}^{t+1}={Q}_{k}^{t}$ since it minimizes $D_{KL}(Q_{k}\|P_{k})$.

Link with personalized federated learning. Since the prior $\hat{P}_{k}^{t+1}={Q}_{k}^{t}$ and ${Q}_{k}^{t}$ is equal to the aggregated global posterior at epoch $t-1$, the prior can be viewed as the global knowledge. Optimizing the objective function (6) minimizes the disparity between the global and local knowledge, which is a a prevalent personalization strategy in FL [17].

Re-parameterization trick. Utilizing the Bayesian neural network [18] as the local model aligns with our setting where all parameters are random, and we are optimizing their posterior distribution. In particular, the prior $P_{k}$ and posterior $Q_{k}$ are defined as follows:

with every model parameter $w_{i}$ being independent. Computing the Gibbs posterior directly can be challenging, hence we select the gradient descent as an alternative way. The update rule of (6) at round $t+1$ is:

where the parameter $\lambda$ can be regarded as the learning rate of the GD and the KL-divergence is the regularization term. Calculating the gradients $\nabla_{\mu_{Q_{k},i}}$ and $\nabla_{\mu_{Q_{k},i}}$ directly can be intricate, but the re-parameterization trick is capable of tackling this issue. Concretely, we translate the $h\sim Q_{k}$ to $\varepsilon\sim\mathcal{N}(0,I_{d})$ and then compute the deterministic function $h=\mu+\sigma\odot\varepsilon$, where $\odot$ signifies an element-wise multiplication. As a result, we have $\nabla_{\mu_{Q_{k},i}}\underset{h\sim Q}{\mathbb{E}}l(h)=\nabla_{\mu_{Q_{k},i}}\underset{\varepsilon\sim q(\varepsilon)}{\mathbb{E}}l(\epsilon)$, indicating its computability using an end-to-end framework with back-propagation.

In this section, we demonstrate our algorithm and our theoretical arguments with the FL non-IID setting. Specifically, the aggregation weight $p(k)$ is defined as the sample ratio of client $k$ relative to the entire data size across all clients. For the global aggregation, the global mean and covariance are calculated by $\bar{\mu}=\sum_{k=1}^{K}p(k)\mu_{k}\sigma_{k}^{-2}/\sum_{k=1}^{K}p(k)\sigma_{k}^{-2}$ and $\bar{\sigma}=1/\sum_{k=1}^{K}p(k)\sigma_{k}^{-2}$, respectively. Furthermore, we utilize two real-world datasets: MedMNIST (medical image analysis) [19] and CIFAR-10 [20]. For each dataset, we adopt three distinct data-generating approaches for local clients: 1) Balanced: each client holds an equal number of samples; 2) Unbalanced: varying sample counts per client (e.g., $[0.05,0.05,0.05,0.05,0.1,0.1,0.1,0.1,0.2,0.2]$ for 10 clients); 3) Dirichlet: differing sample counts per client following a Dirichlet distribution [21]. Besides, the entire FL system encompasses $K=10$ clients, initializing their posterior models uniformly from a global posterior model.

Additionally, we deploy two versions of Bayesian neural networks: one with 2 convolutional layers for MedMNIST and another with 3 layers for CIFAR-10. The CrossEntropy loss serves as our loss function, and it is optimized by the Adam optimizer [22] with a learning rate of 1e-3.