Improving Generalization in Federated Learning with Model-Data Mutual Information Regularization: A Posterior Inference Approach
Authors: Hao Zhang
Chenglin Li
Nuowen Kan
Ziyang Zheng
Wenrui Dai
Junni Zou
Hongkai Xiong
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated learning
posterior inference
model-data mutual information
Primary area: other
TL;DR: We propose FedMD for federated posterior inference, which incorporates a global model-data mutual information term as regularization to enhance generalization of the global posterior by alleviating bias of inferred local posteriors.

Abstract: Most of existing federated learning (FL) formulation is treated as a point-estimate of models, inherently prone to overfitting on scarce client-side data with overconfident decisions. Though Bayesian inference can alleviate this issue, a direct posterior inference at clients may result in biased local posterior estimates due to data heterogeneity, leading to a sub-optimal global posterior. From an information-theoretic perspective, we propose FedMDMI, a federated posterior inference framework based on model-data mutual information (MI). Specifically, a global model-data MI term is introduced as regularization to enforce the global model to learn essential information from the heterogeneous local data, alleviating the bias caused by data heterogeneity and hence enhancing generalization. To make this global MI tractable, we decompose it into local MI terms at the clients, converting the global objective with MI regularization into several locally optimizable objectives based on local data. For these local objectives, we further show that the optimal local posterior is a Gibbs posterior, which can be efficiently sampled with stochastic gradient Langevin dynamics methods. Finally, at the server, we approximate sampling from the global Gibbs posterior by simply averaging samples from the local posteriors. Theoretical analysis provides a generalization bound for FL w.r.t. the model-data MI, which, at different levels of regularization, represents a federated version of the bias-variance trade-off. Experimental results demonstrate a better generalization behavior with better calibrated uncertainty estimates of FedMDMI.

Improving Generalization in Federated Learning with Model-Data Mutual Information Regularization: A Posterior Inference Approach
Hao Zhang, Chenglin Li, Nuowen Kan, Ziyang Zheng, Wenrui Dai, Junni Zou, Hongkai Xiong School Electronic Information and Electrical Engineering Shanghai Jiao Tong University
Abstract
Most of existing federated learning (FI.) formulation is treated point-estimate of models, inherently prone to overfitting on scarce client-side data with overconfident decisions. Though Bayesian inference can alleviate this issue, direct posterior inference at clients may result biased local posterior estimates due to data hetero- geneity, leading sub-optimal global posterior. From information-theoreli perspective, we propose FedMDMI, a federated posterior inference framework based on model-data mutual information (MI). Specifically, global model-data MI term introduced as regularization to enforce the global model to learn essen- tial information from the heterogeneous local data, alleviating the bias caused by data heterogeneity and hence enhancing generalization. To make this global MI tractable, decompose it into local MI terms the clients, converting the global objective with MI regularization into several locally optimizable objectives based local data. For these local objectives, posterior is Gibbs posterior, which can be efficiently sampled with stochastic gradient Langevin dynamics methods. Finally, the server, we approximate sam- pling from the global Gibbs posterior by simply averaging samples from the local posteriors. Theoretical analysis provides generalization bound for FL w.r.t. the model-data MI, which, at different levels of regularization, represents federated version of the bias-variance trade-off. Experimental results demonstrate better generalization behavior with better calibrated uncertainty estimates of FedMDMI.
further show that the optimal local
Introduction
Federated learning (FL), as an emerging distributed learning framework, has garnered considerable attention [McMahan al. 2017]. In this paradigm, clients collaborate to train single global model, under the central coordination of a server. Notably, such collaborative training proceeds without the necessity of sharing exchanging the raw data of clients, thus providing a basic level of privacy protection. However, two primary challenges arise, specifically revolving around the intra-client data scarcity and inter-client data heterogeneity, since client models locally trained on these scarce and heterogeneous data are prone to overfitting and bias, leading to diminished generalization performance for the globally aggregated model at the server.
multitude of efforts have been made to address this issue from various perspectives, including federated optimization [Acar al. 2021 Karimireddy al. 2020). federated domain generaliza- tion [Nguyen et al. [0222] de Luca et al. 2022], federated knowledge distillation (Zhu et al., 2021 1 Afonin and Karimireddy 2021] J. etc. Despite their progress, majority of the existing FL formulations consistently treat it as a point-estimate (i.e., a single value estimated for each model weight) based
"Correspondence to: Chenglin Li <LCL. sjtu.edu.cn.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

on the loss function optimized on local client data. From probability theory perspective, however, using these single point-estimates as weights inherently poses a risk of overfitting on the scarce client-side training data. More crucially, this may also lead to overly confident decisions, due to the failure to provide reliable assessment of the uncertainty for models [Shridhar et al. 2019). which is indispensable in some safety-critical applications of federated learning, e.g., autonomous driving, healthcare, and finance.
Instead of single point-estimates, an alternative approach is to apply posterior inference over the model weights [Jospin et al., 2022]. which is theoretically attractive for preventing overfitting to scarce training data and providing natural way to assess uncertainty in weight estimates that can be further propagated into the model's prediction. Under the FL settings, however, little work has explored inferring a global posterior from the heterogeneous data across clients. To achieve the global posterior inference in FI., a typical framework comprises: i) each client initializing local model based on global posterior parameters and independently conducting local posterior inference; ii) the server receiving local posteriors and multiplicatively aggregating them based on global posterior decomposition. These two steps will then iterate until the global posterior converges. Following this framework, FedPA [Al-Shedivat et al. 2020] approximates these local posteriors via the Markov chain Monte Carlo (MCMC) method, and further reduces the computation cost. However, the data heterogeneity among clients may still lead to a biased local posterior inference, which in turn results in sub-optimal global posterior. To mitigate this bias. FedEP [Guo et al., 2023] approximates the global posterior using an expectation propagation method, which, however, incurs extra storage and communication overhead. As stateful method, FedEP is also not suitable solution for low- participation FL scenarios. Consequently, the issue of biased local posteriors incurred by the data heterogeneity in federated posterior inference remains unresolved.
In this paper, from an iiformation-theoreee in FL by incorporating a global mutual information (MI) regularization between the model and data (FedMDMI), which has been proved to be an effective measure of generalization capability of the learning algorithms [Xu and Raginsky 2017]. Under the FL settings, we show that the proposed global MI regularization can effectively alleviate the bias of local posterior incurred by data heterogeneity, and thus improve generalization capability of the global posterior. However, due to the non-exchange restriction of raw data in FL, it is impractical for the server to measure this global MI explicitly. We therefore turn decomposing the global MI into local MI terms, converting the global objective with MI regularization into several locally optimizable objectives based on the client data. For these local objectives, we show that the optimal local posterior is Gibbs posterior, a conclusion well-established in the field of PAC-Bayesian learning [Alquer [021]] Then, we employ the stochastic gradient Langevin dynamics (SGLD) method (Welling and Teh [011]] to sample from this local Gibbs posterior, which provides an unbiased and efficient sampling-bsse posterior approximation. Finally, at the server, we approximate the sampling from the global Gibbs posterior by simply taking the average of samples from the local posteriors. This aggregation method has been successfully employed and proved have a non-asymptotic convergence guarantee with the true global posterior in decentralized and federated scenarios [Gürbüzzbllaban et al., 2021 Plassier al. 2023 Through in-depth analysis, we provide a theoretical guarantee of our FedMDMI by establishing a generalization bound for FL w.r.t. the model-data MI. As a byproduct, differential privacy protection can also be brought by our FedMDMI. Our main contributions are as follows.
perspective, we propose to infer the global posterior
We introduce an cettmattion-hhooeetti approach for the global posterior inference in FL. It incorpo- rates global model-data MI term as regularization, which enhances generalization by alleviating bias of inferred local posteriors, and offers certain client-level privacy protection as byproduct.
We establish generalization bound for FL w.r.t. the global model-data MI, showing that reg- ularizing this global MI leads to reduetion in the generalization error. At different levels of regularization, it also represents a federated version of the bias-variance trade-off.
Extensive empirical results validate that FedMDMI outperforms the other point-estimate and Bayesian inference-basee baselines, while providing well-calibrated uncertainty estimates.
2 Related Work
Model-Data Mutual Information (MI. The mutual information (MI) between model and data quantifies the information that a model contains about the raw data, and serves as an effective measure of the model's complexity. Extensive research [Xu and Raginsky 2017 Asadi et al. 2018 has

---

revealed theoretical connection between this MI and the generalization ability of learning algorithms in centralized learning, Russo and Zou 2016 further demonstrate that this MI can effectively bound and reduce the bias in data analysis for centralized learning. For distributed and federated learning, prior works (Yagli et al., 2020 Barnes et al. 2022 Sefidgaran et al. [202]] also establish generalization bounds based on local MI at clients. Recent studies [Chor et al. 2023, Sefidgran et al., [023) further explore the relationship between the generalization bound and communication rounds through the local MI. In contrast to these methods, our generalization bounds are based on global MI at the server, resulting in tighter generalization bound. More significantly, we also provide a posterior inference method to estimate this MI.
Bayesian Inference. Compared to point-estimates of models, the posterior estimate is more suitable for FL with scarce local client data. Bayesian model ensemble learning- baaning-bbssdd methods, such as FedBE [Chen and Chao 2020] and FedPPD> (Bhatt et al. 2022]. construct a robust global posterior distribution via ensemble learning and knowledge distillation, which, however, require additional methods, on the other hand, decompose the global posterior into product of client posteriors. Among them, FedPA (Al-Shedivat) ct al. 2020] proposes a computation- and communication-efficin decomposition, where data heterogeneity may still result in biased local inference and sub-optimal global posterior. While FedEP [Guo et al. 2023] introduces expectation propagation at the client side to obtain sound global posterior at the server, this approach incurs additional storage and communication costs and may not be suitable for federated scenarios with low client participation. The work most closely related to ours is FALD Plassierr et al., 2023 which can be viewed special case of our FedMDMI when the hyperparameter a is set to 1. We will further demonstrate that this hyperparameter a is instrumental in controlling the tradeoff between fitting and generalization. Different from FALD, our FedMDMI further incorporates the global model as a more robust prior, and we derive a valuable generalization bound w.r.t. the model-data MI.
auxiliary datasets at the server. Global posterior decomposition-
framework for global posterior
3 Preliminaries and Problem Statement
The objective of standard FL is to learn a single global model w from clients via the following optimization problem:
Table I:
Summary of notations.
Si. S
random variables denoting local, global data random variable denoting learned global model number, index of communication rounds number, index local update step sample of client i'ss model at round / and step k sample of aggregated server model at round /
10
T,
min LD(w)= = weRd
T7 Σ (w), (1)
K, k
wi,k we
where (,) (w) Besxp(s) (wr, 2:)] represents the local expected risk of the i-th client associated with data distribution p(S)) on local data S. and the loss function (ww,, 2,) is usually chosen as a negative log likelihood, i.e., log (z, wz). We further define set of random variables S = {S Jm as the global dataset, consisting of m heterogeneous client datasets. Moreover, we let {z] ) denote the local training dataset of size n, drawn independently from distribution p(S.). This setup also gives rise to the global empirical risk as:
min (w) = m SEM
1
Ls, (wo),
(2)
HERd
where Ls, (w) 1 Ej= Ci(w, 2/)) is the local empirical risk.
n:
From a probability theory perspective, employing single point-estimate as the model weight, as shown in Eq. ( may render the model susceptible to overfitting, particularly when dealing with small and scarce datasets at clients. In addition, this point-estimate also makes the model incapable of correctly evaluating the uncertainty in the client's local training data, leading to overly confident decisions. Motivated by the Bayesian inference, we thus treat the global model as random variable rather than single point, and then estimate the global posterior of model w given the global dataset S. Subsequently, the objective of FI. converts to minimizing the posterior expected loss:
I
min Lss (w) Ep(cs()) P(w(S)
Lis (w)
EEM
where p(w[S)= = p( .I..., .., over p(w]S); we are then able to evaluate the risk of learned global posterior rather than a single value-estimate of model weights w.
Sm)) denotes the global posterior. By taking an additional expectation
3

---

However, is intractable to directly infer the global posterior, since raw data of the clients cannot be exchanged in FL. widely adopted solution is the global posterior decomposition [Neiswanger at 2014 AI-Shedivat et al. 2020), which decomposes the global posterior into a product of client's s
Th: (.(w) denotes the ratio of global prior to the product of client priors, which can be viewed as a constant based on prior assumptions, e.g., the Gaussian priors. See Appendix | for the detailed derivation of this decomposition.
local posteriors, i.e., p( SS) x T IlLEM P (w]S,), where _
p(w)
Based on the global posterior decomposition, the update process to optimize Eq. (3) can then be performed as follows: i) each client initializes the model based on global posterior received from the server and independently learns local posterior (e.g.. by variational inference or Markov Chain Monte Carlo); ii) the server collects the local posteriors from clients and multiplicatively aggregates them to obtaín the global posterior, which is then sent back to the clients for the next round of update. This iterative process continues until the global posterior converges, as depicted in Eq. (4)
Server aggregation: p(s" = II p(w"|S)) Client update p(w= (S.) arg argminn ptols, Ls. (4) p(w(S;)
CCM
This iterative update process is similar to FedAvg (McMahan et al., [017]. but has distinctive features:
i) each client estimates a distribution ((=]S))) instead of a single value of w, and ii) the global aggregation on the server is changed from weighted averaging to multiplication. FedPA (AI-Shedivat ct al. 2020] has enabled this process and further reduced computation and communication costs through federated least squares.
Though this posterior inference helps alleviate overfitting in scenarios with scarce client data as compared to single point-estimatt the heterogeneity of the data remains an issue. The independently learned local posteriors are susceptible to shifts induced by data heterogeneity, hindering their ability to generalize to data from other clients. Subsequently. the aggregation of biased local posteriors may lead to a sub-optimal global posterior [Guo et al., [023). This then raises a fundamental question: how can we achieve globally well-generalized posterior in federated heterogeneous scenarios.
Proposed FedMDMI
To reduce the model's dependency on heterogeneous data and mitigate biased local posteriors incurred by such heterogeneity, we are motivated to introduce an additional regularization information (MI) between the model and data (MD). In this section, we elaborate on our FedMDMI with an overview illustrated in Figure I and procedure summarized in Algorithm iin in Appendix A.1
the mutual
4.1 Compressing Information in Weights by Model- Data MI Constraint
The model-data mutual information is de- noted as T(w;; S), which quantifies relevance between the model and input data S, and also serves as a measure of complexity of the learned model for generalization analysis. By incorporating this cerrmaatio--thooetti con- straint as regularization, the original optimiza- tion formulation in Eq. (3) is converted to:
Gllbbal MI Regularication:
p(u)
Server (moving everage)
3)
Decompostiinn
ml
Local Regularization:
Σ (ays) + wii p(uo(S))
wax r S
Gibbs Fasterior
SGLD sampling
min Ep(u(s) $) p(w[S) s.t. 1(w; < les
1
Σ Lss (w) iEM
(5) Figure/ 1:
Overview of FedMDMI: under MI con- straint, each client uploads the sample from the local posterior to server and subsequently downloads the aggregated global sample from the global posterior.
8
Clicat 2
Cliai
Sg
m
where the introduced MI constraint com- presses information stored in model weights, and Ie denotes the upper bound of the global MI. Namely, if less information is extracted by the model from the global heterogeneous dataset, it is less likely that overfitting occurs. Furthermore, this MI constraint is also helpful in mitigating the bias incurred by data heterogeneity. If we denote the bias factor as 5 that affects the generation of local heterogeneous data (such as the difference
Herr, clients are not required to estimate an optimal local posterior before uploading, Similar to Fe- dAvg [MeMahan et al. 2017 the local posteriors are uploaded after certain number of updates.

---

in clients' locations, preferences, or habits), we then have the Markov chain + S - imply- ing [(w; 6) S [w; S). As consequence, we can diminish I [w; S) to constrain [ww; 6), thereby rendering the model insensitive to the bias factor å from the diverse clients.
In essence, this newly introduced model-data MI term (w;; S) constrains the complexity of model space for searching, and also encourages the model w to learn essential information from the global heterogeneous data S. Note that various studies in FL have also attempted to induce the model to learn invariant information across clients through another crucial MI constraint, which is the MI between input data and output representation. commonly known as the information bottleneck (Shwartz- Ziv and Tishby, [017]] regularization. However, accessing this MI constraint usually requires sharing additional feature representations of the data across clients, which is infeasible in FL (Zhang et al., 2023]. In contrast, our proposed model-data MI term not only requires no additional information to be shared with clients, but also provides client-level privacy protection as a byproduct, which will be discussed in detail in Section 5.3
4.2 Global Model-Data MI Decomposition
While our new formulation in Eq. (5) is conceptually promising, is still infeasible to directly measure the global model-data MI term [(w; S) in FL, where we are constrained to only leverage the distributed data at individual clients. Alternatively, we decompose it into several local MI terms I (w;; St), which can be determined locally by these clients.
Proposition 4.1. (Global Model-Data MI Decomposition). Suppose that S (Si,..., Sm) consists of data from the m clients and S A {Si,.... St-1),,,, then based on the chain rule of MI, we have:
m
m
I(w; S)
I (w; S:) I S:: (:::5t S
I (w; S:).
(6)
Proof. See Appendix A.4.3 for the detailed proof.
Therefore, instead [ww; S), we can constrain its upper bound Σm= I (wo; St) in Eq. introducing Lagrange multiplier 2 0, Eq. (5) then re-formulates to:
5
By further
min E Ep(w P(w|S)
Σ Ls,(w) al al((u;,s)) Si) MEM
m
where a is also viewed to balance the fitting and generalization. Based on the global posterior decomposition, the iterative update process to optimize Eq. (7) becomes:
Server aggregation: p(w" (S)
II p(w" 18));
(8)
Client update: p(w*|S]) = arg minEp(())))) [Ls, (w) + ad (ww S.)] p(w[S.)
(9)
At each round, client update needs not to fully infer optimal local posterior. Consistent with FedAvg, it only updates for certain number of iterations before sending learned local posterior to server.
4.3
Local Posterior Inference
For the client update in Eq. ( the second term (w; S,) can be expressed as:
T T(w; S:) (.s) KL [p(w]S), [p]]]]
(10)
which denotes the expectation of Dallback-Leibler (KL) divergence between the posterior [(w) S;) and marginal distribution p over the local data distribution p(S,). Notably, this KL divergence plays key role in the well-known PAC (Probably Approximately Correct)- -Baees bound FdI PAC-Bayes learning [Alquier 2021] provides tight generalization bound for learning algorithms, where the KI. divergence between the posterior p( w/ .S)) and prior p; (wrw)) is a dominant term within this bound. For a fixed prior, pioneer works [Xu and Raginsky [017]] Alquier al.) 2016) have found an optimal posterior to minimize this bound, which is often referred to as the Gibbs posterior. Namely, the optimal posterior follows typical Gibbs distribution.
'Specifically, Ps (10) A (ps [p(w] [ [S,)], which is also called the oracle prior. 'The PAC- ACyBaye bound still holds when pi(w) and the posterior p(w]S.) are chosen arbitrarily.
5

---

Lemma 4.2. [Xu and Raginsky 2017] The Gibbs posterior is the minimum of the objective for client
update in Eq. 2
1 B:
1 (Iss((w) - po
p [w]St) =
exp
(11)
where B, is a normalizationfacaaa
Proof. See Appendix A.4.4 for the detailed proof.
One possible way to obtain this Gibbs posterior is to seek its variational approximation (VA) [Alquer et al.) [016], which generally relies on simplified posterior distributions. However, primary drawback of VA is its tendency to yield biased posterior estimates for complex posterior distributions. What is worse, in FL the data heterogeneity across clients itself contributes already to biased local posteriors, and this bias may exacerbated by VA. Moreover, this method also results in at least doubling the communication overhead due to transmission of both the mean and covariance matrices. In contrast, Markov chain Monte Carlo (MCMC) methods offer an alternative class of sampling- based posterior approximations that are unbiased (Vadera et al. 2020]. , albeit with slower convergence. Thus, we adopt the stochastic gradient Langevin dynamics (SGLD) (Welling and Tch 2011]. an MCMC method that has been proven effective and scalable in large-scale posterior inference problem.
Specifically, SGLD draws samples from the Gibbs posterior, by using the stochastic gradient update:
wi,x = w - Vi. VU, (wi, + ht,kt
(12)
where wh, represents client i's model at round and step k, 7L 1.k
is the local step size, h4 is noise
variable sampled from NO, v 2 t.l al) with being the identity matrix, and
VU. ( v (wl,,,))) - log (w) lu=
(13)
unbiased estimate of gradient. Note that U, corresponds to the exponent in the Gibbs posterior of Eq. 11 and pt (w) h= represents the prior puw)'s probability density value at w = w
4.4 Global Posterior Aggregation and Prior Selection
In theory, if the step size is annealed as Mi + the client update sequence w converges the local Gibbs posterior in Eq. (TT) with sufficiently large and The question then becomes: how we obtain samples from the global posterior (i.e. the product of local Gibbs posteriors) expressed as
p (w S)
II p(w]S) = B
1
Σ
[1s - loggi(())
(14)
iEM
based on the samples drawn from these local posteriors.
Recent studies Güübiizbalaban al. 2021 Plassier al., 2023 have shown that when clients utilize SGLD for posterior inference in distributed or federated settings, leveraging the mean of samples from local posteriors to approximate samples from the target global posterior p( p(w|S), there is non-asymptotic convergence guaranted In other words, we can approximate samples from the global posterior by simply taking an average of the samples drawn from the local posteriors. Specifically, in the local posterior inference, SGLD introduces uncertainty into the predictive estimates by incorporating Gaussian noises, and samples the local model wi, from the local posterior p(w) SS)) in Eq. [11) through Markov chain with steps:
Awi, N (- (wiA--)) 22 a0) o1)
(15)
where Awi, A w/,x wiA After the client performs K steps and uploads the model change, the server can approximate sample w from the global posterior p(w]S) in Eq. (14) simply by taking an averaging (i.e., Aw, & we+1
wt m ΣEEM Awy), through global Markov chain with steps:
Au N
Σ
Σ
n
VU. (wixe-)) m
1
2 al
(16)
HEM k=
'While this convergence bound O((1 - (1//))/ KT +1/m) is established only for p-strongly convex functions,
empirical evaluations in Plassier al. 2023 and results show that it still holds in non-eo settings.

---

We then discuss how the prior is chosen in Eq. renders the mutual information the tightest. which, however, infeasible obtain. In fact, we can use an arbitrary prior r(w) (e (0, I)) to approximate Pi(w),, based on following upper bound:
13 This oracle prior (i.e., P:(w) - Epts, ())
[ww; S:) = Ep(s)) KL (p(u) (1))] KL (pi(w))lr(w) < Ep(s.) [LL (p(w))))]lr())
(17)
For the arbitrary prior r(w) achieve smaller MI, it must essentially predict the posterior [Dziugaite] ct al. [021]. Thus, we consider using the information of global model We as the prior of clients for up-
date of the next round. Based on Eq. 16 for updating Wr, we then define pi (w) 2 (Wwlue-, 1 where
K
1 K
Pr Wt-
1-1. VU (wi-1,-11, Et-1 41
1- 21
al.
(18)
iEM k=1
k=1
Here the variance captures the uncertainty introduced by all the clients the previous round i 1. Meanwhile, Zhang et al. [2022) also show that using the global model as a local prior can alleviate local overfitting. Additionally, w, only sample drawn from the global posterior. To better estimate the mean of this Gaussian prior, we approximate with a global moving average, which is also utilized to accelerate the convergence as shown in Line 16 of Algorithm in Appendix A.1
Fe Analysis
5.1 Generalization Analysis via Model-Data MI
We first provide an information- FL. In FL settings, it is crucial to consider both the gaps arising from the unseen client data (i.e., participating error), and the gaps stemming from the unseen client distributions (i.e., participation gap). Following the framework proposed by Yuan et al. 2021 and Hu et al. [2023), we re-define the more general population risk in Fl. as:
generalization bound in terms of the model-data MI for
Lp(w)= = EM(S)P [E [[[ +]]]
(19)
where we denote P as meta-distribution D, and D is the set of all distributions p( (S:). This formulation takes into account the participation gap, as compared to Eq. @.. By recalling the global empirical risk (w)) in Eq. (3) we define the expected generalization error E [Lp(w) Ls Ls(w))
Theorem 5.1. (Generalization Bounds for FL). Suppose that l(w,, 25) for all e M is bounded by and independent, then the expected generalization error satisfies:
C*I(w;)) S) +
C²1(w;;)
E [[p(w)) - Ls(w)]
[[([)) S
(20)
2mm
2m
where m is the number of clients, the number of samples on the client (assuming, without loss of generality, that the number of samples for all clients is equal), and D is the set of distributions p(S.).
Proof. See Appendix A.4.5 for the detailed proof.
Remark 5.2. On the RHS of Eq. 20 [ww; S) relates to participating generalization error and serves
as the regularized mutual information term that can be estimated in our FedMDMI, while [(w; D) relates to the participation gap that cannot be estimated due to unavailability of the non-participating
clients. Moreover, this generalization bound O
matches with the current bound HHu
m
et al. 2023
Remark 5.3. Several prior works (Yagli al. 2020 Barnes et al.
2022
Setidgaran et al. 2022
have also established a generalization bound based on clients' local MI
( 5 in distributed (or
federated) learning. Additionally, some studies [Chor et al. 2023 Sefidgaran et al. 2023 in FL have explored generalization upper bounds w.r.t. the local time-varying mutual information (w,, S!). In contrast to them, our generalization bound is based on a participation gap framework, with the first term (global MI (w; 5)) exhibiting a tighter bound than the previous local MI Σ, (w; St). For more details, we provide more elaboration of this distinction in Appendix A.4.2
We have also incorporated this technique into the other baselines in our experiments for fair comparison.
7

---

5.2 Fitting vs. Generalization with Data Heterogeneity
As demonstrated in Theorem 5.1 the MI (w; S), acting as a measure of the effective complexity of a model, controls participating generalization capability. The hyperparameter is employed in Eq. 7 to regulate this MI, thereby balancing fitting and generalization of the learned model. Here, we empirically investigate this trade-off under two non-iid settings using the CIFAR-10 10 dataset. The experiments involve 100 clients with a 10% participation rate (refer to Section 6 for under varying
details). The train and test errors of the global model are plotted in Figure 2 It can be seen that with the increase of a, the train error increases slowly, while the test er- ror decreases initially and then increases. This
Dir (0.2)
T 19.5
Dir (0.2)
Dir (0.7)
13.0
Dir (0.7)
train enoor
8 18.5
suggests that models with lower complexity (larger ()) tend to underfit, while those with higher complexity (smaller a) tend fit. The MI I (w, S) allows us to recover federated version of the bias-variance trade- off. Moreover, as a increases, the gap between the test errors under different data heterogene- ity decreases, indicating that the introduced [w; S) can render w insensitive to data hetero- geneity and alleviate the bias caused by such heterogeneity.
0
17.5
17.0
165
-30 -9
-22 -11 -30
atlog:s)
of a(0910)
(a) train error
(b) test error
Figure 2: Train and test errors of the global model on CIFAR-10 dataset, where Dir (0.2) has a higher degree of data heterogeneity than Dir (0.7).
5.3 Client-Level Privacy Protection
The proposed FedMDMI can also provide client-level privacy protection, according to the analysis on MI and posterior inference based on SGLD sampling. First, the proposed MI regularizer [ww; S) directly quantifies the extent to which the model memorizes data. Conversely, reflects the degree of data information leakage. In the extreme case when [(w; S) = 0, becomes entirely independent of S (c.g., akin to white noise) and leaks no information about the data. Thus, we restrict this MI to some extent for privacy protection. Second, SGLD in the local posterior inference has been widely demonstrated to provide strict differential privacy [Wang et al. 2015 Dziugaite and Roy
2018]. Following these works and assuming that (, z!) is L-smooth, when O
V Tog our
FedMDMII also preserves a client-level (c, 5) -ddiifferettial privacy. See Appendix A.4.6 for details.
5.4 Limitation and Complexity Analysis
Though FedMDMI enjoys multi-fold benefits, including mitigating bias and overfitting on hetero- geneous data to enhance generalization, and differential privacy as a byproduct, it may also have limitations. First, both the global posterior decomposition and global MI decomposition rely on the assumption that the global likelihood is conditionally independent given w. While this assumption is commonly adopted in embarrassingly parallel [Neiswanger et al.) 20144 and federated scenarios [Guo et al,, 2023) Al-Shedivat et al., 2020]. it may not be applicable in certain extreme FL scenarios. Second, we do not delve into the convergence rate of FedMDMI under non-convex settings. While numerous studies [Gürbüzbalaban et al.) 202T Plassier et al.) 2023) have analyzed the convergence of SGLD for strongly convex objective in distributed and federated posterior inference, and empirically demonstrated its effectiveness for non-convex objective, the theoretical convergence of SGLD remains an open question when applied to the non-convex objectives in FL.
Last, regarding the complexity analysis, we begin by defining the dimensions of the neural network as d. At each communication round, we adopt the SGLD to estimate the local posterior. Specifically, compared to the SGD employed in FedAvg and FedBE, our FedMDMI entails an additional step of generating the Gaussian noise with d' dimensions, which is subsequently incorporated into each model update iteration. This results in an additional (dd time and (dd) memory. In contrast, at each client, FedPA uses dynamic programming to approximate the inverse matrix d of the neural network, introducing an additional ((I2d) d) time and (ld) memory, where 1 is the number of posterior samples. Similarly, FedEP(L) also requires approximating the covariance as the inverse Hessian, introducing an additional (dd) time and (dd²) memory. For communication and aggregation at the server, our FedMDMI, along with FedAvg and FedPA, requires O(md) time and O (md) memory, where m denotes the number of clients. In contrast, FedEP requires O(md) time and (mdd) memory. Note
8

---

Table 2: Performance comparison (with 5 random seeds) under various settings, where a smaller Dirichlet parameter indicates higher data heterogeneity, and L and H indicate low and high participation rates, respectively. Bold numbers indicate the best performance.
Top Test Accuracy
Duraset
Setting
FodAng
FedM
MimeLite
SCAFFOLD
FedB
PeaPA
FediiP
RALD
PedMDMI (Ouns)
(0.244
79. 777
892 08.12:0.45
82.31
81.49-0.57
80. 70.
82.68
825660063
83.4260330
CIFAR-
60.01++1.00
69.46-0.00
92:25++0044
82.05-0.05
63.18+0.54
8300-+0.67
R3,1 +0
10.2)-4
79.08
94.17.0.62
N2 022.66690.00
81.25
84.78
80.0120777
N2845092
K3.5650044
(0.7)-1
62.00-1.00
65.04+00.62
82.53-0.35
60.05-0.70
83,70
1311241.33
63.76+0061
00.20-
47.15=0010
47. 00:
46.55-00.20
18.50:0.00
03:80:11
13-0000
49.46 10.55
CIFAR
(0.70
41. 11.29.00.00
07.89
41.99.10.00
45.76L10.47
40.45.00.07
49.55.0.85
48.6610.000
50.45 10.55
47 0:202:0:02
00882-0000
67 10:10:91
46,82:0000
18.51:0.87
49.08 67
23-0077
49.70
(0.7):11
02.35
02.07
47.5901/0.622
46.291:0.45
40.08.200447
5000230.40
49.56110.83
50.71
Shake-
non-3d-L
46.11:210:01
50
06.00:00:77
09. 00.7888:0:33
50.11:20.08
20.81
51.08-00.78
52.93 20.00
speare
mon-dd-H
40.99+00.0
51.10-0:52
47.19-00778
68.00++0.70
81.26-0064
52.06
51.12-0.57
53.37
0.5
80
Tace
0.4
0.2
FIL
0.0
67 OR
16
D1 ne 07 Confeonce
1003 2033
3002
4053
10:00
2300
3032
4000
(a) CIFARIO
(b) CIFARI00
(c) CIFARIO
(d) CIFARI00
Figure 3: (a) and (b): Reliability diagrams for uncertainty estimates, where confidence the value of the maximum softmax output. A perfectly calibrated model shows no difference between accuracy and confidence. which is represented by a dashed brown line. Points below this line indicate under- confident predictions, while points above this line correspond to overconfident predictions. (c) and (d): Test accuracy vs. number of communication rounds with Dir (0.2) and participation rate 5%, where curves are averaged over 5 random seeds.
that FedBE not only requires O((m + + 1)d) time and O((m + 1 + 1)d) memory, where denotes the number of global model samples, but also performs the knowledge distillation at the server using unlabeled data, which is both memory- and time-intensive.
Experiments
Datasets and Models. We evaluate our FedMDMI on three benchmark datasets: CIFAR-10, CIFAR- 100, and Shakespeare, with heterogeneous data splits Specifically, for CIFAR-10 and CIFAR- 1000 each client has an uncertain number of classes through setting client sample labels according to the Dirichlet distribution. For example, with Dirichlet (0.2) each client has 80% samples which belong to mostly three or four different classes, while with Dirichlet (0.7) each client has 80% samples which belong to mostly five or six different classes. We use the CNN and RNN models similar to the
prior works [McMahan et al. 2017 Acar al. 202 Detailed experiment setup can be found in Appendix A.2.I
Comparison Methods. Methods for comparison include the single point-estimate-bssee approaches: FedAvg (McMahan et al., 2017). FedM (FedAvg with momentum moving average) [Hsu et al. 2019). ) MimeLite [Karimireddy et al. 2021] and SCAFFOLD (Krimireddy et al.] [2020]. as well as the Bayesian inference-basee
approaches: FedBE [Chen and Chao 2220]] FedPA [Al-Shedivat et al. 2020), FedEP (1) [uuo et al. 2023 and FALD Pllsssier et al.) [023]] Note that Guo et al. 2023 and Plassier et al. 20233 also propose various variants of their methods. In our evaluation, we opt for FedEP (1) and FALD, as they demonstrate superior overall performance. We present additional results of the other bascline (i.c., 3-PredBayes [Hasan et al. 2024)) in Appendix AA3555 In addition, for a fair comparison, we also incorporate the momentum moving average in the global model aggregation step of FedM, FedPA, FedEP, and FALD.
Implementation. We evaluate the performance of the global model after 4000 communication rounds for CIFAR-10 10 and CIFAR- 100, and after 1000 rounds for Shakespeare, which is trained with 100 clients. We set the high (H) and low (L) client participation rates as 10% and 5%, respectively. Clients are uniformly sampled at random without replacement at each round. The learning rate and hyperparameters for all approaches are individually tuned over a grid search. See Appendix A.2.3 for the additionally detailed settings of hyperparameters. Implementable code for evaluation of our FedMDMI is available at: https: :// //gtttttt com/haozzhh/eedHHH1

---

Experiment Setting regarding Uncertainty Quantification for Image Tasks. Uncertainty estima- tion crucial for decision-ne decision-making, Following Guo al.) 2017]. Maddox et al. [2019), we utilize the expected calibration error (ECE) as calibration metric of predictive uncertainty. To calculate ECE for a given model, we divide the test samples into 20 bins based on the model's confidence. Subsequently, we compute the absolute difference between the average confidence and accuracy within each bin and average these differences across all the bins. Additionally, reliability diagrams plotted to illustrate the discrepancy between a method's ' confidence in its predictions and its actual accuracy. For a well-calibrated model, the discrepancy should be close to zero in each bin.
Performance Evaluation. The experi- mental results of all comparison methods under different non-iid settings are pre- sented in Table 2 and Figures c)) and Bjd). In most cases, FedMDMI outper- forms the other algorithms on the three datasets with varying heterogencous data and client participation rates. This supe- rior performance can be attributed to the FaIMDMI
Table 3: ECE (averaged over 5 random seeds).
Method
CIFAR-10 10
CIFAR 100
Dir (0.2)_
Dir (0.75L.
Du (0.23L
Di (0.7)-L
FedAug
0.169 0.0030
0.165
0.0022
0.420
00.030
0.452 0.0.008
FeeM
0.169 a: 0u0r2
0.369 2: 0.0036
0,468
-000027
0.450
4: 11.1000
MimeLite
0.182 4: 000041
0.178 1: 000034
0,461
-0.0229
0.470 1 1..1222
SCAFFOLD
0.192 1: 0.0018
0.194 0.0025
0.472
0..011
0.479 1 1.0.004
RelBE
0.182 -1002)
0.189 + 0.0032
0.640
0.0015
0.463
+ +0.002
FalPA
0.178
0.376 = 0.0020
0.884
= 0.0035
0.871 # #.0.1225
FeJFI
0.121 +-000053
0.118 0.002)
0.250
0.0045
0.273 + +.0.027
FAILD
0.135+00.002
0.127 0.0025
0.267
0.0015
0.260
+ 0.0023
0.015+00.01
0.120 + 0.0051
0.261 - 0.0025
0.263 0.002
proposed model-data MI regularization, which encourages the distributed clients to learn essential information and alleviates bias in the inferred local posteriors, thereby enhancing the generalization capability of the global model. It is worth noting that the convergence rate of the proposed FedMDMI may not be the fastest, especially when compared to the optimization-biz -based method, SCAFFOLD. One potential explanation is because of our use of stochastic gradient Langevin dynamics (SGLD) to approximate the posterior, which often suffers from slow convergence rate due the variance introduced by the stochastic gradient [Wang et al. 2021a). However, compared to SCAFFOLD, our method not only achieves a superior generalization performance, but also offers an improved calibration of uncertainty. Due to space limit, we present some additional experimental results and additional discussion Appendix A..3
Calibration and Uncertainty Estimation. We evaluate the uncertainty estimates of all comparison methods, shown in Table and Figures (a) and In most cases, our FedMDMI attains the lowest expected calibration error (ECE) and provides well-calibrated uncertainty estimates. Notably, the ECE of single point-estimation-basee methods is consistently higher than that of posterior inference- based methods. While all these methods exhibit a tendency to be overconfident in their predictions, posterior inference demonstrates the potential to mitigate this overconfidence.
Choice of appropriate a and B. We have analyzed the impact of a that balances fitting and generalization in Section 5.2 We then analyze the hyperparameter 3, which is used to estimate the mean of Gaussian prior and accelerate convergence. Due to space limit, we present these results in Appendix A.3.2 showing that even if [ 0, FedMDMI still achieves higher accuracy than FedAvg.
7
Conclusion
We have proposed a federated posterior inference approach, which mitigated bias in the posterior estimates and improved generalization by introducing a global model-data MI regularization. To approximate this global MI based on distributed data over the clients, we decomposed it into local MI terms. We showed that the optimal posterior of the local objective with MI regularization was a Gibbs posterior, which could be efficiently sampled by SGLD. We further provided a generalization analysis based on this global MI, and analyzed its impact on fitting and generalization in FL. enabling to present federated version of the bias-variance trade-off.
8 Acknowledgement
This work was supported in part by the National Natural Science Foundation of China under Grant 62125109, Grant T2122024, Grant 62320106003, Grant 62371 1288, Grant 62431017, Grant 62401357, Grant 62401366, Grant 61931023, Grant 61932022, Grant 62120106007, and in part by the Program of Shanghai Science and Technology Innovation Project under Grant 24BC3200800.
10

---

References
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiiin learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pages 1273-1282. PMLR, 2017.
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough. and Venkatesh Saligrama. Federated learning based dynamic regularization. arXiv preprint arXiv:2111.04263, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pages 5132-5143. PMLR, 2020.
Tuan Nguyen, Philip Torr, and Ser-Nam Lim. Fedsr: A simple and effective domain generalization method for federated learning. In Advances in Neural Information Processing Systems, 2022.
Artur Back de Luca, Guojun Zhang, Xi Chen, and Yaoliang Yu. Mitigating data heterogeneity in federated learning with data augmentation. arXiv preprint arXiv:2206.0999 2022.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In International Conference Machine Learning, pages 12878-12889) PMLR, 2021.
Andrei Afonin and Sai Praneeth Karimireddy, Towards model agnostic federated learning using knowledge distillation. arXiv preprint arXiv:21 152/0, 2021.
Kumar Shridhar, Felix Laumann, and Marcus Liwicki. A comprehensive guide Bayesian convolu- tional neural network with variational inference. arXiv preprint arXiv: arXiv:1901.080 1901.0273, 2019.
Laurent Valentin Jospin, Hamid Laga, Farid Boussaid, Wray Buntine, and Mohammed Bennamoun. Hands-on Bayesian neural networks tutorial for deep learning users. IEEE Computational Intelligence Magazine, 17(2):29-48, 2022.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Representations, 2020.
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning as variational inference: scalable expectation propagation approach. In The Eleventh International Conference on Learning Representations, 2023.
Aolin Xu and Maxim Raginsky. Information- -hhooretcc analysis of generalization capability of learning algorithms. Advances in Neural Information Processing Systems, 30, 2017.
Pierre Alquier. User-friendly introduction PAC-Bayes bounds. arXiv preprint arXiv:2110.11216 2021.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688. 2011.
Mert Gürbüzbalaban, Xuefeng Gao, Yuanhan Hu, and Lingjiong Zhu. Decentralized stochastic gradient Langevin dynamics and Hamiltonian Monte Carlo. The Journal of Machine Learning Research, 10844-000011777 2021.
Vincent Plassier, Eric Moulines, and Alain Durmus. Federated averaging Langevin dynamics: Toward unified theory and new algorithms. In International Conference on Artificial Intelligence and Statistics, pages -3366 PMLR, 2023.
Amir Asadi, Emmanuel Abbe, and Sergio Verdú. Chaining mutual information and tightening generalization bounds. Advances in Neural Information Processing Systems, 31, 2018.
Daniel Russo and James Zou. Controlling bias in adaptive data analysis using information theory. In Artificial Intelligence and Statistics, pages 1232-1244 PMLR. 2016.

---

Semih Yagli, Alex Dytso, and H Vincent Poor. Information -thooretic bounds on the generalization eπror and privacy leakage in federated learning. In 2020 IEEE 21st International Workshop on Signal Processing Advances in Wireless Communications (SPAWC), pages 1-5. IEEE, 2020.
Leighton Pate Barnes, Alex Dytso, and Harold Vincent Poor. Improved iiformation-thhoreee generalization bounds for distributed, federated, and iterative learning. Entropy, 24(9):1178 2022.
Milad Sefidgaran, Romain Chor, and Abdellatif Zaidi. Rate-distortion theoretic bounds on gener- alization error for distributed learning. Advances in Neural Information Processing Systems, 35: 19687-19702, 2022.
Romain Chor, Milad Sefidgaran, and Abdellatif Zaidi. More communication does not result in smaller generalization error in federated learning. In 2023 IEEE International Symposium on Information Theoor (ISIT), pages 48 53. IEEE, 2023.
Milad Sefidgaran, Romain Chor, Abdellatif Zaidi, and Yijun Wan. Federated learning you may communicate less often! arXiv preprint arXiv:2306.05866 2023.
Hong- You Chen and Wei-Lun Chao. Fedbe: Making Bayesian model ensemble applicable to federated learning. arXiv preprint arXiv:2009.0199 2020.
Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Bayesian federated learning via predictive distribution distillation. arXiv preprint arXiv:2206.0777 076662 2022.
Willie Neiswanger, Chong Wang, and Eric P Xing. Asymptotically exact, embarrassingly parallel MCMC. Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, pages 623-632, 2014.
Ravid Shwartz-Zii Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arviv: / 703.008100 2017.
Hao Zhang, Chenglin Li, Wenrui Dai, Junni Zou, and Hongkai Xiong. FedCR: Personalized federated learning based on across-client common representation with conditional mutual information regularization. In International Conference on Machine Learning, pages 41314-41330, PMLR, 2023.
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of Gibbs posteriors. The Journal of Machine Learning Research, 17(1):8374-844 2016.
Meet Vadera, Brian Jalaian, and Benjamin Marlin. Generalized Bayesianyesian posterior expectation distillation for deep neural networks. In Conference on Uncertainty in Artificial Intelligence, pages 719-728. PMLR, 2020.
Gintare Karolina Dziugaite, Kyle Hsu, Waseem Gharbieh, Gabriel Arpino, and Daniel Roy. On the role of data in PAC-Bayes bounds. In International Conference on Artificial Intelligence and Statistics, pages 604-612. PMLR, 2021.
Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized federated learning via variational Bayesian inference. In International Conference on Machine Learning. pages 26293 2713002610. PMLR, 2022.
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? arXiv preprint arXiv:21100 14216, 2021.
Xiaolin Hu, Shaojie Li, and Yong Liu. Generalization bounds for federated learning: Fast rates, unparticipating clients and unbounded losses. In The Eleventh International Conference on Learning Representations, 2023.
Yu-Xiang Wang, Stephen Fienberg, and Alex Smola. Privacy for free: Posterior sampling and stochastic gradient Monte Carlo. In International Conference on Machine Learning, pages 2493- 2502. PMLR, 2015.
Gintare Karolina Dziugaite and Daniel Roy. Data-dependent PAC-Bayes priors via differential privacy. Advances in neural information processing systems, 31, 2018.
12

---

Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non- non-identical data distribution for federated visual classification, 2019.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Breaking the centralized barrier for cross-device federated learning. Advances in Neural Information Processing Systems, 34:28663- 28676, 2021.
Mohsin Hasan, Guojun Zhang. Kaiyang Guo, Xi Chen, and Pascal Poupart. Calibrated one round federated learning with bayesian inference in the predictive space. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pages 12313-12321, 2024.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pages 1321-1330. PMLR, 2017.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson A simple baseline for Bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
Bao Wang, Difan Zou, Quanquan Gu, and Stanley J Osher. Laplacian smoothing stochastic gradient Markov chain Monte Carlo. SIAM Journal on Scientific Computing. (3(1):A26-A53, 202
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
Longbing Cao, Hui Chen, Xuhui Fan, Joao Gama, Yew-Soon Ong, and Vipin Kumar. Bayesian federated learning: A survey, 2023.
Minyoung Kim and Timothy Hospedales. Fedhb: Hierarchical bayesian federated learning. arXiv preprint arXiv:2305.00 04979, 2023.
Hui Chen, Hengyu Liu, Zhangkai Wu, Xuhui Fan, and Longbing Cao. Fedsi: Federated subnetwork inference for efficient uncertainty quantification. arXiv preprint arXiv:2404.15657, 2024.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared represen- tations for personalized federated learning. In International Conference on Machine Learning, pages 2089-2099 PMLR, M,,,,22 2021.
Gary Cheng, Karan Chadha, and John Duchi. Federated asymptotics: a model compare federated learning algorithms. In International Conference on Artificial Intelligence and Statistics, pages 10650-10689. PMLR, 2023.
Zifeng Wang, Shao-Lun Huang, Ercan E Kuruoglu, Jimeng Sun, Xi Chen, and Yefeng Zheng PAC-Bayes information bottleneck. arXiv preprint arXiv:21099 14509, 2021b.
Cynthia Dwork, Aaron Roth, al. The algorithmic foundations of differential privacy. Foundations and Trends in Theoretical Computer Science, 933-4):21-409 2014.
13

---

A Appendix
A.1 Detailed Algorithm Design
Here, we present a detailed learning procedure for our proposed FedMDMI, as outlined in Algo- r
Algorithm 1: Proposed FedMDMI
Server Initialization: Avo = O;
2 for each round 1.2....T do sample clients P: [ M
for each client E P; in parallel do receive and initialize local model w Wt
set p(w) N
K Whi m =k=1
20 1 al
for each local step ho = 1,2, ... K do
Us (wi = Ls, (wi
log p( w) --
1e N 0, 2011 cik al
wi,s wiik-1
t.k VUs (wi + h
end
Auf w who and send Au? to server
12
end
14
// at server:
Arc. 1 DiEp, Awf
15
Au --33 (3Ar)-1 + B))w)
17
Wi+11 = We + nnu
broadcast W1+1
end
A.2
Detailed Experiment Setup
The models, data splitting methods, and hyper-parameter settings employed in this paper closely adhere to those specified in the previous empirical benchmarks, which will be detaied the following.
A.2.1 Datasets
We utilize the visual datasets, including CIFAR-10 and CIFAR-100, and the language dataset Shakespeare. These datasets comprise 10, 100, and 80 different labels, respectively. The train and test splits for CIFAR-10, CIFAR-100, and Shakespeare are detailed in Table 4
Table 4: Train and test splits.
Dataset
No. Train
No. Test
No. Train per client (100 clients)
Batch size
Rounds
CIFAR-10
50000
10000
500
100
4000
CIFAR- 100
50000
10000
500
100
4000
Shakespeare
200000
40000
2000
100
1000
To generate heterogeneous splits for CIFAR- 10 and 100, we adopt method where the training
samples are divided by classes and assigned to clients. This approach aligns with the methodology used in FedDyn (Acar et al. 2021 Specifically, we employ a Dirichlet distribution over the labels of the dataset to create federated dataset. this process, vector of the same size as the number of classes is generated for each client using the Dirichlet distribution. These vectors represent the class
14

---

priority of each client. Subsequently, a label is sampled for each client based on these vectors, and images are sampled according to the selected label without replacement. This process repeated until all the data are assigned to clients. It is important to highlight that the factor of Dirichlet distribution corresponds to the degree of data heterogeneity. For Dirichlet factor of 0.2, each client has 80% of samples predominantly belonging to three four different classes. And for Dirichlet factor of 0.7, each client has 80% of samples mainly belonging to five or six different classes.
For Shakespeare, we employ LEAF to generate heterogeneous datasets by limiting the number of data points per client 2000, methodology akin to FedDyn [Acar et al., [202T]. the heterogeneous splits, each client is associated with speaking role comprising few lines.
A.2.2 Models
For CIFAR-10 10 and CIFAR-100 100, we employ CNN model comprising two convolutional layers with sixty four 5 filters, two 2 x 2 max pooling layers, two fully connected layers with 384 and 192 neurons, and softmax layer. A comprehensive description of the model is available in Table 5 Our
CNN model closely resembles those utilized in FedAvg [McMahan et al. 2017 and FedDyn Acar
et al.
202T
For Shakespeare, utilize RNN that takes sequence of characters input, embedding each into an 8-dimensional learned space. This sequence is then processed through two-layer LSTM with hidden size of 100 nodes, followed by a Softmax layer. A comprehensive description of the model is presented in Table 6 This neural network model is identical to the one utilized in FedProx
al. 2020 and FedDyn (Acar al.
2021
Table 5: CNN Architecture for CIFAR-IO & CIFAR- 100.
Table 6: Shakespeare model ar- chitecture.
Layer Type
Size
Convolution ReLu
5x5x64
Layer Type
Size
Max Pooling
2x2
Embedding
(80, 8)
Convolution ReLu
5x5x64
LSTM
(80, 100)
Max Pooling
2x2
LSTM
(80, 100)
Fully Connected + ReLU
x600x384
Fully Connected
(100, 80)
Fully Connected + ReLU
384 x 192
Fully Connected
192x & 192x 100
A.2.3 Hyper-Parameterr
All different comparison approaches are implemented in PyTorch 1.4.0 and CUDA 9.2, with GEFORCE GTX 1080 Ti throughout our experiments. We tune the hyper-parameter over grid search to compare the performance of different methods. For all these methods, we set the local learning rate as 0.1, turn the local learning rate decay over (0.9999, 0. 0.9995, 0.999, 0.995, 0.992). Note that when aggregating model changes the server, we introduce a global learning rate, and conduct a grid search for this global learning rate over the set 31,3, 6, 9, 12, 15, 18) for all methods. Moreover, set up 5 epochs of local updates for CIFAR-10 and CIAR-100, 10 epochs of local updates for Shakespeare. We also set 3 = 0.9 for the momentum moving average used in the global aggregation step of FedM, FedPA. FedEP, FALD, and FedMD. For our FedMD,
turn
{10-6 10-7,11100 10-1,011000 10-9, 10 10-10). }. For FedEP, we turn Ocoo over {10-2 2 10-3 10-4, 10-5 10-5,10-9) Note that in FALD, the original work employs extremely small learning rate (10-9), leading to exceedingly slow convergence. In this paper, we utilize a larger learning rate and introduce an noise-scale factor control the noise. For FedBE, We focus on Gaussian and allocate 10% of the training data of each client to be available at the server
A.3
Additional Experimental Results
A.3.1 Calibration and Uncertainty Estimation for Other Settings
Here, we present the additional experimental results uncertainty estimates, as detailed in Table
15

---

Table 7: ECE (averaged over 5 random seeds).
CIFARIO
CIFAR100
Method
Dir (0.2)-H
Dir (0.7)-H
Dir (0.2)-H
Dir (0.7)-H
FedAvg
0.168
0.170
0.434
0.440
FedM
0.164
0.162
0.471
0.461
MimeLite
0.177
0.186
0.469
0.478
SCAFFOLD
0.186
0.195
0,480
FedPA
0.167
0.174
0.382
0,380
FedEP
0.130
0.124
0.298
0.284
FALD
0.139
0.126
0.274
0.259
FedMDMI
0.125
0.122
0.267
0.262
A.3.2 Experimental Results Under Different 3
In this section, we provide the test accuracy of the proposed FedMDMI vs. FedAvg, under different choices of B, shown in Figures 4 4c).
FodMOOti
Fedling
40
15
50
02
0.1
3.6
1.0
02
2.5
0.8
6/3
02
0.1
3.5
0.
3.0
20
430
1203
Rounds
(a) CIFAR-10
(b) CIFAR-100
(c) Shakespeare
(d) Shakespeare
Figure 4: (a), (b) and (c): sensitivity analysis of 8; and (d): test accuracy of clients vs. number of communication rounds achieved by different comparison methods.
A.3.3 Convergence Curve for Shakespeare
In this section, we provide the convergence curves for different methods on the Shakespeare dataset, as shown in Figure Ad).
A.3.4 Other Experimental Results with More Settings and Model Structures
The survey paper [Cao et al. 2023 provides a overview of Bayesian federated learning (BFL), including the client- and server-side BFL and various other categories of BFL. For example, FedHB [Kim and Hospedales 2023 proposes hierarchical BFL approach, using hierarchical Bayesian modeling to describe the generative process of clients' local data with local models governed by higher-level global variate. FedSI [Chen et al. [2224] proposes personalized BFL approach, performing posterior inference (i.c., Linearized Laplace Approximation) over an individual subset of network parameters for each client, while keeping other parameters deterministic.
Extension to a PFL Context:
Inspired by FedRep [Collins et al., 2021] and FedSI [Chen et al. 2024), we consider applying algorithm to personalized Bayesian FL., named PerFed-MDMI. FedRep learns shared data representation across clients and unique local heads for each client to fulfill their personal objectives. FedSI further updates the distributions of model parameters these updated distributions to the server for global aggregation throughout the entire training. Model parameters of the decision layers then fine-tuned during the evaluation phase. For subnetwork (i.e., representation layers) inference, instead of using the Linearized Laplace Approximation from FedSI, PerFed-MDMI employs our model-data mutual information regularization technique. Additionally, since our posterior inference method does not involve covariance matrix estimation, we do not consider selecting a smaller subnetwork to reduce computational and storage overhead.
the representation layers and sends
For experimental comparison, we first evaluate our algorithm against FedProx and FedHB on a traditional FL task, namely global prediction. Results in Table 8 demonstrate that our algorithm generally outperforms FedProx and FedHB.
For the personalized FL task, we compare our newly designed algorithm, PerFed-MDML with FedAvg and FedRep, demonstrating clear improvement Table where all data each client split into 70% training and 30% test set. This indicates that mutual information regularization
16

---

Table 8: Global Performance comparison under
Table 9: Personalized Performance comparison under various settings.
various settings.
Top-l Test Accuracy (5)
Averaged accuracy (8) local
Dutavet
Settngg
FodPors
FedHB-NIW
FeIMDMI (Ours)
Ditaset
Setting
FodAAg
FodRRpp
FodMPor-MDMI (Ours)
Dir (0.2)-1.
79,82
83.19
83.42
Dir (0.2)-L
75,46
76.78
79.87
CIFAR-10
(0.71-L
80.59
63.25
83.81
CIFAR-10
Dir (0, (0:7)-1.
76.89
77.45
80.49
(0.2-H
(0.7)-H
(0.7)-H
77.10
78.97
Die (0.2)-L
42.89
48.15
49.46
Dir (0.2)-L
31,02
40.52
48.06
CI0AR-100
Dir (0.2)-11
44.12
49.26
90.45
CIFAR-100
Dir (0.2)-11
32.87
41.87
48,76
Dir (0.23-H
43.07
49,02
49.70
(0.2)-H1
(0.7)-H
44.85
49.38
50.71
(0.7)-H
33.88
42.74
49.42
technique compatible with FedRep and FedSI, effectively promoting personalized FL. Additionally, since our focus in this paper is on training a global posterior rather than the personalized posteriors, we did not directly compare the accuracy with FedSI, FedPop, and other baselines in personalized FI. due to the time constraint.
Influence of the Degree of Data Heterogeneity: We also conduct additional experiments, observing that as the degree of heterogeneity increases, the performance of all the comparison algorithms declines, as shown in I0 However, the performance of our FedMDMI decreases to lesser extent in the most cases. This indicates that our FedMDMI is more robust to the increasing heterogeneity as compared to the other baselines. Since we cannot display images here, we present these results in the following table form for now, and we will convert it to line plots in the final version of our manuscript.
We have not yet reached a conclusion on how data heterogeneity affects the ECE, as shown in It is clear that uncertainty estimation is influenced by local sample size, and smaller sample sizes tend to lead to overconfident decisions. On the other hand, data heterogeneity has a more pronounced effect on accuracy.
Table 10:
Test Accuracy (%) v.s. Data heterogeneity.
CIFARIO
CIFARTOU
Method
iid-H
Dir (0.7)-H
Dir (0.2)-H
Dir (0.1)-H
iid-H
Dir (0.7)-H
Dir (0.2)-H
Dir (0.1)-H
FedAvg
83.22
80.31
79.68
77.96
48.65
42.35
40.32
38.14
FedBE
84.06
82.33
81.25
79.19
51.68
46.29
44.82
42.06
FedPA
84.82
82.93
82.78
80.37
52.44
49,66
48.51
46.79
FedEP
84,93
83.79
83.30
81.23
$2.99
50.02
49.08
47.32
FedMDMI
85.05
83.76
83.56
82,28
53.28
50.71
49.70
48.41
Table 11: ECE v.s. Data heterogeneity.
CIFARIO
CIFARIO0
Method
iid-H
Dir (0.7)-H
Dir (0.2)-H
Dir (0.
iid-H
Dir (0.7)-H
Dir (0.2)-H
Dir (0.1)-H
FedAvg
0.174
0.170
0.168
0.160
0.437
0.440
0.434
0.430
FedBE
0.181
0.184
0.187
0.169
0.452
0.467
0.438
0.442
FedPA 0.179
0.174
0.167
0.168
0.387
0.380
0.382
0.361
FedEP 0.137
0.124
0.130
0.125
0.284
0.298
0.277
FedMDMI
0.133
0.122
0.125
0.117
0.272
0.262
0.267
0.259
Note Within the comparison methods, FedEP utilize the variational approximation obtain the global posterior. However, determining whether MCMC or VA is more suitable for f'ederated posterior inference is in general a complex problem, as their effectiveness depends on the specific context. Due to the data heterogeneity issue, regularization at the client level is often necessary to mitigate bias in the local posteriors. Our FedMDMI employs the model-data mutual information regularization alongside the MCMC for posterior inference. In contrast, FedEP uses the posterior from the previous round to regularize the current local posterior. However, in some practical federated scenarios with low client participation rates, this retained local posterior in FedEP can become very stale, thus reducing its effectiveness addressing the data heterogeneity issue. This also contributes to the less stable and slower convergence observed for FedEP compared to our FedMDMI.
17

---

Influence of Model Architectures: We evaluate the performance of our FedMDMI by using ResNet18 on the CIFAR-10 and CIFAR- 1-00 datasets. Table [2 shows that our FedMDMI still outperforms other comparison methods in terms of the generalization performance, with the ResNet18 architecture. Here, we replace the batch norm with group norm, and set the number of clients to 20, such that the 10% and 5% sampling rates correspond to only two and one client participating in training per communication round, respectively. This may highlight the robustness our FedMDMI to some possible
changes and its ability adapt to various models.
Table 12: Performance comparison on Resent-18 under various settings.
Top-1
Test Accuracy (8)
Datnset
Setting
FedAvg
FedM
Mimel.ite
SCAFFOLD
FedRF
FedPA
FedEP ())
FAI.D
FedMDMI (Ours)
(0.2)-L
85.12
86.91
87.14
CIFARIO
Dir (0.7)-L
85.41
87.29
85.71
87,10
87,35
87.91
87.46
87.01
87.82
Resent18
(0.2)-11
87.32
87.09
Dir (0.7)-H
86.20
87:47
86.09
87,98
87.51
87.56
88.06
87.92
88.25
2)-
53.19
55.34
57.21
$7.08
57.28
58.19
CIFAR100
Dir (0. 71-
53.40
55.80
53.76
57.64
57.46
57.16
57.92
57.16
58.26
R8
(0.2)-1
53.40
55.49
Dir (0.7)-H
53.91
56.03
54.16
57.98
37.81
58.24
58.13
57,59
58.81
Influence of the Number of Local Data Samples: We conduct additional experiments to analyze the impact of the number of local data samples. Given that the total number of data samples in CIFAR- 10 and CIFAR-100 is fixed, we controlled the number of samples on each client by adjusting the number of clients. With fewer data samples on a client, the local model is more prone overfitting. The results are shown in Table 13] Our FedMDMI maintains superior performance even with scarce local data, demonstrating that our MI regularization-ban posterior estimation effectively alleviates the overfitting caused by data scarcity.
Table 13: Performance Comparison across various number of data sample on clients
Top-I
Test Accuracy (%)
Dataset
Setting
FedAvg
FedM
Mimel ite
SCAFFOL.D
FedBF
FedPA
FedEP (I)
FALD
FedMDMI (Ours)
50
77.02
80.08
77.15
79.25
78.99
80.42
80.34
80.59
80.86
CIFARIO
100
77.51
80.49
77.60
79.57
79.05
80.97
81.09
80:97
81.47
(Dir L)
250
78.49
81.78
78.48
80.97
80.25
81.33
82.15
81.72
82.87
500
79.26
82.12
79.54
82.11
81.49
82.29
82.63
82.56
83,42
50
36.76
40.49
35,87
42.24
38.84
42.26
43.15
42.14
43.84
CIFARI00 (Di (0.2)-L)
100
37.05
42.86
37.88
44.41
39.71
44.81
45.41
43.89
46.08
250
39.61
45.62
39.15
46.97
43.08
46.80
47.15
45.79
47.90
Influence of Batch Size: For the batch size, we also conduct additional experiments to demonstrate the effect of batch size the performance. The results are shown in Table [4] Our method demonstrates a greater robustness across different batch sizes compared to other baselines.
Table 14: Performance Comparison across various batch-size for local optimaization
Top-1 Test Accuracy (%)
Dataset
Setting
FedAvg
Mimel
SCAFFOLDD
FedBE
FedPA
FedEP (T)
FedMDMT (Ours)
73.54
76.78
73.84
77.89
79.23
78.43
78.45
77.05
77.45
CIFARIO
79.26
82.12
79.54
82.11
81.49
82.29
82.63
82.56
83.42
(Dir (0.2)-L)
100
79.18
82.08
79.80
82,02
81.30
82.18
82.41
82.29
83.49
200
79.05
81.49
78.78
81.83
80.16
81.28
80.97
80.78
81.91
33.97
42.54
43,97
CIFARI00
50
40.35
47.13
40.25
47.76
44.35
48.59
48.89
48.43
49,46
(Dir (0.2)-L)
46.91
200
39.06
43.67
38.34
45,49
44.73
45.56
45.54
45.34
46.43
Extension to Other Dataset: We further used TensorFlow Federated to generate the Stack Overflow dataset. Similar to Cheng et al. [2023). due to the limited graphics memory and rebuttal-time, we utilize only sample of 300 clients from the original dataset, with the following observations.
18

---

Table 15: Performance comparison Stack Overflow under various settings.
Top-1
Test Accuracy (%)
Setting
FodAvg
FedM
Mimel
SCAFFOLDD
FedPA
FedEPP (1)
FAL.D
FedMDMI (Ours)
Stack Overtlow
non-iid-L
40.08
41.24
40.16 41.34
42.11
42,66
42.79
42.34
43,00
non-iid-H
41.23
42.01
42.50
42.67
43.01
42.41
43.11
As shown in Table 15 on the Stack Overflow dataset, our FedMDMI continues outperform the majority of other optimization algorithms designed address the data heterogeneity. This can be attributed to our proposed model-data mutual information regularization that enhances generalization.
To provide empirical evidence on the complexity analysis, we conduct experiments to evaluate the running time of each communication round for these algorithms. Taking CIFAR 100 dataset with participation rate m 0.1 as an example, the average time required to execute a communication round for different models LeNet and ResNet- 18 with GEFORCE GTX 1080 Ti is follows:
Experimental Results of Complexity Analysis:
For LeNet: -vgeAvg: 5.29 seconds; FedPA: 6.94 seconds; FedEP: 11.25 seconds; FedBE: 17.36 seconds; FedMDMMI 5.36 seconds.
For ResNet-18: FedAvg: 13.05 seconds; FedPA: 16.87 seconds; FedEP: 25.10 seconds; FedBE:
43.06 seconds; FedMDMI: 13.88 seconds.
This shows that the time consumed by our algorithm per communication round does not exhibit significant difference compared to FedAvg. The time required by FedBE is much larger than that of other algorithms. Moreover, as the network size increases, the time taken by all baselines also increases, which aligns with our intuition.
A.3.5 Other Experimental Results with 3-PredBayes
Table 16: ECE (averaged over 5 random seeds).
CIFAR-10
CIFAR 100
Method
(0.25-1.
(0:7)-I.
(0.2)-H
Dir (0.7)-
Dir (0.2)-1.
Dir (0.7)-1.
Dir (0.25H
Dir (0:7)-H
FedAvg
0.169
0.165
0.168
0.170
0.429
0.432
0.434
0.440
A-PredBayes
0.124
0.126
0.124
0.126
0,262
0.242
0.260
0.266
FadMDMI (Ours)
0.115
0.120
0.125
0.122
0.261
0.263
0.267
0.262
We present the experimental results of our algorithm compared with the additional baseline, 3- PredBayes (Hasan et al. 2024]. [8-PreBBayes is designed to provide well-calibrated uncertainty estimates other than a good prediction accuracy. In addition, 3-PredBayes is also designed for single round of communication, making it unfair to directly compare its accuracy with other algorithms. Moreover, the B-PredBayes requires auxiliary data for the knowledge distillation at the server. To meet this requirement, allocate 10% of the training data of each client to be available at the server.
Experimental results in Table 16 show that our algorithm outperforms 3-PredBayes in terms of the ECE in most cases. It is worth noting that enhancing the quality of the auxiliary data could improve performance of these algorithms.
We note that 3-PredBayes sometimes performs better for CIFAR-100, 100, we consider that there are two possible reasons.
federated heterogeneous scenario involving CIFAR-100, which comprises 100 categories compared to CIFAR-10, each client typically handles a subset of 13-16 (or 20-25) categories when setting Dir 0.2 (or Dir 0.7). Consequently, with such a high degree of heterogeneity, clients' training tends to incline towards overconfident predictions, possibly neglecting some categories entirely. Then, when performing the aggregation globally (e.g., in our FedMDMI algorithm), it may produce overconfident predictions. ii) 3 13-PredBayes additionally utilizes subset of the training data encompassing all categories to perform the global knowledge distillation (which may though incur the additional storage/computation and privacy leakage issues at server). This strategy significantly mitigates the overconfident predictions stemming from the training on locally heterogencous data.
19

---

A.4 Proof
A.4.1 Proof of Global Posterior Decomposition
Global posterior decomposition is widely adopted approach in federated Bayesian inference Al
Shedivat et al. [020 and embarrassingly parallel [ al. 2014 which
decomposes the global posterior into a product of client posteriors:
P(w]S) x p(S)w)p(w) (A1
p (w)St) Pi(w)
(S (w) p(w) &
p(w)
(21)
pl w
(w)S))
(s(w))
where (A1) is due to the assumption of the global likelihood being conditionally independent given w Al-Shedivat et al. 2220), Guo et al. [2023) Neiswanger et al. (2014), and r - p(w)
I P.(w) represents
the ratio the global prior to the product chinn priors. Note that there exists another posterior
decomposition Sefidgaran et al. [2023 which assumes independence among the local posteriors given
the global model of previous communication round, rather than based on the Bayesian inference principle, as shown in the following section.
A.4.2
Difference between the Bayesian Inference Approach and Model Distribution Update over Communication Rounds for FL
There an another posterior decomposition, based on the Model Distribution Update over communi- cation rounds [Chor et al. 2023 Sefidgaran et al. [2023]. To mitigate confusion, we will delineate the difference between the "Model Distribution Update over communication rounds" and "Bayesian Inference" for FL.
Posterior Decomposition based on Model Distribution Update over Communication Rounds [Chor et al., 2023 Sefidgaran et al. 2023
In this context, along this update procedure, the training process of FL is viewed as Markov process indexed by time t. That is, at each communication round t, wi is regarded as a random variable, determined by the previous round's random variable w (initialization of the global model) and the
current round's state S) i.e., p set of random variables encompasses model from all clients [M] and server across all communication rounds [T], we directly have
p (wESS), w)-1) By defining W =
(w
[M] w(rr): which denotes the
II p (wilS) w)-1) p(an loo } (E(T] [E(M)
P(WSS)
II
(22)
This equation corresponds to Eq. (2) Sefidgaran al. 2023 where w, and wi are the random variables rather than samples. It is crucial to note its primary distinction from Eq. (21) in our posterior decomposition: this equation directly assumes independence among the posteriors p(w)(Ss) given Wt-1 (the initial model at current round), without relying on Bayes' theorem or the assumption of likelihood independence given w. Additionally, it is worth mentioning that the distribution p(w:) of the global model varies over time.
ii) Global posterior decomposition based Bayesian inference
On the contrary, our entire work based on Bayesian inference. In our context, the parameters of learned global model estimate the global posterior p( (S) of model weights w given the global dataset S. Due to the inability exchange raw client data in FL, we decompose the global posterior into a product of clients' posteriors:
treated as random variables rather than single points, then aim to
m
p(w[S) x Il
p (w St),
(23)
20

---

which is based on the assumption of the global likelihood being conditionally independent given w,
i... p(S,..... Sm w) I P (S, w).
This equation indicates that under the independence assumption of likelihood and prior assumption, the posterior distribution of the full data is the product of local posterior p( [[]]S Consequently, our subsequent objective is to estimate the Icoal posterior p(w[S;) and then approximate global p(w]S). Alternative methods include obtaining samples drawn from the posterior (MCMC), or fit the optimal distribution by Gaussian assumption (Variational Inference).
Specifically, in our work, following the introduction of our newly proposed Model-Data MI regu- larization, there is an optimal global posterior p(w) St), also referred to as the Gibbs posterior. To obtain samples form global posterior p(w[S) I[ (w[S,)) through optimization (or training iteration with t), we employ the MCMC sampling (SGLD). After communication rounds and local iterations, wie denotes locally sampled model of client i, following the distribution p (wc SS)), and W2 denotes globally aggregated sample, following the distribution p(w)SS)
iii) Further clarification
In summary, we would like to highlight the difference between "Bayesian Inference" " and "Model Distribution Update over Communication Rounds" for FL from the following three perspectives.
First, within the context of our work, there is an estimated posterior p(w) (S) (i.e., Gibbs posterior) when given task S and prior assumption. We use the sample wy to approximate sample from this posterior through the MCMC sampling methods. Conversely, within the context of "Model Distribution Update over Communication Rounds", the distribution of a global model p(w,) (S) evolves with time t.
Second, within the context of our work, w is random variable of learned global model, while w, and w represent samples obtained through the MCMC sampling, following the distribution p(w)S) and p(w) S.). respectively. Consequently, the notation p(wj,,]S wt-1) is inaccurate and meaningless in our context. Conversely, within the context of "Model Distribution Update over Communication Rounds", wi and W, are both random variables at time t, and pl w |S,, w;-1) holds true.
Third, the global posterior decomposition in our work does not directly assume independence among the posteriors p( St). Instead, it relies on the conditional independence of the likelihood function and prior assumptions. Conversely, within the context of "Model Distribution Update over Communication Rounds", the posterior decomposition directly assumes independence among the posteriors p(u:|S S.)) given wt-1-
Thus, literature [Chor al. 2023, Sefidgaran et al. analyzing the relationship between the generalization bound and number of communication rounds t. However, it worth noting that these analyses fundamentally diverge from the focus of our work, which is in essence Bayesian inference-bart based FL. algorithm.
2023 offers new and intriguing direction for
A.4.3 Proof of Global Model-Data MI Decomposition
Proposition A.1. (Global Model- Date-Dataata MI Decomposition), Suppose that S (Si...., Sm } consists of data from clients, then based on the chain rule of mutual information, have:
Sw = Σ [T; - (
(24)
m
(w; S.).
where Si-1 2 (
Proof. To improve readability, we first present the typical chain rule for mutual information here. For any random variables x, y and z, we have T(x, y: =) = T(y; 2) + T(x; zly)). By iteratively applying this chain rule of mutual information, we have:
21

---

m Σ (w; S:(sst))) 1=1
[wu; S)
where Sh-1 _ St-)) Then, based on the assumption of the global likelihood being conditionally independent given w. we have
7 (w;; (w: [ ((: (S;;
(A1) (w. - (S; -
(A2) (wr; S.) 15 (S;
(25)
(Aa) I (w;; 5,) - (5;;
S I (w; St),
where (Al) and (A2) are due to the the chain rule of mutual information, and (A3) is from the assumption of the global likelihood being conditionally independent given w Al-Shedivat al.
2020 Guo et al. 2023]. That is,
(s 1, S(w) p(w) KI. (s/ (w) =0
(26)
A.4.4 Proof of Lemma 4.2
Lemma A.2. Xu and Raginsky 2017| The Gibbs posterior is the minimizer of the objective for client
update in Eq. 9
[-a [.s; (w) logpe(w))]
(w]()))
exp
(27)
B,
where B, is normalization factior.
Proof. order for our paper to self-contained, we re-state the proof in Xu and Raginsky (2017).
Wang et al. 2021b here.
For the local objective, we have:
min F((w)= E (( [[s,(w)) + (Tu; St) p(w[S;)
= -fotw p(w (i) (s)) (w)) du + / p(w, S:)[log((w (8:) log pi (w:) dw....
=
(28)
Consequently, differentiating F((w) w.r.t. p(w)S,) results in:
((() = (w) + a logp(w S.) a log (og)((w) + a.
Setting V ((()) = O, we then have:
logp (w Si) = - - Ls. (w) + logpi(w)
-
1
p (w Si) = Pi(w) exp
(w)
exp 1-1}
(29)
p (w St) x (w(w) exp
Ls, (w)
1
[Ls, (w) a log (w))]
x exp
To integrate the distribution to 1, we add an additional normalization factor B..
22

---

A.4.5 Generalization Analysis
Theorem A.3. (Generalization Bounds for FL). Suppose that Es(w, for all E M is bounded by C' and independent, then the expected generalization evvor statisfies:
C= S)
+ V
(C
E [[p((w) - Ls(w)] S
(30)
2mm
2m
where m the number of clients, v the number of samples on the client (assuming, without loss of generality, that the number of samples for all clients is equal), and D is the set of distributions p(S.),
In the subsequent proof, we primarily leverage the theorem and lemma drawn from PAC-Bayesian learning in Alquier [021] Note that while the derivation about the model-data mutual information for centrilized learning Xu and Raginsky [017] is concise, it is only applicable i.i.d. cases and may not be extended directly to federated heterogeneous scenarios.
Proof. Under the FL settings, it crucial to consider not only the gaps arising from the unseen client data (i.e., participating error), but also the gaps stemming from the unseen client distributions (i.e., participation gap). Following the two-level framework proposed by Yuan et al. [2021] and Hu et al. [2023] we define the loss functions separately as:
the population risk:
(p(w)= = Ep (p))))( = Ep(sieer (1: x]]]
(31)
2) the semi-empirical risk:
Lo(w)= Σ Lp(s:)(u== =
E 1e,
(32)
m
m
iEM
3) the empirical risk:
Ls(w) 1 m iEM
Σ (s) -s(
1
& (w, 2
(33)
74 j=1
iEM
Then, the generalization
be written as:
E [Lp (w) [s(w]
Ep.. Ep(w s) [Tp(ww) - [p(w)] + EPEp Eo(uss) [T((w) Ls(w))
= -Ep EE Ep 5) Lp(w) 1
Σ (w) REMM
+ EEpEp(S))
Σ EM
(Lp(s)(w) Ls, (w))
(34)
We first derive the upper bound of the second term in Eq. 34
Lemma A... (Generalization Bound for the second term). Let pl (a)) be a distribution over examples in client and let P be mata distribution. Let m be the number of clients and n be the number of samples on the client. Then, we have:
C2 (2I(w;)) S)
EpEp(s, Bpts())
Σ (Lp(s:)(w) - Ls, Ls,(ww) iEM
$
(35)
m
2mn
Proof. Note that Hoeffding's inequality only requires that the random variables be independent and not necessarily identically distributed. We apply Hoeffding's 's inequality to single variable Un with values the interval (a, b) for the case of n 1, which simply states that:
E ( ((UG)-E(UN))
(36)
When U1 = Ee 21)1 Eg 2x), we have:
(e Sea8
R62
Ep(s.)
(37)
23

---

That is:
E
Exi
B < (4:(10173)) E=((0, -((
Epp
e
m
B ( (e (ei(wr,z)) (12,2)))
pSp))
IIII
i=1==1
(38)
m
e c (t -c ex(we:22)
E ((s))))
2
When we have:
i D.n 2i [e (t//ott (<((ws2))) (((m))))))((m=)))
1202 8mm
Ep.p(s.)
e
<
(39)
which can be equivalently expressed as:
Epp)(s.)
(em (o. S
Pe S
(40)
on which we can build generalization bound of the second term. Then, we apply the Donsker and Varadhan's 's variational formula (Equation 2.3 on Page 13 of literature Alquier 2021 to get:
Eppp(s.)
sup
e E (ons (L(p)m))
22,2
(41)
(EP(()
where P(O) denotes the of all probability distributions and - p(w) is a prior distribution over hypothesis. Rearranging terms, we have:
EpEp(s)) sup e E. E. (ooc 4s, (L(p|())) R22 pEP(()
Somm
1.
(42)
Further, replacing À with i and p(w(S) with p. it becomes evident that:
EpEp(S) Bu
1
Σ
[[(s)[W)-
Σ
[Ls(ww) + XC2
+
K KL(p(w|S)(lz)
0.
8mn
REM
SEM
(43)
That is:
EpEp(s (ww
1
Σ (L)(s)(w) (w))
EEpsps (s,) Ee
AC2
KI(p(w(S)[lm)
m
8mn
VEM
XC 8mn
Iww; S)
+
A
(44)
The choice
/ 8mml (w; S) leads to:
1
Σ
(L)ss)(w)-Lsss
(LM())(w) - Ls. (s,(w))
C
(21(w;))
(45)
2mn
iEM
24

---

Then, we give the upper bound of the first term in Eq. (34) as follows.
This first term corresponds the participation gap. Let P be meta distribution D, where D is the set of all probability distributions, i.e., D _ we define p(w (D as ( In this context, 5 = drawn from p(S)) and p(S,) is drawn from meta distribution P, we then have:
With slight abuse of notation,
,,...., Smh, where S is
E. plss)) B p(w|s)
(T:(w) [w))
= ps,
E Ep(S))~ (ts,) [1]]
Σ iEM
E. Es,
(46)
Ep(S.)- =(pE)(u(D)
(E 2:]l
Σ ME
p(s) (, Si>l
m
For loss function (ww, 2.) bounded by upper bound C, its expectation B2.) (4 S,]] also forms bounded loss function with an upper bound C. Let (sp))(() [w, S;)] be collection of independent random variables from P, wwe can directly derive the following conclusion based on variant of the PAC-Bayes theorem (i.e., Theorem 2.8 in [Alquier [2021)).
Let P be a meta distribution over D, where (p(S)),.... p(S..)) And let å p(w) be a prior distribution over hypothesis. Then, for any > 0 and for from distribution p[w] D).
the set of all probability distributions, i.e., A
AC²
(L(p(u|D)]lt
Ep Ep [p @EpBuwppee D) (D(W)
+
(47)
8mn
A
where p(w D) = Es, (p([[()) We then have:
AC²
KL(p( D)][T]
E Ep [pp(w LD(W)] SE
+
8mn
(48)
2C
1ww; D) X
8m
where Ep KL [ D). In particular, this mutual information term [(w; D) cannot be "estimated", since it depends on the statisties of the ponttatiiiipaiing clients" which is not available. Furthermore, / (w; D) is tighter than (ww S) due the convexity of the KL divergence.
The choice À = / / (w; D leads to:
[3; D)
T. Euonptll))) (T) Ln()) S
(49)
2m
Based on Eq.
49 and Eq. 35 we have:
C2 C2I(w; S)
C C2I D) 2m
E (p() Ls(w))
+
(50)
2mn
A.4.6 Proof of Client-level Differential Privacy
Assuming that ((w, z/) is I-smooth, when t 2 serves client-level (e, diffferrrnt privacy.
vora [og(2/8) our proposed FedMDMI also pre-
Proof. We have the following Gaussian mechanism for differential privacy Dwork et al.) [2044] Let E E (0, 1) be an arbitrary number. For c² > 2 l1(1225/8), the Gaussian mechanism with standard deviation parameter
2Lc/s is (e, diffferreetllll private. Wang et al. (0055 show that under this 2'n
Gaussian mechanism, in standard SGLD,
327 log(2/8) ensures the privacy loss to smaller than
</N with probability > 32rk
r5
where is the total number of samples and represents
2nk
25

---

the samples sampled in the current iteration. We have an extra parameter va in the noise variance.
Thus, we need
32r (2g(2 for preserving the client-level (e, diifffrrrelt privacy.
voe"n
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The Summary and Introduction of the manuscript accurately reflect the scope of the research, centered around the federated learning, and delineate contribution, as evidenced in Line 73.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the "Limitations" of work in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiaatinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
27

---

Answer: [Yes]
Justification: In section of Appendix, we provide the assumption and proofs related to the theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: In Section A.1 we provided all the details of the algorithm.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
28

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide the reproducible code in the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ praips.cc/ubblic/guuueee CoodeSumeeieeeeee ssi iessonPolll molyicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section A.2 of the Appendix, we elucidate the network architecture, hetero- geneous data splits, and hyperparameter configuration utilized in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We re-run experiments over different seeds to get the variance measurements of test accuracy eliminate the impact of random seeds, shown in Table 2
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
29

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative етrоr rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Section A.2.3 provide the type of computing workers GPU and the version PyTorch.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification: Our work adheres the NeurIPS Code of Ethics in all aspects.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Federated learning an important paradigm. Our work proposes new algorithm for this paradigm from the perspective federated posterior inference and provides differential privacy as byproduct. This research has the potential to benefit fields such engineering, economics, biology, and epidemiology, which often involve multi-organization collaboration but do not allow for the sharing of their private data.
machine learning
Guidelines:
30

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The code package, dataset, and network architecture employed in the experi- ment are appropriately referenced.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
31

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.com//aaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Our work does not involve the release of new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
32

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "uP79pdHwoa", "reviewer": "Reviewer_4nGu", "summary": "The paper considers the problem of Bayesian Federated Learning when there’s data heterogeneity and class imbalance across clients and develops a posterior inference approach for model parameters through mutual information regularization of the data and global parameters in local posteriors. This is achieved via using the KL formulation of the mutual information and showing that the optimal local posterior is a Gibbs distribution. To infer local posterior Stochastic Gradient Langevin Dynamics is used and to aggregate the local posteriors into the global posteriors a simple average of the local samples is considered. The theoretical results suggest that the proposed algorithm results in the convergence of the local and global posteriors and corresponding generalization bounds are provided. Results are shown on a few datasets where competitive test performance and uncertainty calibration are achieved.", "questions": "Even though I went through the paper in full I’m still not fully convinced that mutual information regularization is a good idea. The bound provided in Eq. 20 suggests that the generalization error is bounded by two mutual information terms one of which is the target of estimation of the paper. How can we argue that there’s no tradeoff between lowering the first and second terms? Why does reducing the first term result in a lower generalization error?", "limitations": "See the weaknesses section.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "*   The problem considered is timely and important. In almost every application of federated learning, client data heterogeneity and scarcity are inevitable. For data scarcity, considering the model uncertainty (through Bayesian modeling) and for data heterogeneity removing the biases of local and global posteriors due to biases in local datasets are sensible directions.\n    \n*   Accompanying the empirical results and intuitions with theoretical arguments is another strength of the paper.\n    \n*   The writing is clear and understandable, and the organization of the arguments is intuitive and logical.", "weaknesses": "*   Although competing (Bayesian) methods are not specifically designed to deal with data heterogeneity there seems to be very little to no gain achieved by the model. The results in the current form seem weak to me. The performance gains are quite marginal (e.g. Fig. 3) both for uncertainty calibration and test performance. Are the performance differences provided in tables 2, and 3 statistically significant? Shouldn’t we assume that as the degree of heterogeneity increases the proposed method is more effective? Why don’t we see a larger gap as a function of $\\\\alpha$?\n    \n*   A few important papers seem to be missing from the introduction such as \\[1\\]. Some methods are not considered for the comparisons such as \\[2,3\\].\n    \n*   The only information about the time (and memory) complexity is a sentence somewhere close to the end of the paper. While an important motivation of the paper is to reduce the computational and memory cost, empirical timing results aren’t shown. Can you include a detailed timing comparison as a function of data heterogeneity, dimension and size of the network, etc? Intuitively variational methods should do a much better job compared to sampling-based methods. Although the stochastic version of the Langevin Dynamics is used here still they should take much longer than variational methods to converge.\n    \n\n\\[1\\] Cao, Longbing, et al. \"Bayesian federated learning: A survey.\" arXiv preprint arXiv:2304.13267 (2023).\n\n\\[2\\] Chen, Hui, et al. \"FedSI: Federated Subnetwork Inference for Efficient Uncertainty Quantification.\" arXiv preprint arXiv:2404.15657 (2024).\n\n\\[3\\] Kim, Minyoung, and Timothy Hospedales. \"Fedhb: Hierarchical Bayesian federated learning.\" arXiv preprint arXiv:2305.04979 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank for the comments, and our response is as follows.\n\n**Comment 1-Further Clarification on Fig.3:** Fig. 3 visualizes uncertainty calibration and convergence, with quantitative results in Tables 2 and 3 highlighting improvements of our FedMDMI.\n\ni) Figs. 3(a) and 3(b) provide an effective visual representation of model calibration. In a well-calibrated model, the difference between confidence and accuracy should be close to zero for each bin. Our FedMDMI (black curve) shows a closer alignment to zero across most bins, indicating a superior calibration. This is further supported by Table 3, where our FedMDMI exhibits the lowest expected calibration error in most cases, particularly on CIFAR-100 dataset.\n\nii) Figs. 3(c) and (d) illustrate convergence behavior. Although FedMDMI converges slower than SCAFFOLD, it slightly outperforms other Bayesian baselines. This is explained in Line 358 of the manuscript, attributing the slower convergence to the variance from stochastic gradient Langevin dynamics (SGLD).\n\niii) Figs. 3(c) and (d)  also show that the final convergence of our FedMDMI achieves higher accuracy. This was further demonstrated in Table 2. Since CIFAR-10 dataset is relatively simple, the baselines can easily achieve high accuracy, making our improvement less significant. However, for the more complex datasets like CIFAR-100 and Shakespeare, our FedMDMI shows a more significant improvement.\n\niv) Regarding our response to Comment 6, the computational and storage complexity of our FedMDMI is lower than that of most other Bayesian FL algorithms.\n\n**Comment 2-Performance differences in Tables 2 and 3:** Please note that results presented in Tables 2 and 3 are averaged over five random seeds, as indicated in the caption of both tables.\n\n**Comment 3- Degree of heterogeneity and effectiveness:** Yes, we agree with the reviewer. As the degree of heterogeneity increases, performance of all baselines decreases, but our method's performance decreases less significantly. This indicates that our FedMDMI is more robust to increasing heterogeneity compared to the baseline.\n\nSpecifically, for CIFAR-10, as the degree of heterogeneity increases from Dir(0.7)-H to Dir(0.2)-H, performance of our FedMDMI and FedAvg decreases by 0.2% and 0.63%, respectively. Similarly, for CIFAR-100, performance of our FedMDMI and FedAvg decreases by 1.01% and 2.03%, respectively. This trend holds in most cases, indicating that our proposed posterior inference approach based on model-data mutual information regularization can effectively alleviate the impact of data heterogeneity.\n\n**Comment 4-Discussion of $\\alpha$:** The Lagrange multiplier $\\alpha$ in Eq. (7) balances fitting ($L_{S_i}(w)$) and generalization ($I(w; S_i)$). Fig. 2 illustrates that as $\\alpha$ increases, the gap between test errors across varying data heterogeneity decreases, indicating the efficacy of our proposed regularization in addressing data heterogeneity. Specifically, increasing $\\alpha$ leads to a gradual rise in train error and an initial drop followed by an increase in test error, suggesting that lower complexity (larger $\\alpha$) may cause underfitting, while higher complexity (smaller $\\alpha$) risks overfitting. Thus, larger $\\alpha$ mitigates the bias from data heterogeneity.\n\n**Comment 5-Additional Baselines:**  Please see common comment 2.\n\n**Comment 6-Complexity Analysis:**  Please see common comment 1.\n\n**Comment 7-Comparison of MCMC and VA:** Please see common comment 3.\n\n**Comment 8-Effectiveness of MI regularization:**  As discussed in the Related Work section, Model-Data Mutual Information (MDMI) regularization is a well-established concept in centralized learning, supported by both theoretical analyses and extensive experimental verification. For instance, [A1] demonstrates that MDMI regularization outperforms L2-norm regularization and dropout through Fisher information matrix estimation.\n\nFor FL, we propose a federated posterior inference method to estimate the global mutual information via global MDMI Decomposition. Our experiments indicate that MDMI regularization is equally effective in FL.\n\n[A1] Z. Wang *et al.*, \"Pac-Bayes information bottleneck.\" in *Proc. ICLR*, 2022.\n\n**Comment 9-Difference between mutual information terms:**\nIn Eq. (20), $I(w; S)$ relates to the participating generalization error, i.e., the difference between the empirical and expected risk for participating clients. This serves as the regularized term that can be estimated in our FedMDMI. On the other hand, $I(w; D)$ relates to the participation gap, i.e., difference in the expected risk between participating and non-participating clients. This term, however, cannot be estimated due to unavailability of the non-participating clients.\n\nHere, we will illustrate the difference with a simple example. Consider an extreme case with 10 clients, each holding only one label of MNIST dataset. In this case, we restrict federated training on the first 9 clients, which contain train dataset labeled 0-8. During the test phase, the error measured on test sets from labels 0-8 is referred to as the participating generalization error $I(w; S)$. The error measured on data from the $10$-th client, which holds label 9, is termed as the participation gap $I(w; D)$.\n\nThus, $I(w;D)$ and $I(w;S)$ do not represent a compromise or antagonistic relationship. By optimizing $I(w;S)$, we can improve performance of the learned model on the distribution seen during training. Consequently, if the distribution of clients that have never participated is the same as the distribution seen by trained clients, then $I(w;D)$ degenerates to $I(w;S)$. In other words, our FedMDMI can only guarantee improved generalization on the distribution seen during training, but not on the out-of-distribution data from non-participating clients. This could motivate an interesting future direction: reducing $I(w;D)$ through techniques like transfer learning or domain generalization.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8oqjLEtggZ", "reviewer": "Reviewer_re5T", "summary": "In this paper, the authors introduce a method for federated learning to bypass problems caused by inter-client data heterogeneity.\nFor this, they introduce a Bayesian approach with information-theoretic regularizer, that will prevent local models from overfitting. Specifically, the authors add a model-data regularizer at a global level and then show how it could be computed in a federated fashion. The local optimal posterior appeared to be the Gibbs posterior, and the authors employ SGLD to sample from it. To show the efficacy of the approach, authors conduct a series of experiments on image and text data at federated datasets.", "questions": "- In Line 76 it is said that the MI term (...) \"offers a certain client-level privacy protection as a byproduct.\"\nI think the authors imply that the MI term regularizes the fitness of a local model. I wonder, how overfitness can compromise the privacy of training data.\n- In lines 171-175 there is a reasoning about some bias factor $\\delta$ that affects the generation of local data. What is $\\delta$? Why does inequality in 174 hold? Can you elaborate more on the last sentence in lines 174-175?  \n- In lines 212-213 authors use the term Variational Approximation (VA). I have a feeling that the term Variational Inference (VI) is a more common alternative to MCMC met in literature.", "limitations": "The authors have a separate section on limitations. However, not all of them are addressed (see the Weaknesses section, about upper-bound on a loss).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the paper has the following strengths:\n- I find the paper very easy to follow, and the suggested idea is very interesting and natural;\n- It tackles an important problem of data heterogeneity in Federated Learning;\n- The paper provides detailed theoretical derivations and experimental evaluation;", "weaknesses": "I find the following things are downsides:\n- The loss used in the optimization is a result of several levels of approximations.\nFirst, the global model-data MI term is upper-bounded by the sum of local model-data MI terms. \nSecond, each local model-data term is itself upper-bounded by the RHS of Eq. 17. \nI think that the discussion on the tightness of the upper bound is missing.\n\n- I feel that some baselines are missing. Probably the first and classical approach to combat the problem of data heterogeneity was FedProx [1] (which was cited), but it is not compared with.\nAlso, it would be interesting to see, where the method is placed if compared with personalized approaches, that are specially built to deal with data heterogeneity. E.g. FedPop [2] (Bayesian), FedRep [3] (not Bayesian).\n\n[1] Li T. et al. Federated optimization in heterogeneous networks //Proceedings of Machine learning and systems. – 2020. – Т. 2. – С. 429-450.\n\n[2] Kotelevskii N. et al. Fedpop: A bayesian approach for personalised federated learning //Advances in Neural Information Processing Systems. – 2022. – Т. 35. – С. 8687-8701. \n\n[3] Collins L. et al. Exploiting shared representations for personalized federated learning //International conference on machine learning. – PMLR, 2021. – С. 2089-2099.\n\n- Minor typos:\n1) Lines 142-143 collapsed (negative vspace?);\n2) Figure 2: The left y-axis lives in [0, 1], right y-axis in [0, 100].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank for the comments, and our response can be found in the following.\n\n**Comment 1-Discussion on tightness of the upper bound:** \n\nFirst, we are constrained in practice to only leverage local data $S_i$ at individual client $i$ under the FL settings. Alternatively, based on the Global Model-Data MI Decomposition in Proposition 4.1, we have: $$I(w ; S) = \\sum_{i=1}^m \\left[ I\\left(w ; S_i\\right) - I\\left(S_i ; S^{i-1}\\right) \\right] \n\\le \\sum_{i=1}^m I\\left(w ; S_i\\right).$$\nHere, we can upper bound $I(w ; S)$ by the sum of local MI terms $I(w ; S_i)$ that can be estimated locally by these clients. By doing so, we also introduce an estimation error, i.e., the data correlation $I\\left(S_i ; S^{i-1}\\right)$ between the individual clients, which determines the tightness of the proposed upper bound $\\sum_{i=1}^m I\\left(w ; S_i\\right)$. This data correlation $I\\left(S_i ; S^{i-1}\\right)$ between the individual clients is related to the clients' data generation and collection, which is intractable to estimate and optimize.\n\nThen, each local MI term $I(w; S_i)$ can be further expressed as:\n$$ I(w ; S_i) = \\mathbb{E}\\_{p(S_i)}[\\operatorname{KL}(p(w \\mid S_i) \\| p_i(w))] \n  = \\mathbb{E}\\_{p(S_i)}[\\operatorname{KL}(p(w \\mid S_i) \\| r(w))]-\\operatorname{KL}[p_i(w) \\| r(w)] \\leq \\mathbb{E}\\_{p(S_i)}[\\operatorname{KL}(p(w \\mid S_i) \\| r(w))]\\],$$\nwhere $p_i(w) \\triangleq \\mathbb{E}_{p(S_i)}[p(w | S_i)]$ denotes the oracle prior and  $r(w)$ is an arbitrary prior distribution that is used to approximate $p_i(w)$. This may incur another estimation error, which is the difference between the oracle prior $p_i(w)$ and the prior $r(w)$ we actually use, i.e., $\\operatorname{KL}\\left[p_i(w) \\| r(w)\\right]$.\n\nAs indicated in Line 240, [B1] emphasizes that to achieve a small KL divergence, the prior must, in essence, predict the posterior. Specifically, [B1] employs a distribution of pre-trained models derived from a portion of the untrained data as prior $r(w)$.\n\nIn our work, we propose using the global model in the previous round as the mean $\\mu$ of Gaussian prior and the uncertainty introduced by all clients in the previous round as the covariance $\\Sigma^{-1}$ in the prior, as shown in Eq. (18) and Line 6 in Algorithm 1 of our FedMDMI. This incorporates the global data information and potentially helps predict the local priors based on the global posterior decomposition.\n\n[B1] G. Dziugaite _et al._, \"On the role of data in PAC-Bayes bounds,\" in _Proc. ICML_, 2021.\n\n**Comment 2-Additional baselines:** \n\nPlease see our response to common Comment 2.\n\n**Comment 3-Minor typos:**\n\n Thank you for pointing out the small issue with collapsed Lines 142-143, and we will fix it. Additionally, in Figure 2, the ordinate represents the percentage error. Specifically, the left figure indicates that the training error is within the range of [0, 1%], while the right figure indicates that the test error is within the range of [0, 20%].\n\n**Comment 4-Relationship between Overfitness and Privacy of Training Data:**\n\nThe MI term $I(w;S)$ regularizes the fitness of a local model. In fact, the hyperparameter $\\alpha$ controls the degree of the MI regularization and embodies a trade-off between privacy and fitting (or trade-off between generalization and fitting). As highlighted in Line 275, the proposed MI regularizer $I(w;S)$ directly quantifies the extent to which the model memorizes data. A larger value of $I(w;S)$ implies that the model fits the data well, thus signifying the fitting. Conversely, a smaller $I(w;S)$ indicates that the model avoids memorizing excessive data details, thereby promoting generalization and, from another perspective, privacy protection.\n\nIn fact, some work [B2] demonstrates that in certain cases, knowing the previous model and the gradient update from a client can allow one to infer a training example held by that user. Therefore, for stronger privacy protection, differential privacy can be used to encrypt these models or gradient updates. Our work shows that regularizing the MI term can provide a form of differential privacy protection as a byproduct.\n\n[B2] P. Kairouz _et al._, Advances and open problems in federated learning, Foundations and Trends in machine learning, 2021.\n\n**Comment 5- Further Clarification on Factor $\\delta$ :** \n\nHere, we would like to explore and offer further insights, from the perspective of FL, on how minimizing mutual information can help mitigate bias resulted from the data heterogeneity.\n\nIn FL, we denote the bias factor as $\\delta$ that affects the generation of local heterogeneous data (such as the difference in clients' locations, preferences, or habits). If we can effectively eliminate the influence of the bias factor $\\delta$, the entire dataset would become homogeneous (similar), thereby reducing bias among local posteriors.\n\nIndeed, directly eliminating the impact of bias factor is impractical. But leveraging Markov chain $\\delta \\rightarrow S \\rightarrow w$ allows us to directly infer that $I(w; \\delta) \\le I(w; S)$, where the inequality holds due to the **Data Processing Inequality (DPI)**. As a consequence, we can diminish $I(w; S)$ to constrain $I(w; \\delta)$, thereby rendering the model $w$ insensitive to the bias factor $\\delta$ from the diverse clients.\n\nFurthermore, similar analyses are prevalent in other contexts. For instance, in centered learning, literature [B3] (Section 2.2 on Page 6) assumes that nuisances affect the observed data, and these nuisances are mitigated through information bottleneck regularization (another well-known mutual information).\n\n[B3] A. Achille _et al._, \"Emergence of invariance and disentanglement in deep representations,\" Journal of Machine Learning Research, 2018.\n\n\n**Comment 6-Comparison between MCMC and Variational Approximation (VA):** \n\nPlease see our response to common Comment 3.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BQWGbk90EO", "reviewer": "Reviewer_XkDq", "summary": "The paper proposes an approach to mitigate training failure in the heterogeneous federated learning setup. The approach combines Bayesian perspective of posterior inference on the client side and regularization of mutual information between weights and data in order to reduce the effect of difference in the local datasets. The authors provide computable approximations of the values involved and provide information theoretic bound on the generalization in federated setup. Experiments show that the method succeeds and outperforms several baselines.", "questions": "1 - In equation (20) there are two different mutual information terms, I(w;S) and I(w;D). What is the difference between them?\n\n2 - What are the conditions under which it is possible to decompose posterior of the global model into the product of posteriors of the local models?\n\n3 - In the conclusion you claim to show that the optimal posterior is Gibbs posterior, but as I understand you used this result from previous research?", "limitations": "Limitations are addressed.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper gives a detailed explanation of the approach for heterogeneous federated learning with mutual information regularization and posterior estimation. A generalization bound is provided and extensive experiments are performed.", "weaknesses": "The motivation to introduce posterior inference in federated setup is not clear: The problem of the point estimation and inability to have uncertainty of the predictions is equally valid for the centralized setup as well. The further description of the approach is convoluted, it seems that the main reason to resort to posterior estimation is to be able to obtain tractable computation for the mutual information term and corresponding generalization bound from PAC-Bayes perspective.\n\nThe method is motivated by scarce local data that might lead to overfitting and prevent from training a global model when aggregated, but this setup is not checked empirically. Only the heterogeneity with respect to labels distribution is evaluated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank for the comments, and our response can be found in the following.\n\n**Comment 1-Further clarification on Our Motivation:**\n\nWe agree with the reviewer that in centralized learning, posterior inference is proposed to provide a more reliable assessment of model uncertainty than point estimation. This uncertainty estimation is also crucial in the safety-critical applications of federated learning, such as autonomous driving, healthcare, and finance. Consequently, there has been some recent literature, such as FedPA and FedEP, considering the introduction of posterior inference into the federated setup.\n\nHowever, we found that these Bayesian inference-based federated learning algorithms do not adequately address the data heterogeneity issue. Therefore, we are motivated to propose a mutual information regularization to enforce the global posterior to learn essential information from the heterogeneous local data, thereby improving the generalization capability. To optimize this regularization, we further employ a series of techniques, including the global mutual information decomposition, PAC-Bayesian conclusions, and stochastic gradient Langevin dynamics (SGLD) sampling.\n\n**Comment 2-Additional Experiments on Scarce Local Data:** \n\nThank you for your comments. We have conducted additional experiments to analyze the impact of the number of local data samples. Given that the total number of data samples in CIFAR-10 and CIFAR-100 is fixed, we controlled the number of samples on each client by adjusting the number of clients. With fewer data samples on a client, the local model is more prone to overfitting. The results are shown in Table 11. Our FedMDMI maintains superior performance even with scarce local data, demonstrating that our MI regularization-based posterior estimation effectively alleviates the overfitting caused by data scarcity.\n\n**Comment 3-The difference between mutual information terms, I(w; S) and I(w; D):** \n\nIn Eq. (20), $I(w; S)$ relates to the participating generalization error, i.e., the difference between the empirical and expected risk for the participating clients. This serves as the regularized term that can be estimated in our FedMDMI. On the other hand, $I(w; D)$ relates to the participation gap, i.e., the difference in the expected risk between the participating and non-participating clients. This term, however, cannot be estimated due to the unavailability of the non-participating clients. \n\nHere, we will illustrate the difference with a simple example. Consider an extreme case with 10 clients, each holding only one label of the MNIST dataset. In this scenario, we restrict federated training on the first 9 clients, which contain train dataset labeled 0-8. During the test phase, the error measured on test sets from labels 0-8 is referred to as the participating generalization error $I(w; S)$. The error measured on the data from the $10$-th client, which holds label 9, is termed as the participation gap  $I(w; D)$.\n\nThus, $I(w;D)$ and $I(w;S)$ do not represent a compromise or antagonistic relationship. By optimizing $I(w;S)$, we can improve the performance of the learned model on the distribution seen during training. Consequently, if the distribution of clients that have never participated is the same as the distribution seen by trained clients, then $I(w;D)$ degenerates to $I(w;S)$. In other words, our FedMDMI can only guarantee improved generalization on the distribution seen during training, but not on the out-of-distribution data from non-participating clients. This could motivate an interesting future direction: reducing $I(w;D)$ through techniques like transfer learning or domain generalization.\n\n**Comment 4-What are the conditions under which it is possible to decompose the posterior of the global model into the product of posteriors of the local models?**  \n\nThe decomposition of the global posterior into the product of posteriors of the local models is based on two assumptions:\n\n1. The global likelihood is conditionally independent given $w$, i.e., $p(S_1,\\dots,S_m | w) = \\prod_{i=1}^m p(S_i | w)$.\n2. The ratio of the global prior to the product of client priors, $\\frac{p(w)}{\\prod_{i=1}^m p_i(w)}$, is considered a constant based on some prior assumptions, such as Gaussian priors or uniform priors.\n\nThese assumptions are widely recognized and significant across various fields, including in PoE [C1] and (EP)-MCMC [C2, C3], and Bayesian FL [C4].\n\n[C1] Hinton, Geoffrey E., \"Training products of experts by minimizing contrastive divergence,\" Neural computation 14(8): 1771-1800, 2002.\n\n[C2] Neiswanger, Willie, Chong Wang, and Eric P. Xing, \"Asymptotically exact, embarrassingly parallel MCMC,\" in _Proc. UAI_, 2014.\n\n[C3] Wang, Xiangyu, and David B. Dunson, \"Parallelizing MCMC via Weierstrass sampler,\" arXiv preprint arXiv:1312.4605, 2013.\n\n[C4] Al-Shedivat, Maruan, _et al._, \"Federated learning via posterior averaging: A new perspective and practical algorithms,\" in _Proc. ICLR_, 2021.\n\n**Comment 5-In the conclusion, you claim to show that the optimal posterior is the Gibbs posterior, but as I understand, you used this result from previous research?**\n\nWe agree with the reviewer that this is not our contribution. To avoid confusion, we emphasized in the original manuscript that \"this conclusion is well-established in the field of PAC-Bayesian learning\" in the Introduction (Line 64). We also stated that, \"In order for our paper to be self-contained, we re-state the proof from [12, 44] here for the optimal posterior\" in the Proof (Line 988). To avoid overclaiming our contribution, we will further emphasize that the result of the optimal posterior being the Gibbs posterior stems from previous research.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wx6KDBeNND", "reviewer": "Reviewer_hP14", "summary": "This paper proposed a Federated Learning model that support Bayesian inference. To alleviate potential bias induced from local client data, a regularisation constraint on model-data mutual information is introduced. The authors show that the MCMC inference with the regularisation can be implemented through the stochastic gradient Langevin dynamics. The author also proves a generalisation bound.", "questions": "- Figure 3 is very busy. Can you elaborate how they show the proposed method \"ourperforms the other algorithms\"? In particular, if the MI constraint regularises the effect of local client bias, why does the proposed method show volatile learning curve?\n- Discuss the computational complexity of the proposed approach and how adjusting the batch size affects the performance.", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Information-theoretical modelling of federated learning is not very common. More exploration in this area is important.\n- The authors provide justifications to most of the design decisions and they look sound to me.", "weaknesses": "- The writing is sometimes difficult to follow. The relevance of some theoretical analyses is not always clear. It would be better to introduce Algorithm 1 in the main text early on. Steps in line 10 and line 16 seem to be the key differences of the proposed method; They should be reflected in Figure 1. \n- The proposed method is computationally more complex than the point estimates. The experimental results of FedMDMI do not show significantly better performance compared to FedEP. A comparison of computational cost is required to understand the potential advantages and compromises compared with all baselines.\n- FedMDMI has been evaluated on small models that may not fully reflect its performance in scenarios requiring larger and more complex architectures such as ResNets. Evaluating FedMDMI on such models would provide a more comprehensive evaluation of its scalability and generalization capabilities.\n- The Dirichlet distribution is used to simulate class imbalance data heterogeneity. Using heterogeneous datasets provided by TensorFlow Federated and LEAF would ensure that the method is evaluated under natural heterogeneous data distributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank for the comments, and our response can be found in the following.\n\n**Comment 1-The writing is sometimes difficult to follow:** \n\n1) Due to the space limit of the main text, we placed the theoretical proof, more detailed analysis, and Algorithm 1 in the Appendix of our submission. To improve readability, we will include these core steps of theoretical analysis and Algorithm 1 in the main text in the final version of this paper, where there is one additional page allowed.\n\n2) Line 10 of Algorithm 1 represents the local SGLD update, which has been reflected in Figure 1 as the “SGLD for sampling” step. In the modified version, we will also reflect the global sliding average step (i.e., Line 16 of Algorithm 1) in Figure 1.\n\n\n**Comment 2-A Comparison of Computational Cost:** \n\nPlease see our response to common Comment 1. Here, we first provide a complexity analysis concerning the computation time, storage, and communication for our method compared to other Bayesian methods. Subsequently, our experiments examine the relationship between computation time and model dimension within a single communication round. It is shown that the computation time of our algorithm is significantly superior to that of FedEP.\n\n**Comment 3-Additional Experiments on ResNet-18:** \n\nFollowing this suggestion, we have evaluated the performance of our FedMDMI by using ResNet18 on the CIFAR-10 and CIFAR-100 datasets, with the following observations.\n\nTable 13 shows that our FedMDMI still outperforms other comparison methods in terms of the generalization performance, with the ResNet-18 architecture. Here, we replace the batch norm with group norm and set the number of clients to 20, such that the 10% and 5% sampling rates correspond to only two and one client participating in training per communication round, respectively. This may highlight the robustness of our FedMDMI to some possible model-architecture changes and its ability to adapt to various models.\n\n**Comment 4-Additional  datasets provided by TensorFlow Federated and LEAF:**\n\nPlease note that in the original manuscript, we used LEAF to generate the heterogeneous Shakespeare dataset, with each client associated with a speaking role comprising a few lines. Following this advice, we have further used TensorFlow Federated to generate the Stack Overflow dataset. Similar to [D1], due to the limited graphics memory and rebuttal time, we utilize only a sample of 200 clients from the original dataset, with the following observations.\n\nAs shown in Table 14, on the Stack Overflow dataset, our FedMDMI continues to outperform the majority of other optimization algorithms designed to address the data heterogeneity issue. This can be attributed to our proposed model-data mutual information regularization that enhances generalization.\n\n[D1] G. Cheng _et al._, \"Federated asymptotics: a model to compare federated learning algorithms,\" in _Proc. AISTATS_, 2023.\n\n**Comment 5-Further Clarification on Fig. 3:**\n\ni) Figures 3(a) and 3(b) provide an effective visual representation of model calibration. In a well-calibrated model, the difference between confidence and accuracy should be close to zero for each bin. Our FedMDMI (black curve) shows a closer alignment to zero across most bins, indicating superior calibration. This is further supported by Table 3, where our FedMDMI exhibits the lowest expected calibration error in most cases, particularly on the CIFAR-100 dataset.\n\nii) Figures 3(c) and 3(d) show the convergence behavior of different algorithms. Here, we would like to show that though our FedMDMI does not converge the fastest, especially compared to the traditional optimization-based algorithm SCAFFOLD, it still slightly outperforms the other Bayesian inference-based baselines. In fact, we have also given an explanation on this observation in Line 358 of the original manuscript: \"One potential explanation is our use of stochastic gradient Langevin dynamics (SGLD) to approximate the posterior, which often suffers from a slow convergence rate due to the variance introduced by the stochastic gradient.\"\n\niii) Figures 3(c) and 3(d) also show that the final convergence of our FedMDMI achieves higher accuracy. This was further demonstrated in Table 2. Since the CIFAR-10 dataset is relatively simple, the baselines can easily achieve high accuracy, making our improvement less significant. However, for the more complex datasets like CIFAR-100 and Shakespeare, our FedMDMI shows a more significant improvement.\n\nThe volatile learning curve mentioned by the reviewer may be due to the fact that it was plotted averaged over 5 random seeds and drawn by sampling 25 points apart, resulting in a curve that is not completely smooth. In addition, we will show here the standard deviation of the accuracy on each convergence curve when it is close to convergence (i.e., 2000-4000 rounds), as follows.\n\nFor **CIFAR10 on Dir (0.2)-L**: FedAvg (0.44%), FedPA (0.74%), FedEP (0.34%), FedBE (0.51%), FALD (1.17%), FedMDMI (0.28%).\n\nFor **CIFAR100 on Dir (0.2)-L**: FedAvg (0.58%), FedPA (1.19%), FedEP (0.85%), FedBE (0.69%), FALD (0.96%), FedMDMI (0.79%).\n\nSuch experimental results show that, in most cases, our FedMDMI is less volatile during the convergence stage compared to other posterior inference-based algorithms.\n\n**Comment 6-Discussion on Computational Complexity:**\n\nPlease also see our response to common Comment 1. \n\nFor the batch size, in the original manuscript it was set to the default value of 50 for CIFAR-10 and CIFAR-100, following the same setting as in the other baselines. We have also conducted additional experiments to demonstrate the effect of batch size on performance. The results are shown in Table 12. Our method demonstrates greater robustness across different batch sizes compared to other baselines.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thanks for the comments. Below is our response to common cocerns, with new tables in the attached PDF.

**Comment 1-Complexity analysis :** We provide complexity analysis w.r.t. time, storage, and communication of different methods. We further empirically examine the relationship between computation time, data heterogeneity, and model dimension within a communication round.

1) We begin by defining the dimensions of neural network as $d$. At each round, we adopt SGLD (Stochastic Gradient Langevin Dynamics) to estimate the local posterior. Specifically, compared to SGD employed in **FedAvg** and **FedBE**, **our FedMDMI** entails an additional step of generating Gaussian noise with $d$ dimensions, which is subsequently incorporated into each model update iteration. This causes additional $\mathcal{O}(d)$ time and $\mathcal{O}(d)$ memory. In contrast, at each client, **FedPA** uses dynamic programming to approximate the inverse matrix $d \times d$ of neural network, incurring additional $\mathcal{O}(l^2d)$ time and $\mathcal{O}(ld)$ memory, where $l$ is number of posterior samples. Similarly, **FedEP** also requires approximating the covariance as the inverse Hessian, introducing additional $\mathcal{O}(d^3)$ time and $\mathcal{O}(d^2)$ memory.

For communication and aggregation at server, our **FedMDMI**, along with **FedAvg** and **FedPA**, requires $\mathcal{O}(md)$ time and $\mathcal{O}(md)$ memory, where $m$ denotes number of clients. In contrast, **FedEP** requires $\mathcal{O}(md)$ time and $\mathcal{O}(md^2)$ memory. **FedBE** not only requires $\mathcal{O}((m+1+l)d)$ time and $\mathcal{O}((m+1+l)d)$ memory, where $l$ denotes number of global model samples, but also performs knowledge distillation at server using unlabeled data, which is both memory- and time-intensive.

2) We also conduct experiments to evaluate running time of each communication round for these algorithms. Taking CIFAR100 dataset with participation rate $\frac{p}{m}=0.1$ as example, the average time to execute a communication round for LeNet and ResNet-18 with GEFORCE GTX 1080 Ti is:

**LeNet**:
FedAvg: 5.29 s
FedPA: 6.94 s
FedEP: 11.25 s
FedBE: 17.36 s
FedMDMI: 5.36 s

**ResNet-18**:
FedAvg: 13.05 s
FedPA: 16.87 s
FedEP: 25.10 s
FedBE: 43.06 s
FedMDMI: 13.88 s

This shows that time consumed by our FedMDMI per communication round does not exhibit significant difference than FedAvg. The time required by FedBE is much larger than other methods. As network size increases, the time taken by all baselines also increases, aligning with our intuition.

Finally, data heterogeneity does not affect local computation time when numbers of local samples and updates are fixed.

**Comment 2-Additional Baeslines:** We will include these references in the revised manuscript. The survey provides an overview of Bayesian federated learning (BFL). For example, FedHB proposes a hierarchical BFL approach, using hierarchical Bayesian modeling to describe the generative process of clients' local data with local models governed by a higher-level global variate. FedSI proposes a personalized BFL, performing posterior inference over an individual subset of network parameters for each client, while keeping other parameters deterministic.  

Inspired by FedRep and FedSI, we apply our algorithm to personalized Bayesian FL, named PerFed-MDMI. FedRep learns a shared data representation across clients and unique local heads for each client to fulfill their personal objectives. FedSI further updates the distributions of model parameters over the representation layers and sends these updated distributions to the server for global aggregation in training. Model parameters of the decision layers are then fine-tuned during evaluation phase. For subnetwork (i.e., representation layers) inference, instead of using Linearized Laplace Approximation from FedSI, PerFed-MDMI employs our model-data mutual information regularization. Also, since our posterior inference method involves no covariance matrix estimation, we do not consider selecting a smaller subnetwork to reduce computation and storage overhead.

Empirically, we first evaluate our method against FedProx and FedHB on traditional FL task, named global prediction. Table 9 shows that our method outperforms FedProx and FedHB.

For personalized FL task, we compare our newly designed method, PerFed-MDMI, with FedAvg and FedRep, presenting improvement in Table 10, where all data at each client is split into 70\% training set and 30\% test set. This indicates that our MI regularization is compatible with FedRep and FedSI, effectively promoting personalized FL. Additionally, since our focus is on training a global posterior rather than the personalized posteriors, we did not directly compare the accuracy with FedSI, FedPop, and other baselines in personalized FL due to the time constraint.

**Comment 3-Comparison of MCMC and Variational Approximation (VA):**

We use SGLD instead of VA for three reasons.

1) Optimal posterior obtained in our FedMDMI is the Gibbs posterior, and SGLD has been proven efficient and effective in large-scale Gibbs posterior inference.

2) As shown in Line 212, there has also been pioneering work on obtaining a Gibbs posterior through VA. But this method results in at least doubling the communication overhead due to transmission of both the mean and covariance matrices. Another primary drawback of VA is its tendency to yield biased posterior estimates for complex posterior distributions. This issue is further compounded in FL, where data heterogeneity across clients already contributes to biased local posteriors, and this bias may be exacerbated by VA.

3) As shown in Figs. 3(c)(d) and 4(d), convergence rate of our FedMDMI may not be the fastest, especially when compared to the optimization-based method, SCAFFOLD. But it still slightly outperforms Bayesian-based baselines in most cases, showing that SGLD in FL is efficient and does not significantly reduce the convergence rate.

Meta-review: The paper addresses a timely and important problem in federated learning by proposing a novel approach to tackle data heterogeneity and scarcity using Bayesian modeling for uncertainty estimation and mutual information regularization. The proposed methodology is well-motivated and combines theoretical arguments with empirical results, making it a strong contribution to the field. The clear and logical writing, along with the detailed theoretical derivations, further enhances the paper's strength. Notably, the paper provides a thorough exploration of an underexplored area in federated learning, offering sound justifications for most design decisions and demonstrating the potential of information-theoretical modeling in this context.

Despite its strengths, the paper has some limitations that reviewers noted. The performance gains observed were marginal, and there were concerns about the lack of comprehensive comparisons with recent state-of-the-art baselines specifically designed to address data heterogeneity in federated learning. Some important baselines were missing from the experiments, and the empirical evaluations were conducted on small-scale models, raising questions about the method's scalability to more complex architectures. 

Overall, while there are areas for improvement, the paper's strengths in tackling a crucial problem with novel and theoretically grounded methods justify its acceptance.
Decision: Accept (poster)