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_cc_2024_Conference
Published: 2024-05-15 15:31:59

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.

Main paper: 
 introduction: 
 Federated learning (FL), as an emerging distributed learning framework, has garnered considerable attention [McMahan et al., 2017] . In this paradigm, clients collaborate to train a single global model, under the central coordination of a server. Notably, such a collaborative training proceeds without the necessity of sharing or 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 a diminished generalization performance for the globally aggregated model at the server.A multitude of efforts have been made to address this issue from various perspectives, including federated optimization [Acar et al., 2021 , Karimireddy et al., 2020] , federated domain generalization [Nguyen et al., 2022 , de Luca et al., 2022] , federated knowledge distillation [Zhu et al., 2021, Afonin and Karimireddy, 2021] , etc. Despite their progress, a majority of the existing FL formulations consistently treat it as a point-estimate (i.e., a single value estimated for each model weight) based on the loss function optimized on local client data. From a 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 a 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 FL, a typical framework comprises: i) each client initializing a local model based on global posterior parameters and independently conducting a 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 a 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 a stateful method, FedEP is also not a suitable solution for lowparticipation 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 information-theoretic perspective, we propose to infer the global posterior 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 to 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 a Gibbs posterior, a conclusion well-established in the field of PAC-Bayesian learning [Alquier, 2021] . Then, we employ the stochastic gradient Langevin dynamics (SGLD) method [Welling and Teh, 2011] to sample from this local Gibbs posterior, which provides an unbiased and efficient sampling-based 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 to have a non-asymptotic convergence guarantee with the true global posterior in decentralized and federated scenarios [Gürbüzbalaban et al., 2021 , Plassier et 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.• We introduce an information-theoretic approach for the global posterior inference in FL. It incorporates a 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 a generalization bound for FL w.r.t. the global model-data MI, showing that regularizing this global MI leads to a reduction 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-based baselines, while providing well-calibrated uncertainty estimates.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: 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.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the 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 or 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
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 sample of client i's model at round t and step k wt sample of aggregated server model at round tThe objective of standard FL is to learn a single global model w from m clients via the following optimization problem:EQUATIONwhere L p(Si) (w) = E zi∼p(Si) [ℓ i (w, z i )] represents the local expected risk of the i-th client associated with data distribution p(S i ) on local data S i , and the loss function ℓ i (w, z i ) is usually chosen as a negative log likelihood, i.e., -log p (z i |w). We further define a set of random variables S = {S 1 , . . . , S m } as the global dataset, consisting of m heterogeneous client datasets. Moreover, we let {z j i } ni j=1 denote the local training dataset of size n i drawn independently from distribution p(S i ). This setup also gives rise to the global empirical risk as:EQUATIONwhere L Si (w) = 1 ni ni j=1 ℓ i (w, z j i ) is the local empirical risk. From a probability theory perspective, employing a single point-estimate as the model weight, as shown in Eq. ( 2), 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 w as a random variable rather than a single point, and then estimate the global posterior of model w given the global dataset S. Subsequently, the objective of FL converts to minimizing the posterior expected loss:EQUATIONwhere p(w|S) = p(w|S 1 , . . . , S m ) denotes the global posterior. By taking an additional expectation 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.However, it is intractable to directly infer the global posterior, since raw data of the clients cannot be exchanged in FL. A widely adopted solution is the global posterior decomposition [Neiswanger et al., 2014 , Al-Shedivat et al., 2020] , which decomposes the global posterior into a product of client's local posteriors, i.e., p(w|S) ∝ τ i∈M p (w|S i ) , where τ ≜ p(w) m i=1 pi(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 A.4.1 for the detailed derivation of this decomposition.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 a 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 obtain 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(w * |S) = i∈M p(w * |Si); Client update 2 : p(w * |Si) = arg min p(w|S i ) E p(w|S i ) LS i (w). (4)This iterative update process is similar to FedAvg [McMahan et al., 2017] , but has distinctive features: i) each client estimates a distribution p(w|S i ) instead of a single value of w, and ii) the global aggregation on the server is changed from weighted averaging to multiplication. FedPA [Al-Shedivat et 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-estimates, 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., 2023] . This then raises a fundamental question: how can we achieve a globally well-generalized posterior in federated heterogeneous scenarios.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 5.4.Guidelines:• The answer NA means that the paper has no 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 (e.g., independence assumptions, noiseless settings, model well-specification, 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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 on the mutual information (MI) between the model and data (MD). In this section, we elaborate on our FedMDMI with an overview illustrated in Figure 1 , and procedure summarized in Algorithm 1 in Appendix A.1. The model-data mutual information is denoted as I(w; S), which quantifies relevance between the model w and input data S, and also serves as a measure of complexity of the learned model for generalization analysis. By incorporating this information-theoretic constraint as regularization, the original optimization formulation in Eq. ( 3) is converted to:EQUATIONwhere the introduced MI constraint compresses information stored in model weights, and I c 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 δ that affects the generation of local heterogeneous data (such as the difference in clients' locations, preferences, or habits), we then have the Markov chain δ → S → w, implying I(w; δ) ≤ I(w; S). As a consequence, we can diminish I(w; S) to constrain I(w; δ), thereby rendering the model w insensitive to the bias factor δ from the diverse clients.In essence, this newly introduced model-data MI term I(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, 2017] 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.While our new formulation in Eq. ( 5) is conceptually promising, it is still infeasible to directly measure the global model-data MI term I(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; S i ), which can be determined locally by these clients. Proposition 4.1. (Global Model-Data MI Decomposition). Suppose that S = {S 1 , . . . , S m } consists of data from the m clients and S i-1 ≜ {S 1 , . . . , S i-1 }, then based on the chain rule of MI, we have:EQUATIONProof. See Appendix A.4.3 for the detailed proof.Therefore, instead of I(w; S), we can constrain its upper bound m i=1 I (w; S i ) in Eq. ( 5). By further introducing a Lagrange multiplier α ≥ 0, Eq. ( 5) then re-formulates to:EQUATION)where α is also viewed to balance the fitting and generalization. Based on the global posterior decomposition, the iterative update process to optimize Eq. ( 7) becomes:EQUATIONAt each round, client update needs not to fully infer optimal local posterior. Consistent with FedAvg, it only updates for a certain number of iterations before sending learned local posterior to server.For the client update in Eq. ( 9), the second term I(w; S i ) can be expressed as:EQUATIONwhich denotes the expectation of Kullback-Leibler (KL) divergence between the posterior p(w|S i ) and marginal distribution p i (w)3 over the local data distribution p(S i ). Notably, this KL divergence plays a key role in the well-known PAC (Probably Approximately Correct)-Bayes bound4 . PAC-Bayes learning [Alquier, 2021] provides a tight generalization bound for learning algorithms, where the KL divergence between the posterior p(w|S i ) and prior p i (w) is a dominant term within this bound. For a fixed prior, pioneer works [Xu and Raginsky, 2017, Alquier et 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 a typical Gibbs distribution.Lemma 4.2. [Xu and Raginsky, 2017] The Gibbs posterior is the minimum of the objective for client update in Eq. ( 9):EQUATIONwhere B i is a normalization factor.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) [Alquier et al., 2016] , which generally relies on simplified posterior distributions. However, a 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 be 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 a slower convergence. Thus, we adopt the stochastic gradient Langevin dynamics (SGLD) [Welling and Teh, 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:EQUATIONwhere w i t,k represents client i's model at round t and step k, η t,k L is the local step size, h t,k is a noise variable sampled from N 0, 2η t,k L αI with I being the identity matrix, andEQUATIONis an unbiased estimate of gradient. Note that U i corresponds to the exponent in the Gibbs posterior of Eq. ( 11), and p i (w) | w=w i t,k-1 represents the prior p i (w)'s probability density value at w = w i t,k-1 .In theory, if the step size is annealed as η t,k L → 0, the client update sequence w i t,k converges to the local Gibbs posterior in Eq. ( 11) with sufficiently large k and t. The question then becomes: how can we obtain samples from the global posterior (i.e. the product of local Gibbs posteriors) expressed asEQUATIONbased on the samples drawn from these local posteriors.Recent studies [Gürbüzbalaban et al., 2021 , Plassier et 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(w|S), there is a non-asymptotic convergence guarantee5 . 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 w i t,k from the local posterior p(w|S i ) in Eq. ( 11) through a Markov chain with steps:EQUATIONwhere ∆w i t,k ≜ w i t,k -w i t,k-1 . After the client performs K steps and uploads the model change, the server can approximate a sample w t+1 from the global posterior p(w|S) in Eq. ( 14) simply by taking an averaging (i.e., ∆w t ≜ w t+1 -w t = 1 m i∈M ∆w i t ), through a global Markov chain with steps:EQUATION)We then discuss how the prior is chosen in Eq. ( 13). This oracle prior (i.e., p i (w) ≜ E p(Si) [p(w|S i )]) renders the mutual information the tightest, which, however, is infeasible to obtain. In fact, we can use an arbitrary prior r(w) (e.g., N (0, I)) to approximate p i (w), based on following upper bound: I(w; S i ) = E p(Si) KL p(w|S i )∥r(w) -KL p i (w)∥r(w) ≤ E p(Si) KL p(w|S i )∥r(w) .(17) For the arbitrary prior r(w) to achieve a smaller MI, it must essentially predict the posterior [Dziugaite et al., 2021] . Thus, we consider using the information of global model w t as the prior of clients for update of the next round. Based on Eq. ( 16) for updating w t , we then define p i (w) ≜ N (w|µ t-1 , Σ t-1 ), whereEQUATION)Here the variance captures the uncertainty introduced by all the clients at the previous round t -1. Meanwhile, Zhang et al. [2022] also show that using the global model as a local prior can alleviate local overfitting. Additionally, w t is only a sample drawn from the global posterior. To better estimate the mean of this Gaussian prior, we approximate it with a global moving average, which is also utilized to accelerate the convergence6 , as shown in Line 16 of Algorithm 1 in Appendix A.1.5 FedMDMI Analysis 
generalization analysis via model-data mi: 
 We first provide an information-theoretic generalization bound in terms of the model-data MI for 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:EQUATIONwhere we denote P as a meta-distribution on D, and D is the set of all distributions p(S i ). This formulation takes into account the participation gap, as compared to Eq. (1). By recalling the global empirical risk L S (w) in Eq. (3), we define the expected generalization error as E [L P (w) -L S (w)].Theorem 5.1. (Generalization Bounds for FL). Suppose that ℓ i (w, z j i ) for all i ∈ M is bounded by C and independent, then the expected generalization error satisfies:EQUATION)where m is the number of clients, n is 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 i ).Proof. See Appendix A.4.5 for the detailed proof.Remark 5.2. On the RHS of Eq. ( 20), I(w; S) relates to participating generalization error and serves as the regularized mutual information term that can be estimated in our FedMDMI, while I(w; D) relates to the participation gap that cannot be estimated due to unavailability of the non-participating clients. Moreover, this generalization bound O 1 √ mn + 1 √ m matches with the current bound [Hu et al., 2023] . Remark 5.3. Several prior works [Yagli et al., 2020 , Barnes et al., 2022 , Sefidgaran et al., 2022] have also established a generalization bound based on clients' local MI I(w; S i ) 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 I(w t , S t i ). In contrast to them, our generalization bound is based on a participation gap framework, with the first term (global MI I(w; S)) exhibiting a tighter bound than the previous local MI i I(w; S i ). For more details, we provide more elaboration of this distinction in Appendix A.4.2.As demonstrated in Theorem 5.1, the MI I(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 dataset. The experiments involve 100 clients with a 10% participation rate (refer to Section 6 for details). The train and test errors of the global model are plotted in Figure 2 under varying α. It can be seen that with the increase of α, the train error increases slowly, while the test error decreases initially and then increases. This suggests that models with lower complexity (larger α) tend to underfit, while those with higher complexity (smaller α) tend to overfit. The MI I(w; S) allows us to recover a federated version of the bias-variance tradeoff. Moreover, as α increases, the gap between the test errors under different data heterogeneity decreases, indicating that the introduced I(w; S) can render w insensitive to data heterogeneity and alleviate the bias caused by such heterogeneity.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 I(w; S) directly quantifies the extent to which the model memorizes data. Conversely, it reflects the degree of data information leakage. In the extreme case when I(w; S) = 0, w becomes entirely independent of S (e.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 ℓ i (w,z j i ) is L-smooth, when k ≥ O √ αϵ 2log(2/δ) , our FedMDMI also preserves a client-level (ϵ, δ)-differential privacy. See Appendix A.4.6 for details.Though FedMDMI enjoys multi-fold benefits, including mitigating bias and overfitting on heterogeneous 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., 2014] 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., 2021 , 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. 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 are 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. One potential explanation is because of 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 [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 in Appendix A.3.Calibration and Uncertainty Estimation. We evaluate the uncertainty estimates of all comparison methods, as shown in Table 3 and Figures 3(a ) and 3(b). 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-based 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 α and β. We have analyzed the impact of α that balances fitting and generalization in Section 5.2. We then analyze the hyperparameter β, 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. 
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 a federated version of the bias-variance trade-off. 
acknowledgement: 
 This work was supported in part by the National Natural Science Foundation of China under Grant 62125109, Grant T2122024, Grant 62320106003, Grant 62371288, 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. 
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 is an important privacy-protecting machine learning paradigm. Our work proposes a new algorithm for this paradigm from the perspective of federated posterior inference and provides differential privacy as a byproduct. This research has the potential to benefit fields such as engineering, economics, biology, and epidemiology, which often involve multi-organization collaboration but do not allow for the sharing of their private data.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 a 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 to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?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 a 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 to 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 a 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 experiment 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
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 to 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS 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. 
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 crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• 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. 
 
 Appendix: A.1 Detailed Algorithm DesignHere, we present a detailed learning procedure for our proposed FedMDMI, as outlined in Algorithm 1.Algorithm 1: Proposed FedMDMI 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 in the following.w i t,0 = w t 6 set p(w) = N w w t , 1 m K k=1 2η t-1,k L αI 7 for each local step k = 1, 2, . . . , K do 8 U i w i t,k-1 = L Si (w i t,k-1 ) -α log p(w)| w=w i t,k-1 9 h t,k ∼ N 0, 2η t,k L αI 10 w i t,k = w i t,k-1 -η t,k L ∇U i w i t,k-1 + h t,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 . To generate heterogeneous splits for CIFAR-10 and CIFAR-100, we adopt a 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 a federated dataset. In this process, a vector of the same size as the number of classes is generated for each client using the Dirichlet distribution. These vectors represent the class 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 is 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 a Dirichlet factor of 0.2, each client has 80% of samples predominantly belonging to three or four different classes. And for a 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 to 2000, a methodology akin to FedDyn [Acar et al., 2021] . In the heterogeneous splits, each client is associated with a speaking role comprising a few lines.For CIFAR- For Shakespeare, we utilize an RNN that takes a sequence of characters as input, embedding each into an 8-dimensional learned space. This sequence is then processed through a two-layer LSTM with a 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 [Li et al., 2020] and FedDyn [Acar et al., 2021] . 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 a 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.9995, 0.999, 0.995, 0.992}. Note that when aggregating model changes on the server, we introduce a global learning rate, and conduct a grid search for this global learning rate over the set {1, 3, 6, 9, 12, 15, 18} for all methods. Moreover, we set up 5 epochs of local updates for CIFAR-10 and CIAR-100, 10 epochs of local updates for Shakespeare. We also set β = 0.9 for the momentum moving average used in the global aggregation step of FedM, FedPA, FedEP, FALD, and FedMD. For our FedMD, we turn α over {10 -6 , 10 -7 , 10 -8 , 10 -9 , 10 -10 }. For FedEP, we turn α cov over {10 -2 , 10 -3 , 10 -4 , 10 -5 , 10 -6 }. Note that in FALD, the original work employs an 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 to 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.Here, we present the additional experimental results on uncertainty estimates, as detailed in Table 7 . In this section, we provide the convergence curves for different methods on the Shakespeare dataset, as shown in Figure 4 (d).Extension to a PFL Context: 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 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 [Chen et al., 2024] proposes a personalized BFL approach, performing posterior inference (i.e., Linearized Laplace Approximation) over an individual subset of network parameters for each client, while keeping other parameters deterministic.Inspired by FedRep [Collins et al., 2021] and FedSI [Chen et al., 2024] , we consider applying 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 throughout the entire training. Model parameters of the decision layers are 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.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-MDMI, with FedAvg and FedRep, demonstrating a clear improvement in Table 9 , where all data at each client is split into 70% training set and 30% test set. This indicates that our mutual information regularization 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 10. However, the performance of our FedMDMI decreases to a 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 11.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. Note Within the comparison methods, FedEP utilize the variational approximation to obtain the global posterior. However, determining whether MCMC or VA is more suitable for federated 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 in addressing the data heterogeneity issue. This also contributes to the less stable and slower convergence observed for FedEP compared to our FedMDMI.We evaluate the performance of our FedMDMI by using ResNet18 on the CIFAR-10 and CIFAR-100 datasets. Table 12 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 of our FedMDMI to some possible model-architecture changes and its ability to adapt to various models. Influence of Batch Size: For the batch size, we also conduct additional experiments to demonstrate the effect of batch size on the performance. The results are shown in Table 14 . Our method demonstrates a greater robustness across different batch sizes compared to other baselines. As shown in Table 15 , on the Stack Overflow dataset, our FedMDMI continues to outperform the majority of other optimization algorithms designed to 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 CIFAR100 dataset with participation rate p 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 as follows:For LeNet: -FedAvg: 5.29 seconds; -FedPA: 6.94 seconds; -FedEP: 11.25 seconds; -FedBE: 17.36 seconds; -FedMDMI: 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 a 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. We present the experimental results of our algorithm compared with the additional baseline, β-PredBayes [Hasan et al., 2024] . β-PredBayes is designed to provide well-calibrated uncertainty estimates other than a good prediction accuracy. In addition, β-PredBayes is also designed for a single round of communication, making it unfair to directly compare its accuracy with other algorithms. Moreover, the β-PredBayes requires auxiliary data for the knowledge distillation at the server. To meet this requirement, we 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 β-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 β-PredBayes sometimes performs better for CIFAR-100, we consider that there are two possible reasons. i) In a 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) β-PredBayes additionally utilizes a 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 heterogeneous data.Global posterior decomposition is a widely adopted approach in federated Bayesian inference Al-Shedivat et al. [2020] and embarrassingly parallel (EP)-MCMC Neiswanger et al. [2014] , which decomposes the global posterior into a product of client posteriors:EQUATION)where ( A1) is due to the assumption of the global likelihood being conditionally independent given w Al-Shedivat et al.[2020], Guo et al. [2023] , Neiswanger et al. [2014] , and τ ≜ p(w) m i=1 pi(w) represents the ratio of the global prior to the product of client 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.There is an another posterior decomposition, based on the Model Distribution Update over communication 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.i) 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 a Markov process indexed by time t. That is, at each communication round t, w i t is regarded as a random variable, determined by the previous round's random variable w t-1 (initialization of the global model) and the current round's state S t i , i.e., p w i t |S t i , w t-1 . By defining W = w[M ][T ] , w [T ] , which denotes the set of random variables encompasses model from all clients [M] and server across all communication rounds [T] , we directly haveEQUATIONThis equation corresponds to Eq. ( 2) in Sefidgaran et al. [2023] , where w t and w i t 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 i t |S i ) given w t-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 t ) of the global model varies over time.On the contrary, our entire work is based on Bayesian inference. In our context, the parameters of a learned global model w are treated as random variables rather than single points, then we aim to estimate the global posterior p(w|S) of model weights w given the global dataset S. Due to the inability to exchange raw client data in FL, we decompose the global posterior into a product of clients' posteriors:EQUATIONwhich is based on the assumption of the global likelihood being conditionally independent given w, i.e., p(S 1 , . . . , S m |w) = m i=1 p (S i |w). This equation indicates that under the independence assumption of likelihood and prior assumption, the posterior distribution of the full data S is the product of local posterior p(w|S i ). Consequently, our subsequent objective is to estimate the lcoal posterior p(w|S i ) and then approximate global p(w|S). Alternative methods include obtaining samples drawn from the posterior (MCMC), or to fit the optimal distribution by Gaussian assumption (Variational Inference). Specifically, in our work, following the introduction of our newly proposed Model-Data MI regularization, there is an optimal global posterior p(w|S i ), also referred as the Gibbs posterior. To obtain samples form global posterior p(w|S) = m i=1 p (w|S i ) through optimization (or training iteration with t), we employ the MCMC sampling (SGLD). After t communication rounds and k local iterations, w i t,k denotes a locally sampled model of client i, following the distribution p (w|S i ), and w t denotes a globally aggregated sample, following the distribution p(w|S).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 a task S and prior assumption. We use the sample w t 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 t |S) evolves with time t.• Second, within the context of our work, w is a random variable of learned global model, while w t and w i t,k represent samples obtained through the MCMC sampling, following the distribution p(w|S) and p(w|S i ), respectively. Consequently, the notation p(w i t,k |S i , w t-1 ) is inaccurate and meaningless in our context. Conversely, within the context of "Model Distribution Update over Communication Rounds", w i t,k and w t are both random variables at time t, and p(w i t,k |S i , w t-1 ) holds true. • Third, the global posterior decomposition in our work does not directly assume independence among the posteriors p(w|S i ). 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(w i t |S i ) given w t-1 .Thus, literature [Chor et al., 2023 , Sefidgaran et al., 2023] offers a new and intriguing direction for analyzing the relationship between the generalization bound and number of communication rounds t. However, it is worth noting that these analyses fundamentally diverge from the focus of our work, which is in essence a Bayesian inference-based FL algorithm.Proposition A.1. (Global Model-Data MI Decomposition). Suppose that S = {S 1 , . . . , S m } consists of data from m clients, then based on the chain rule of mutual information, we have:EQUATIONwhere S i-1 ≜ {S 1 , . . . , S i-1 }.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 I(x, y; z) = I(y; z) + I(x; z|y). By iteratively applying this chain rule of mutual information, we have:Theorem A.3. (Generalization Bounds for FL). Suppose that ℓ i (w, z j i ) for all i ∈ M is bounded by C and independent, then the expected generalization error statisfies:EQUATION)where m is the number of clients, n is 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 i ).In the subsequent proof, we primarily leverage the theorem and lemma drawn from PAC-Bayesian learning in Alquier [2021] . Note that while the derivation about the model-data mutual information for centrilized learning in Xu and Raginsky [2017] is concise, it is only applicable to i.i.d. cases and may not be extended directly to federated heterogeneous scenarios.Proof. Under the FL settings, it is 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:1) the population risk:EQUATION2) the semi-empirical risk:EQUATION3) the empirical risk:EQUATIONThen, the generalization error can be written as:EQUATIONWe first derive the upper bound of the second term in Eq. (34).Lemma A.4. (Generalization Bound for the second term). Let p(S i ) be a distribution over examples in client i and let P be a mata distribution. Let m be the number of clients and n be the number of samples on the client. Then, we have:EQUATION)Proof. Note that Hoeffding's inequality only requires that the random variables be independent and not necessarily identically distributed. We apply Hoeffding's inequality to a single variable U 1 with values in the interval [a, b] for the case of n = 1, which simply states that:EQUATION)When U 1 = E zi∼p(Si) [ℓ i (w, z i )] -ℓ i (w, z i ), we have:EQUATION)That is:EQUATION)When t = t mn , we have:EQUATIONwhich can be equivalently expressed as:EQUATIONon which we can build generalization bound of the second term. Then, we apply the Donsker and Varadhan's variational formula (Equation 2.3 on Page 13 of literature Alquier [2021] ) to get:EQUATIONwhere P(Θ) denotes the set of all probability distributions and π ≜ p(w) is a prior distribution over hypothesis. Rearranging terms, we have:EQUATIONFurther, replacing λ with t and p(w|S) with ρ, it becomes evident that:EQUATIONThat is:EQUATIONThe choice λ = 8mnI(w; S)/C 2 leads to:EQUATIONThen, we give the upper bound of the first term in Eq. ( 34), as follows.This first term corresponds to the participation gap. Let P be a meta distribution over D, where D is the set of all probability distributions, i.e., D ≜ {p(S 1 ), ..., p(S m )}. With a slight abuse of notation, we define p(w|D) as E Si∼p(Si) [p(w|S) ]. In this context, S = {S 1 , . . . , S i , . . . , S m }, where S i is drawn from p(S i ) and p(S i ) is drawn from meta distribution P , we then have:EQUATIONFor a loss function ℓ i (w, z i ) bounded by an upper bound C, its expectation E zi∼p(Si) [ℓ i (w, S i )] also forms a bounded loss function with an upper bound C. Let E zi∼p(Si) [ℓ i (w, S i )] be a collection of independent random variables from P , wwe can directly derive the following conclusion based on a variant of the PAC-Bayes theorem (i.e., Theorem 2.8 in Alquier [2021] ).Let P be a meta distribution over D, where D is the set of all probability distributions, i.e., D ≜ {p(S 1 ), ..., p(S m )}. And let π ≜ p(w) be a prior distribution over hypothesis. Then, for any λ > 0 and for w from distribution p(w|D),EQUATIONwhere p(w|D) = E Si∼p(Si) [p(w|S) ]. We then have:EQUATIONwhere E P [KL(p(w|D)∥π)] = I(w; D). In particular, this mutual information term I(w; D) cannot be "estimated", since it depends on the statistics of the "non-participating clients" which is not available. Furthermore, I(w; D) is tighter than I(w; S) due the convexity of the KL divergence.The choice λ = 8mI(w; D)/C 2 leads to:EQUATIONBased on Eq. ( 49) and Eq. ( 35), we have:EQUATIONA.4.6 Proof of Client-level Differential PrivacyAssuming that ℓ i (w,z j i ) is L-smooth, when k ≥ O √ αϵ 2log(2/δ) , our proposed FedMDMI also preserves client-level (ϵ, δ)-differential privacy.Proof. We have the following Gaussian mechanism for differential privacy Dwork et al. [2014] . Let ε ∈ (0, 1) be an arbitrary number. For c 2 > 2 ln(1.25/δ), the Gaussian mechanism with standard deviation parameter σ ≥ 2Lc/ε is (ε, δ)-differentially private. Wang et al. [2015] show that under this Gaussian mechanism, in standard SGLD, k ≥ ϵ 2 n 32τ log(2/δ) ensures the privacy loss to be smaller thanϵ √ N √ 32τ k log(2/δ)with probability > 1 -τ δ 2nk , where n is the total number of samples and τ represents the samples sampled in the current iteration. We have an extra parameter √ α in the noise variance. Thus, we need k ≥ √ αϵ 2 n 32τ log(2/δ) for preserving the client-level (ϵ, δ)-differential privacy.NeurIPS Paper ChecklistHere, clients are not required to estimate an optimal local posterior before uploading. Similar to Fe-dAvg[McMahan et al., 2017], the local posteriors are uploaded after a certain number of updates.Specifically, pi(w) ≜ E p(S i ) [p(w|Si)], which is also called the oracle prior.The PAC-Bayes bound still holds even when pi(w) and the posterior p(w|Si) are chosen arbitrarily.While this convergence bound O((1-γµ/8) KT +1/m) is established only for µ-strongly convex functions, empirical evaluations inPlassier et al. [2023] and our results show that it still holds in non-convex settings.We have also incorporated this technique into the other baselines in our experiments for a fair comparison. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: *   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$?
    
*   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\].
    
*   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.
    

\[1\] Cao, Longbing, et al. "Bayesian federated learning: A survey." arXiv preprint arXiv:2304.13267 (2023).

\[2\] Chen, Hui, et al. "FedSI: Federated Subnetwork Inference for Efficient Uncertainty Quantification." arXiv preprint arXiv:2404.15657 (2024).

\[3\] Kim, Minyoung, and Timothy Hospedales. "Fedhb: Hierarchical Bayesian federated learning." arXiv preprint arXiv:2305.04979 (2023).
questions: 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: limitations: See the weaknesses section.
paper_summary: paper_summary: In this paper, the authors introduce a method for federated learning to bypass problems caused by inter-client data heterogeneity.
For 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.
strength_weakness: strength_weakness: I find the following things are downsides:
- The loss used in the optimization is a result of several levels of approximations.
First, the global model-data MI term is upper-bounded by the sum of local model-data MI terms. 
Second, each local model-data term is itself upper-bounded by the RHS of Eq. 17. 
I think that the discussion on the tightness of the upper bound is missing.

- 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.
Also, 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).

[1] Li T. et al. Federated optimization in heterogeneous networks //Proceedings of Machine learning and systems. – 2020. – Т. 2. – С. 429-450.

[2] Kotelevskii N. et al. Fedpop: A bayesian approach for personalised federated learning //Advances in Neural Information Processing Systems. – 2022. – Т. 35. – С. 8687-8701. 

[3] Collins L. et al. Exploiting shared representations for personalized federated learning //International conference on machine learning. – PMLR, 2021. – С. 2089-2099.

- Minor typos:
1) Lines 142-143 collapsed (negative vspace?);
2) Figure 2: The left y-axis lives in [0, 1], right y-axis in [0, 100].
questions: questions: - In Line 76 it is said that the MI term (...) "offers a certain client-level privacy protection as a byproduct."
I 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.
- 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?  
- 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: 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).
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

The 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.
questions: 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?

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

3 - 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: Limitations are addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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. 
- 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.
- 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.
- 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.
questions: 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?
- Discuss the computational complexity of the proposed approach and how adjusting the batch size affects the performance.
limitations: limitations: N/A

Comments: title: Official Comment by Authors:
Thanks for the response
Official Comment by Authors:
Reply
Official Comment by Authors:
Official Comment by Authors:
Official Comment by Authors:
Official Comment by Authors:
Official Comment by Authors:
Updated review
comment: Dear Reviewer hP14,

Many thanks for the time and effort that you have dedicated to reviewing our paper and providing these insightful comments, which will further help enhance the quality of the final version of our manuscript.
I appreciate the author's effort in improving the paper. I will raise my score.
Dear Reviewer 4nGu,

We would like to express our gratitude once again for the time and effort you have dedicated in reviewing our paper, including the rebuttal and discussion phases. Your insightful comments are invaluable for enhancing our work. In the final version of our manuscript and supplementary material, we will include these additional experimental results, complexity analysis regarding the time and memory, comparison of MCMC and variational inference, and the variances and error bars in the tables and plots.
Thanks for the additional explanations and clarifications. Your responses clarify most of my questions and confusion. Please include this additional information about the time and memory complexity in the supplementary. Also please add the variances and error bars to your tables and plots. In the discussion section of your paper, elaborating on MCMC vs. variational methods would be very helpful to place your work in the broader literature.

I'm going to increase my rating of the paper. Although there are clear and significant improvements in the performance and calibration results compared to competing methods, the improvements look marginal. Nevertheless, I think the paper deserves acceptance as it provides thorough experimental and theoretical arguments and results show clear advantages to other models (in terms of performance and runtime).
Dear Reviewer XkDq,

We sincerely appreciate the time and effort that you have dedicated in reviewing our paper and providing these insightful comments, which will further help improve the quality of the final version of our manuscript.
Dear Reviewer re5T,

We sincerely appreciate your time and effort in reviewing our paper and providing valuable comments. We will incorporate these additional baselines and new supplementary experiments into the revised version. Additionally, personalized Bayesian federated learning will also be included in the discussion. These insights will further help enhance the quality of the final version of our manuscript.
[A3] Plassier _et al._, "Federated averaging langevin dynamics: Toward a unified theory and new algorithms," in _Proc. AISTATS_, 2023.

[A4] Karimireddy, Sai Praneeth _et al._, "SCAFFOLD: Stochastic controlled averaging for federated learning," in _Proc. ICML_, 2020.

**Comment 3-Comparison of MCMC and Variational Approximation (VA):** Within the comparison methods, FedEP does utilize the variational approximation to obtain the global posterior. However, determining whether MCMC or VA is more suitable for federated 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 in addressing the data heterogeneity issue. This also contributes to the less stable and slower convergence observed for FedEP compared to our FedMDMI.
**Table 3: ECE v.s. Data heterogeneity.**
| Method  | **CIFAR10** |   |   |   | **CIFAR100** |   |  |  |
|---------|----------------|----------------------|----------------------|----------------------|-----------------|-----------------------|-----------------------|-----------------------|
|  | iid-H (CIFAR10) | Dir (0.7)-H (CIFAR10) | Dir (0.2)-H (CIFAR10) | Dir (0.1)-H (CIFAR10) | iid-H (CIFAR100) | Dir (0.7)-H (CIFAR100) | Dir (0.2)-H (CIFAR100) | Dir (0.1)-H (CIFAR100) |
| 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.281           | 0.284                 | 0.298                 | 0.277                 |
| FedMDMI | **0.133**       | **0.122**            | **0.125**            | **0.117**            | **0.272**       | **0.262**             | **0.267**             | **0.259**             |

**Comment 1-Complexity analysis:** We did provide the complexity analysis in terms of the computation and memory for a single communication round. To calculate the total time consumption required to reach the convergence, we then have to multiply the computation time per round with the total number of communication rounds $T$ needed for convergence. In the following, we begin by providing a theoretical analysis of the total number of communication rounds required for convergence.

**_i) Theoretical convergence analysis._** For our proposed **FedMDMI**, as shown in Lines 230-232 of our manuscript, and leveraging the insights from [A3], we provide the convergence rate of our **FedMDMI** under the $\mu$-strongly convex objective as: $\mathcal{O}\left((1-\gamma \mu /8)^{KT} + {1/m} + d\right)$, where $m$ is the number of clients, $K$ is the number of local updates, $d$ is dimensions of the model, and $T$ is the number of communication rounds. Thus, the convergence rate scales linearly with the model dimension $d$. Reference [A4] shows that the convergence rate of **FedAvg** on the $\mu$-strongly convex objective is: $\mathcal{O}\left(\mu\exp({-\frac{\mu}{16(1+B^2)}T}) +{1/(mKT)}\right)$, where $B$ is the gradient dissimilarity. Notably, neither **FedPA** nor **FedEP** presented the convergence rate in their original papers.

**_ii) Empirical convergence behavior._** From the convergence curves illustrated in our manuscript (i.e., Figs. 3(c) and 3(d)), it can be empirically observed that the number of communication rounds required for convergence of our **FedMDMI** is less than that of **FedPA** and **FedEP**.

Then, we try to respond to the reviewer's concern on variational approximation (VA) methods being generally faster than sampling-based approaches in real-world problems. It is however observed in our rebuttal that FedEP, which uses the expectation propagation to obtain a VA of the global posterior, consumes more time per communication round and requires a greater number of communication rounds compared to our FedMDMI. The reasons are explained in the following.

**_i) Complexity comparison per communication round._** As stated in our first-round response to the common comment 1, **FedMDMI** only requires generating and computing the Gaussian noise at each round, while the sequence of global model $w_t$ converges to Gibbs posterior with sufficiently large $t$. In contrast, **FedEP** requires approximating the covariance as the inverse Hessian at each round, which introduces an additional $\mathcal{O}(d^3)$ time complexity and $\mathcal{O}(d^2)$ memory complexity.

**_ii) Total communication rounds._** In fact, convergence analysis for the federated variational approximation is still an open question, with no established results at present. So we try to give an intuitive explanation on why VA may converge slower than MCMC in FL. While VA may outperform sampling-based methods in terms of the convergence in centralized learning or real-world cases, the situation is different in FL. The main drawback of VA is that it typically results in biased posterior estimates for complex posterior distributions. This issue is further exacerbated in FL, where data heterogeneity across clients already contributes to biased local posteriors, making the use of VA even more challenging. Consequently, using VA in this context can result in an unstable and slow convergence of the aggregated global posterior.
Dear Reviewer 4nGu,

Thank you for the updated review. We have provided our response to each of these concerns in the following.

**Comment 1-Further Clarification on Fig. 3:** Regarding the plots in Figs. 3(a) and 3(b), as noted by the reviewer, the lines all lie very close together and very far from zero. This observation actually highlights a key challenge in the FL settings: the scarcity of data samples at each client increases the risk of local model overfitting, leading to overconfident decisions and poor uncertainty estimation in the aggregated global model. This also underscores the necessity of incorporating Bayesian inference into FL. As further illustrated in Table 3, the posterior-based approach demonstrates a significantly better performance in uncertainty estimation as compared to the point-based approach.

Besides, following this suggestion, we will also include the shaded error bars representing variance in all plots of Fig. 3 in the final version of our manuscript.

**Comment 2-Performance differences in Table 3:** Following this suggestion, we will include the variance of the ECE values in Table 3, as follows.

**Table 1: ECE (with  the variance of the ECE values) under various setting.**
| Method   | **CIFAR-10** |  | **CIFAR-10** | |
|----------|-------------|-------------|-------------|-------------|
| |Dir (0.2)-L | Dir (0.7)-L | Dir (0.2)-L | Dir (0.7)-L  |
| FedAvg   | 0.169 ± 0.0039 | 0.165 ± 0.0042 | 0.429 ± 0.0030 | 0.432 ± 0.0028 |
| FedM     | 0.159 ± 0.0025 | 0.169 ± 0.0036 | 0.468 ± 0.0027 | 0.459 ± 0.0035 |
| MimeLite | 0.182 ± 0.0041 | 0.178 ± 0.0034 | 0.461 ± 0.0029 | 0.470 ± 0.0022 |
| SCAFFOLD | 0.192 ± 0.0018 | 0.194 ± 0.0025 | 0.472 ± 0.0016 | 0.479 ± 0.0034 |
| FedBE    | 0.182 ± 0.0029 | 0.189 ± 0.0032 | 0.440 ± 0.0015 | 0.463 ± 0.0021 |
| FedPA    | 0.173 ± 0.0031 | 0.176 ± 0.0020 | 0.374 ± 0.0033 | 0.371 ± 0.0025 |
| FedEP    | 0.121 ± 0.0033 | **0.118 ± 0.0021** | 0.289 ± 0.0045 | 0.273 ± 0.0027 |
| FALD     | 0.135 ± 0.0028 | 0.127 ± 0.0023 | 0.267 ± 0.0018 | 0.269 ± 0.0023 |
| FedMDMI  | **0.115 ± 0.0019** | 0.120 ± 0.0031 | **0.261 ± 0.0023** | **0.263 ± 0.0029** |

**Comment 3- Degree of heterogeneity and effectiveness:** 
First, we would like to clarify that most of the Bayesian-based comparison methods, including FedBE, FedPA, and FedEP, are specifically designed to address the data heterogeneity issue in FL. We apologize for not clarifying this point earlier due to the word limit on the first-round response.

Besides, following the reviewer's suggestion, we have also conducted additional experiments (iid and Dir (0.1)), observing that as the degree of heterogeneity increases, the generalization performance of all the comparison algorithms declines. However, the generalization performance of our FedMDMI decreases to a 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. 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 2: Test Accuracy (\%) v.s. Data heterogeneity.**
| Method  | **CIFAR10** |   |   |   | **CIFAR100** |   |  |  |
|---------|----------------|----------------------|----------------------|----------------------|-----------------|-----------------------|-----------------------|-----------------------|
|  | iid-H (CIFAR10) | Dir (0.7)-H (CIFAR10) | Dir (0.2)-H (CIFAR10) | Dir (0.1)-H (CIFAR10) | iid-H (CIFAR100) | Dir (0.7)-H (CIFAR100) | Dir (0.2)-H (CIFAR100) | Dir (0.1)-H (CIFAR100) |
| 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                | 52.99           | 50.02                 | 49.08                 | 47.32                 |
| FedMDMI | **85.05**       | 83.76                | **83.56**            | **82.28**            | **53.28**       | **50.71**             | **49.70**             | **48.41**             |
Thank you for your detailed explanation of the figures and results. The new experiments and results are helpful for better understanding the empirical aspects of the contribution. Please see below for an updated review.

**Comment 1-Further Clarification on Fig.3:** I understand what the plots are representing. To me, it’s surprising that the lines all lie very close together and very far from zero (the optimal line). Adding shaded error bars representing variance to all plots in Fig. 3 would be very helpful. 

**Comment 2-Performance differences in Table 3:** Can you report the variance around those mean performances as well?

**Comment 3- Degree of heterogeneity and effectiveness:** The pattern you’re referring to is not consistent across all methods (e.g. FedEP). Since no other model is designed to handle data heterogeneity I’m very surprised that FedMDMI doesn’t outperform all other models by a large margin as the heterogeneity hyperparameter increases. Can you make a line plot of the ECE and top-1 accuracy as a function of heterogeneity to see if the line corresponding to FedMDMI diverges at some point?

**Comment 1-Complexity analysis:** As I understand, the computational and memory complexity is provided for a single round of SGLD. The main question is how long it takes for each method to converge. Depending on the problem this could result in better or worse convergence than optimization-based methods [1,2]. In general, it’s believed that for real-world problems variational methods are faster in high dimensions than sampling-based methods. Therefore I’m surprised that the timing results in the rebuttal are in favor of FedMDMI. Do the authors have any explanation for this?

**Comment 3-Comparison of MCMC and Variational Approximation (VA):** While the motivation is solid and sensible, ultimately it’s an empirical question whether variational methods or sampling-based methods perform best for FL with data heterogeneity. Is any of the compared methods based on a variational framework?

[1] Ma, Yi-An, et al. "Sampling can be faster than optimization." Proceedings of the National Academy of Sciences 116.42 (2019): 20881-20885.

[2] Kim, Kyurae, Yian Ma, and Jacob Gardner. "Linear Convergence of Black-Box Variational Inference: Should We Stick the Landing?." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.
Dear authors,

Thank you for your detailed response. 

I appreciate how you addressed the concerns I raised. Therefore, I am willing to raise my score from 5 to 6.

I suggest adding results from new experiments to the revised version, including the running time, to make your evaluation more complete.

It would also be helpful to mention relevant literature on Bayesian Federated Learning to give readers more context.
I thank the authors for the rebuttal and raise my score.