Federated Learning with Uncertainty and Personalization via Efficient Second-order Optimization
Authors: Shivam Pal
Aishwarya Gupta
Saqib Sarwar
Piyush Rai
Venue: Accepted by TMLR ()

Abstract: Federated Learning (FL) has emerged as a promising method to collaboratively learn from decentralized and heterogeneous data available at different clients without the requirement of data ever leaving the clients. Recent works on FL have advocated taking a Bayesian approach to FL as it offers a principled way to account for the model and predictive uncertainty by learning a posterior distribution for the client and/or server models. Moreover, Bayesian FL also naturally enables personalization in FL to handle data heterogeneity across the different clients by having each client learn its own distinct personalized model. In particular, the hierarchical Bayesian approach enables all the clients to learn their personalized models while also taking into account the commonalities via a prior distribution provided by the server. However, despite their promise, Bayesian approaches for FL can be computationally expensive and can have high communication costs as well because of the requirement of computing and sending the posterior distributions. We present a novel Bayesian FL method using an efficient second-order optimization approach, with a computational cost that is similar to first-order optimization methods like Adam, but also provides the various benefits of the Bayesian approach for FL (e.g., uncertainty, personalization), while also being significantly more efficient and accurate than SOTA Bayesian FL methods (both for standard as well as personalized FL settings). Our method achieves improved predictive accuracies as well as better uncertainty estimates as compared to the baselines which include both optimization based as well as Bayesian FL methods.

Published in Transactions on Machine Learning Research (12/2025)
Federated Learning with Uncertainty and Personalization via Efficient Second-order Optimization
Shivam Pal
pshivam@cse..itk.cm
Dept. of Computer Science and Engineering
II Kanpur
Aishwarya Gupta
uishwaryaagucse..lkkk
Dept. of Computer Science and Engineering IT Kanpur
Saqib Sarwar
sagib@cse..üle
Dept. of Computer Science and Engineering IIT Kanpur
Piyush Rai
piyush@cse. wit.seeeiiit......oo.coooooooo
Dept. of Computer Science and Engineering and Dept. of Intelligent Systems
11T Kanpur
Reviewed on OpenReview: https: // openreview. net/ form? id= TehCnGBK4F
Abstract
Federated Learning (FL) has emerged as a promising method to collaboratively learn from decentralized and heterogeneous data available at different clients without the requirement of data ever leaving the clients. Recent works on FL have advocated taking Bayesian ap- proach to FL it offers a principled way to account for the model and predictive uncertainty by learning posterior distribution for the client and/or server models. Moreover, Bayesian FL also naturally enables personalization in FL to handle data heterogeneity across the dif- ferent clients by having each client learn its own distinct personalized model. In particular, the hierarchical Bayesian approach enables all the clients to learn their personalized models while also taking into account the commonalitics via a prior distribution provided by the server. However, despite their promise, Bayesian approaches for FL can be computationally expensive and have high communication costs well because of the requirement of computing and sending the posterior distributions. We present novel Bayesian FL method using an efficient sccond-order optimization approach, with computational cost that is similar to first-order optimization methods like Adam, but also provides the various benefits of the Bayesian approach for FL (e.g., uncertainty, personalization), while also being sig- nificantly more efficient and accurate than SOTA Bayesian FL methods (both for standard as well personalized FL settings). Our method achieves improved predictive accuracies well as better uncertainty estimates compared the baselines which include both optimization based well as Bayesian FL methods.
1 Introduction
Federaled Learning (FL) (McMahan et al. 2017) aims at learning global model collaboratively across clients without compromising their privacy. It involves multiple client-server communication rounds, where in each round the selected clients send their local models (trained on their private datasot). to the server and the server aggregates the received models followed by its broadcasting to all clients. Thus, the global model, approximation to the model obtained if all the data was accessible, depends significantly both

---

Published in Transactions on Machine Learning Research (12/2025)
Server
p(0 D) = (w, diag(h))
D U Dx
h = Σ Tahe
P2
w = (Ek= Tehy o we) h
0: pointwise multiplicaation polut wise division
(w, h))
(w, A)
(wg, , hg)
(wo, h)
(w2, h2)
(wi, As)
Client 1
Client 2
Client K
(1), h1) = IVON(DI, h)
(w², h2) IVON(D,, w, h)
(wK, hx) IVON(DR: w, h)
p(0 D1) = N(w,, diag(hi) w1, h1 Rd
p(0 D2) = (ww,, (hag( w2, h2 C Rd
p(0 Dx) =s N(xx, diag(hx)) WKy hg R2
Figure Illustration of FedIvon.
the quality of the received clients' models and the chosen aggregation strategy at the server. As a result, a straightforward approach like FedAvg (McMahan et al., 2017) can yicld a high-performing global model if the data is i.i.d. distributed among clients but performs imally in case of non-i.i.d. data distribution. Moreover, the challenges are compounded if the individual clients' private dataset sizes are small.
The limitations of standard FL become even more apparent with data heterogeneity, where clients have distinct data distributions. single global model might fail to represent all clients well, leading to poor performance. This motivates personalized FL (pFL) Tan et al. (2022). which aims to adapt models to individual clients while leveraging shared global knowledge.
In such settings, learning the posterior distribution instead of a point estimate at each client results in enhanced performance and uncertainty mcasures, as demonstrated in several recent works, such (AL- Shedivat et al.] 2020 Liu et al., 2024 Bhatt et al. 2023 Guo et al. 2023) which have advocated taking a Bayesian approach to FL. Moreover, Bayesian FL is also natural for personalization because the server model can serve as a prior distribution in a hierarchical Bayesian framework, enabling easy personalization of client models using their respective client-specific likelihoods. However, existing Bayesian FL and pFL methods usually rely on running compulationally expensive routines on the clients (e.g., requiring expensive MCMC sampling (Al-Shedivat et al., 2020), expensive Laplace's approximation which requires Hessian computations and inversions (Liu et al. 2024) the clients, methods based on learning deep ensembles (Linsner et al., 2021)), as well expensive client-server communicatin (Kassab & Simeone, 2022) and aggregation at the server (note that, unlike standard FI., Bayesian FL, would require sending the whole client posterior to the server). Due to such computational bottlenecks and communication overhead, Bayesian approaches lack scalability, especially for clients with limited resources and bandwidth.
Thus, to bridge this gap, we propose a novel Bayesian FL algorithm FedIvon (with its high-level idea lustrated in Fig. [}, that balances the benefits of Bayesian inference, such as enhanced performance, and quantification of predictive uncertainty, with minimal increase in computational and communication over- head. In particular, 2024h to perform highly efficient variational inference (VT) on each client by approximating its local postcrior using a Gaussian with diagonal covariance. It uses the natural gradient to capture the geomelry of the loss function for faster convergence. Moreover, it computes the Hessian implicitly, making our method computa- tionally cheaper than other existing Bayesian FL and pFL methods that use explicit Hessian computation, e.g., Laplace's approximation (Liu et al. 2024), expensive MCMC sampling (Al-Shedivat et al. 2020 Bhatt
leverage the IVON (Improved Variational Online Newton) algorithm (Shen et al,
et al. 2023). even VT Kassab & Simeone 2022) at the clients. These local posteriors can be efficiently sent
2

---

Published in Transactions on Machine Learning Research (12/2025)
to the server and the global posterior can be computed for which also present local posterior aggregation strategies. Our main contributions are:
We introduce an efficient Bayesian FL method FedIvon that uses an efficient second-order opti- mization approach for variational inference whose per-round compute matches first-order optimizers (e.g., Adam) while delivering higher predictive accuracy and better uncertainty estimation than statc-of-the-art Bayesian and non-Baycsian FL methods.
Our method also supports client-level model personalization naturally by leveraging hierarchical Bayesian framework. Clients can use the server's posterior as prior to learn their private models, effectively balancing local adaptation with global knowledge sharing.
2
Related Work
FedAvg (McMahan et al. 2017), the foundational federated learning algorithm, approximates the global model as the weighted aggregation of locally trained client models, performing effectively with i.i.d. data distributions. Since then, numcrous sophisticated and efficient algorithms have been proposod to handle more realistic challenges such as non-i.i.d. data distribution, heterogeneous and resource-consttaainee clients, and multi- modal data as explored in recent works (Lin et al. 2020 Pfeiffer et al., 2023 Zhang et al 2023 Che et al. 2023 Liu et al. 2023). However, here, we will primarily restrict our discussion to Bayesian FI. and personalized FL algorithuss as they are the most relevant to our work.
Bayesian Federated Learning In federated learning, a key limitation of conventional point-estimate- based approaches is their susceptibility to overfitting in limited data settings and lack of predictive uncertainty estimates. To address this, Bayesian approaches have been advocated for federated learning, which in- volves the computation of clients' local posterior distribution followed by their aggregation at. the server to compute the global posterior distribution, offering enhanced performance and quantification of predictive uncertainty. Unfortunately computing full posterior distribution on resourco-constriing clients can be in- tractable, and also poses communication overheads due to the need of communicating the full posterior to the server. FedBE (Chen & Chao 2020) mitigates the commuicatiin overhoad by leveraging SWAG (Mad- dox et al.] 2019 to learn each client's 's posterior but communicating only its mean. The server then fits a Gaussian Dirichlet distribution to the clients' posterior mean and distills it into a single model to be communicated in the next round. However, the server-side aggregation in FedBE does not incorporate the client model uncertainty since the clients do not transmit the respective posteriors to the server. FedPA (AL- Shedival el al.) 2020) addresses this by learning Gaussian distribution for each client and computing the mean of the global posterior at the server via posterior averaging. However, it ultimately discards the global covariance as the server broadcasts only a point estimate in subsequent rounds, so the aggregated postcrior is not reused as a prior on the clients. In contrast, our method FedIvon maintain a diagonal Gaussian posterior at the server and explicitly sends both mean and (diagonal) precision back to the clients, which then used as the prior at the client for the next round. This preserves a fully hierarchical Bayesian update across communication rounds. Similarly, FedLaplace (Liu et al.) 2024) approximates each client's posterior Gaussian distribution, ion, modeling the global posterior as mixture of Gaussian, though further approximated by a single Gaussian by minimizing KL divergence. Moreover, dependence on Laplace's approximation makes FedLaplace heavy in terms of computation and communication costs.
Second-order Optimization for Federated Learning shows promise for improving convergence but is often limited by efficiency and communication overhcad. Methods such as FedNL (Safaryan et al.] 2021). which use privacy-preserving Hessian learning and compression, and second-order approaches incorporating global line search (Bischoff et al. 2021b), offer potential solutions to these challenges.
Personalized Federated Learning In the case of non-iid data distribution among clients, a single global model represents the average data distribution and diverges substantially from each client's local distribution. Consequently, the global model, though benefitted from collaborative learning, performs suboptimally for individual clients. Personalized federated learning addresses this by adapting a part or the whole model to the local data distribution explicitly. A typical approach is to split the model into two parts, a base model for global representation learning and a head model for personalized learning. FedPer (Achituve et al.] 2021
3

---

Published in Transactions on Machine Learning Research (12/2025)
and FedRep CColins et al. 2021) use this strategy, applying FedAvg for collaborative learning of the base model leveraged by the head for local data adaptation. Similarly, FedLG (Liang et al. 2020) splits the model into local and global components to learn local and shared representations, respectively. It shares the global parameters with the server while enhancing local parameters further using the unsupervised or self-supervisod approach. Avg (Fallah et al. 2020) applies a Aodol-Aguostic Meta-Learuing (MAML) (Finn et al. 2017] inspired framework to learn a shared model for faster adaptation to the client's data. pFedME (T Dinh et al. 2020) decouples personalized adaptation from shared learning by regularizing each client's loss function using Moreau envelopes. pFedBayes (Zhang et al., 2022) is Bayesian approach that aims at learning the personalized posterior distribution of each client. In each round, pFedBayes computes client's 's posterior using the global model as the prior and sends it to the server for updating the global model. pledVEMM (Zhu et al., 2023) also computes the client's posterior by restricting it to the Gaussian family. However, it leverages the collaborative knowledge of other clients by assuming conditional independence among clients' models given the global model.
3 Bayesian FL via Improved Variational Online Newton
The standard formulation of FL is similar to distributed optimization except some additional constraints, such no data-sharing among clients and server and a limited communication budget. Assuming K clients, let D UsE[K] D& be the total available data where D, denotes the private data of client k. The objective of standard FL is to solve 0" = arg ming ΣnE(K) - log p(Dx 0). However, this optimization problem is not trivial it requires access to each client's data which is not permitted in the federated setting. Thus, multi-round approach is usually taken where clients learn their local models, send these local models to central server which aggregates them into global model, and send the global model to the clients to continue the next round of learming.
Unlike standard FL, which only learns a point estimate of 8, an alternative is to take a Bayesian approach and learn the distribution of conditioned on data. The posterior distribution of can be wrillen as
D) x p(0) II P(D KE[R]
(1)
where p(0) is prior distribution I and (DDx 0) is the likelihood w.r.t. client k. Assuming uniform prior p(0), it can be trivially shown that optimizing the standard FL objective function is equivalent to finding the mode of the posterior p(0 | 'D,, i.e., 0" = arg maxo log p(g((0) (D).
Computing the full posterior p(0 DD) is more useful than computing just the point estimate 0 bocause the posterior helps take into account model uncertainty. However, is computationally intractable to compute the posterior exactly. Directly approximating p(0 (D) using approximate inference methods such as MCMC or variational inference (Angelino et al.) 2016) is also non-trivial, as it requires computing each client's likelihood which in turn requires global access to all the client's data.
Claim 1 Given common prior and conditional independence, the global posterior p(0 D) can be approx- imated at the server by the product of local client posteriors without requiring access to any client 's local data.
Proof of the claim ] is given in the appendix B If local posteriors p(0 Dx) are also being approximated, multiple rounds of optimization are needed to reduce the aggregation error in the global posterior (AI- Shedivat et al. 2020). In Bayesian FL., another challenge to make the computation of the local posteriors, their aggregation at the server, and the client-server communication, efficient, which in general can be difficult in general (Al-Shedival et. al. 2020).
4 Client's posterior approximation
1 Σ 4(0), and we Assuming client & has Ng training examples, its local loss can be defined as (x) = Ne compute the point estimate of the parameters 0% arg ming (0). However, in our Bayesian FL.
4

---

Published in Transactions on Machine Learning Research (12/2025)
setting, we will compute the (approximate) posterior distribution for each client using variational inference, which amounts to solving the following optimization problem
T (0) arg min (x(4)
(2)
(0)
(g) = Eg (0)
+ DKL (xxx()()))x())))
(3)
where P& (0) is the prior and Dx is the Beellback-Leible divergence. If we use the Gaussian variational m, ()) where m, and of denote the variational parameters that are to be optimized for. Optimizing the objective Equation w.r.t these variational parameters requires making the following updates
family for ()()) with diagonal covariance then (0) = (0
m n+1 m Lx(()
(4)
1+1 o -oVo.
(5)
k
where a > O is the learning rate. Computing exact gradients in the above update equations is difficult due to the expectation term in (x(g). A naïve way to optimize is to use stochastic gradient estimators. However, these approaches are not very scalable due to the high variance in the gradient estimates. [Shen et al. (2024) improved these update equations and provided much more efficient update equations similar to the Adam optimizer, which is essentially the improved variational online Newton (TVON) algorithm, with almost similar computational cost as Adam, and their key differences are summarized below.
Unlike Adam which solves for 0, IVON solves for both the mcan vector m and the variances o² which provides us estimate of the Gaussian variational approximation at each client. Note that the mean plays the role of 0 in Adam. In addition, the variances naturally provide the uncertainty estimates for 0, essential for Bayesian FI. (both in estimating the client models' uncertainties as well as during the aggregation of client models at the server).
Unlike Adam which uses squared minibatch gradients to adjust the learning rates in different dimen- sions, IVON uses reparametrization defined gradient element wise multiplied by (0- m)/o² to get an unbiased estimate of the (diagonal) Hessian. Using this, IVON is able to get cheap estimate of the Hessian, which makes it a second-order method, unlike Adam.
IVON offers the significant advantage of providing estimate of second-order information h with minimal computational overhead. The Hessian h corresponds to the inverse of o², where o² x K. An estimate of h is accessible throughout the training process (see Algorithm [ Moreover, there is no explicit update equation for h. It is computed implicitly using gradient information. In comparison, standard optimization mothods such as SGD, Adam, and SWAG require additional effort estimate second-order information.
5 Posterior aggregation at server
At the server, we can aggregate the client posteriors to compute the global posterior (Fischer et al. 2024) IVON approximates clients' posteriors as Gaussians and product of Gaussian distributions still Gaussian distribution up multiplicative constant. Thus we approximate the global distribution Gaussian whose optimal mean and covariance matrix expressions are given below. Moreover, since each client's variational approximation is Gaussian with diagonal covariance matrix, it makes the aggregation operations efficient. Let's s assume (()) Dx) = N( p4, H where μ4 = m, and H diag(hæ) = diag(1/0% Using results of the product of Gaussians (Liu et al. 2024 Tischer et al. 2024), our aggregation is defined as
log ( () ~
Wh log q Dx))
(6)
1=1
where we = No
q(0) D) R N(O 14, H-1), with H
H and p= _K wi Hap
Σx No
5

---

Published in Transactions on Machine Learning Research (12/2025)
Algorithm Client Update
Input: Local dataset D, model weights m, Hhss local ooeal_epochs(ee (E), learning rates fae), weight decay
8, hyperparaneters B1, 82, batch-size B
Output: Trained model weights m, Hessian o
1: g - 0, A A
DI: n == E
- 1//A(h +).
de dd ae for all e (1. n2.
for c 1 to E do
Sample batch of inputs of size B from D.
6:
g - V((O), where 0 ~ qe-1
b where qe = =N(0 m. diag(o?))
h g (0- m)/o²
g Łig + (1 - 31)
h 8yh (1 32)h + 2
1
-1-122 (h (1)
10:
g g/ 3
11:
0. g htmm/h// 6h
12:
- 1/ A(h 6
13:
end for
Algorithm Fedlvon Algorithm
Input:
Total communication rounds R. total clients K, clients' private datasets (De) weight mo, initial model Hessian ho
initial model
1: for == to R do
Broadcast my h to all K clients Randomly sample & clients
Update selected client models locally
for i- 1 to & do
(m; h.) < CLIENT UPDATE(D; mr, hr)
6: end for
7: Initialize m841 + 0, hr41 6-
Aggregation of client models at server wli] is relative dataset size |D:l/DD
for -ni to do
hr41 t- hr+1 + h4 wlil
10:
M,+1 + mr+1 + (m; @ h,) x wii
11: end for
12: m + mr41/h141
a elementwise division; global weight and Hessian
13: end for
Output: Global model weights and Hessian (m hn)
Other aggregation strategies are also possible Fischer et al. 2024 and we leave this for future work. Note that our aggregation strategy can also be seen as Fisher- weightt model merging (Daheim et al. 2023) where each client model is represented as the mean weights m& and Fisher matrix which depends on local posterior's variances of (although model merging only computes the mean, not the covariance, and thus does not yield global posterior distribution at the server).
The appendix provides further details of IVON and its integration in our Bayesian FL setup.
Notably, Fedlvon is appealing from two perspectives. It can be viewed an an efficient Bayesian FL algorithm offering the various benefits of the Bayesian approach, as well as a federated learning algorithm that easily incorporates second-order information during the training of the client models, while not incurring the usual overheads of second-order methods used by some FL algorithms (Bischoff et. al. 2021a).

---

Published in Transactions on Machine Learning Research (12/2025)
Fedivon
FedLaplace
FedAvg
Test Loss
2.0
2.0
2
1.5
0.5
CIFAR-10
EMNIST
SVHN
Test Accuracy
0.6
0.75
0.75
0.4
0.50
0.25
0.2
0,00
250
500
750
1000
200
400
600
100
200 Round
300
400
Round
Round
Figure 2: Loss and Accuracy of various methods vs rounds for the EMNIST, SVHN, and CTFAR-10 datasets.
6 Personalized Federated Learning
Personalized FL in FedIvon can be achieved straightforwardly. Similar to equation g the personalized loss function for each client is defined as,
(x(g) (0) (000 (0)
(7)
Where B 2 O controls the level of personalization. The term (o) represents the prior distribution for client k. During each communication round, the posterior distribution from the server can be used as the prior pk (0) for the client. This setup enables clients to adapt the global model according to their local data characteristics while leveraging information from the global model.
When B = 0, the model becomes fully personalized, relying solely the client's data without influence from the prior (i.e., no information from the server). Conversely, a higher value of / incorporates more knowledge from the global server model into the client's learning process, balancing between personalization and shared global information. This framework provides a flexible mochanism to adapt client models according to their individual data while still benefiting from collective learning through the shared server posterior. We fixed B 1 in all our pFL experiments.
7 Experiments: Standard FL
We experiment on three publicly available datasets: EMNIST (Cohen al. 2017), SVHN (Netzer et al. 2011) and CIFAR-10 (Krizhevsky] 2009). EMNIST consists of 28x28 grayscale images of alphabets and digits (0-9) with train and test split comprising 124800 and 20800 images respectively; however, in our experiments, we restrict to alphabets only. SVHN consists of 32x32 RGB images of house number plates categorized into 10 distinct classes, each corresponding to one of the ten digits. It has a train and test split of size 73252 and 26032 respectively. CIFAR-10 comprises 32x32 RGB images of objects classified into classes with 50000 training images and 10000 test imagos.
In our experiments, We use Adam optimizer with Saariing_rat==lee and FedLaplace method. IVON (Shen et al.) 2024 optimizer is used for FedIvon with different hyperparam- eters given in Table 0 Linearly decaying learning rate is used in all the experiments.
weight weight_decay=2er for FedAvg
We evaluate Fedlvon in a challenging and realistic scenario involving heterogeneous data distribution among a large number of clients with each client having very few training examples. For each experiment, we consider a total of 200 clients with each client having a small private training set of less than 100 examples. To simulate non-i.idd data distribution, we randomly sample inputs from the training split, partition the sampled inputs into shards, and distribute shards among clients to create imbalanne training data

---

Published in Transactions on Machine Learning Research (12/2025)
similar to (Chen & Chao 2020 For fair comparison, use the same non-ii.. data split across clients for all the baseline methods and FedIvon. We follow the experimental setup of (Bhatt et al. 2023 and train customized CNN models on EMNIST, SVHN, and CIFAR-10 datasets. We compare our proposed method FedIvon with FedAvg (McMahan et al. 2017) (simple aggregation of client models at server) and FedLaplace (Liu et al., 2024) (using the Laplace's 's approximation to fit a Gaussian distribution to each client's local model followed by aggregation at the server). FodAvg serves as a baseline to emphasize the importance of uncertainty quantification without compromising on the performance while FedLaplace serves competitive baseline to evaluate FedIvon's predictive uncertainty measures. For all the baselines and FedIvon, we run the federated algorithm for 2000 communication rounds, selecting a randomly sampled 5% i.e, 10 clients per round. We train each client's model locally for 2 epochs using a batch size of 32. We provide further details on hyperparameeerss model architectures, and split in the appendix.
7.1 Classification Task
We train a classification model in FL setting using all the methods and report the results in Table [ We evaluate all trained models' performance (accuracy and negative log-likelihood) on the test split and use metrics such as Expected Calibration Error (ECE) and Brier score to quantify predictive uncertainty. In our results, FedIvon@mme @mean denotes point estimate based predictions evaluated at the mean of IVON posterior and FedIvon corresponds to Monte Carlo averaging with 500 samples from the posterior.
Metric
Dataset
FodAvg
FedLaplace
FedIvon &mmean
FedIvon
EMNIST
91.66
91.33
93.14
93.09
ACC(1)
CIFAR-10
62.25
62.92
62.54
SVHN
82.14
81.99
84.54
84.76
EMNIST
0.0405
0.0381
0.0349
0.0188
BCE(L)
CIFAR-10
0.0981
0.1072
0.0983
0.0312
SVHN
0.0311
0.0241
0.0148
EMNIST 0.3355
0.3255
0.2821
0.2341
NLL(1)
CIFAR-10
1.199
1,233
1,1500
1.0790
SVHN
0.6857
0.6423
0.5624
0.5303
EMNIST
0.1303
0.1314
0.1075
0.1019
BS(L)
CIFAR-10
0.5191
0.5284
0.5114
0.5021
SVHN
0.2640
0.2627
0.2256
0.2210
Table 1: Test accuracy (ACC), Expected Calib. Error (ECE), Negative Log Likelihood (NLL), and Brier Score (BS)
As shown in Table 1 FedIvon outperforms all the basclines and yiclds the best test performance and calibra- tion scores. FedIvon leverages the Improved Variational Online Newton (IVON) method to approximate the Hessian by continuous updates throughout the training. We also show the convergence of all the methods on all the datasets in Figure 2 As observed, Fedlvon exhibits slightly slower improvements in the early training phase as compared to other basclincs but soon outperforms them owing to its improved Hessian approximation as training progresses. Moreover, unlike FedLaplace which fits Gaussian distributtin to the client's model using Laplace approximation evaluated at MAP estimate, Fedlvon approximates the Hessian over the entire course of its training, resulting in much better predictive uncertainty estimates. Figure 3 shows the reliability diagrams for FedAvg, FedLaplace, and FedIvon on the CIFAR-10 10 and EMNIST datasets. This clearly shows that FedIvon yields better-calibrater predictions than both FedAvg and Fedlaplace. For other dataset reliability diagrams are shown in the appendix B
7.2 Out-of-Distribution Detection Task
Models
EMNIST
CIFAR-10
SVHM
Predictive uncertainty of the model plays a crucial role in uncertainty-driven tasks such as OOD detection and active learn- ing. We evaluate FedIvon and the baselines for distinguishing OOD inputs from in-distribution inputs using their predictive uncertainty. Given any input x, the predictive uncertainty of the model's output is given by its Shannon entropy and is used to filter OOD inputs.
FedAvg
0.8910
0.7896
0.7975
FedLaplace
0.8297
0.7513
0.8222
FedIvvn
0.9032
0.7662
0.8233
Table 2: AUROC (t) score for
DOD/in-domain data detection
8

---

Published in Transactions on Machine Learning Research (12/2025)
FedAug with CIFAN-10
Fediaplace with CIFAR-10
Fedivon with CIPAR-10
3.0
Emptria
Empriial Cup
Erptical
Gap
0.6
Trequency
0.5
09-
Precuency
0.5
0.8
Trqquency
08
Dreang
2.4
0A-
0.4
0.4
0.4
KCS 6.672
SCE 0.515
0.0
0.2
0.0
Corfidence
Figure 3: Reliability diagrams for CIFAR-10 experiments (left: FedAvg, center: Fedlaplace, right: Fedlvon).
We simulate this task by randomly sampling 5000 images from the OOD dataset and mixing it with an equal number of randomly sampled inputs from the test split of the training dataset.
Specifically, we use EMNIST, CIFAR-10, and SVHN as the OOD dataset for the models trained on EMNIST, SVHN, and CTFAR-10 respectively. We report the AUROC (area under the ROC curve) metric for all the methods on all the datasets in Table which shows that FedIvon achieves better or competitive AUROC scores compared to the other baselines.
7.3 Effect of the number of local epochs.
We empirically investigate the impact of varying the number of local epochs E on the convergence behaviour of the different methods the server. Figure shows the convergence trajectories for several values of E. When E 1, Fedlvon converges more slowly than FedAvg, and FedLaplace converges FedIvon, this slower initial progress can be attributed to the fact that gradients are computed using stochastic weight samples; at initialisation, these samples yield noisier gradient estimates, which weakens the quality of early updates. FedTaplace, which requires the computation of a MAP estimate, also suffers at = 1: a single epoch of training produces a suboptimal MAP point and therefore leads to slower improvement. When the number of local epochs increased to E 2, all methods show faster convergence compared with the F = 1 setting. Additional local iterations provide better gradient estimates for FedIvon and a more accurate MAP estimate for FedLaplace, leading to improved server-side progress. After small number of communication rounds, FedIvon surpasses both FedAvg and FedLaplace, reflecting its ability to refine its variational updates as the global posterior becomes more informative.
more slowly. For
However, incrcasing the number of local epochs beyond = 2 leads to unstable training behaviour for FedIvon (and, more generally, for federated posterior learning algorithms which the server posterior is reused as prior for the clients in subsequent rounds). Large values of make each client's posterior highly concentrated (see Algorithm [ sceing the same data multiple times naturally increases posterior precision). When these high-precision client posteriors sent back to the clients as the prior for the next round, this sharp prior can dominate the local likelihood. In this regime, the client update behaves as if it were heavily regularized around the global mean, which limits the clicnt's ability to learn from its own data. As this prior- likelliid imbalance is repeatedly amplified across rounds, the global posterior becomes progressively sharper and the overall training dynamics can become unstable, particularly under highly non-i.i.d. conditions. For this reason, we restrict the number of local epochs to E = 2 for all the experiments, which provides balance between improved local optimization and stable global aggregation.
aggregated at the server and the resulting global posterior is
8 Communication and computational cost
Let P denote the number of model parameters. In each local step, FedIvon draws single Monte Carlo sample 0 m and performs one forward and one backward pass to obtain stochastic gradient.
9

---

Published in Transactions on Machine Learning Research (12/2025)
The subsequent IVON updates (momentum, diagonal curvature h, and inaanvarrinne updates) require only elementwise (P) operations. Thus, the pinimathhatch computational cost. is Credtvon B)) Civ + Cbw + O(P), which is the same asymptotic cost as FedAvg/ Adam. The difference between FedIvon and FedAvg is only by a constant factor: FedIvon performs the same -orward-backward pass but includes several additional ((P) elementwise updales.
practice, this overhead is small because the forward back pass dominates E (Ctw Chw + +O()), the runtime. For client k, the cost per communication round is F round =B identical in scaling to FedAvg and differing only by the small constant factor above.
Communication per round is lincar in model size. Each client uploads its posterior mean and diagonal precision (two vectors in IRE) and downloads the corresponding global quantities, for a total of 4P scalars per round. This is a factor of two more than FedAvg, which transmits only a single parameter vector in each direction. In typical federated-learning settings, where communication budgets are dominated by transmitting full model weights, doubling the number of transmitted vectors (from P to 2P) does not affect feasibility, the order of communication remains (P)).
Table 3 reports the empirical communicatin and computation costs per communication round for all meth- denotes the total amount of data transmitted between server and clients (up- link + downlink) in megabytes, averaged over rounds and participating clients. FLOPs/round is the average number of floating-point operations per round (estimated via PyTorch'' profiler), and Runtime/round the corresponding average wall-clock time in seconds.
ods. Communication/ooni
As expected, FedAvg is the most efficient baseline in all three met- rics, since it only communicates model weights and performs a sin- gle forward-backwar FedLaplace and Fedlvon incur roughly twice the communication cost of non-Bayesian method (FodAvg), because they must transmit both the mean parameters and a diagonal covariance between server and clients. In terms of FLOPs and wall-clock runtime, FedLaplace is the most expensive due to the additional passes required to form a Laplace approximation, whereas FedIvon remains close to FedAvg. Overall, these results indicate that FedIvon achieves Bayesian un- certainty quantification with a runtime overhead that is comparable to standard first-order federated optimizers, at the cost of twice increase in communication.
CC (MB)
FL.PPs (x (x10")
Runtime (s)
Method
pass per batch. Among the Bayesian hods,
FedAvg
9.7
8.31
12.91
FedLaplace FedIvon
19.4
16.71
19.4
8.48
15.05
Table 3: Per-round communication cost (CC), computational cost (FLOPs), and wall-clock runtime for different federated learning algorithms. Values are aver- aged over global rounds and participat- ing clients. The model and dataset were fixed across methods, so the numbers are best interpreted relatively.
9 Experiments: Personalized FL
For personalized FL experiments, focus two types of data heterogencity the clients similar to
Zhu et. al. 2202)) for classification task. We compare our approach FedIvon against personalized federated
baselines (PFedME
T Dinh et al. 2020 PFedBayes (Zhang et al. 2022), and pFedVEM (Zhu et al. 2023)
Class distribution skew: In class distribution skew, clients have data from only a limited set of classos. To simulate this, we use the CIFAR-10 datasct and assign cach client data from a random selection of 5 out of the classes.
Class concept drift: To simulate class concept drift, we use the CIFAR-100 dataset, which includes 20 superclasses, each containing 5 subclasses. For each client, we randomly select one subclass from each superclass (1 out of 5). The client's local data is then drawn exclusively from these selected subclasses, creating a shift in label concepts across clients. We define the classification task as predicting the 20 superclass.
To model data-quantity disparity, we randomly divide the training set into partitions of varying sizes by uniformly sampling slice indices, then assign each partition to different client.
10

---

Published in Transactions on Machine Learning Research (12/2025)
Method
Clients
Model
Local
PRedME
PFyd
PFedVEM
Rmeannn&mean
FedIvon
PM GM
56.9
72.3 0.1
71.4 1:
73.2 0.2
74.4 0.3
0.4
56.6 1.0
52.0 + 1.0
56.0 0.4
67.1 + 1.0
67.8 + 1.6
PM
52.1 d: 0.1
71.4 0.2
68.5
71.9
71.7 0.3
100
GM
60.1 1. 0.3
53.2 d: 0.7
60.1 0.2
68.4 0.2
69.2 0.2
PM
46.6 0.1
68.5 0.2
64.6
70.1 0.3
00.7..0.7
70.8 0.4
200
GM
58.7 1: 0.2
51.4 d: 0.3
59.4 0.3
68.2 1: 0.3
68.7 0.3
Table 4: Comparison Personalized FL Methods on CIFAR-10
Method
Clients
Model
Local
PPedME
pFedBayes
pFedVEM
FedIvon& mean
Fedlvon
PM
34.3 + 0.2
52.5 + 0.5
49.6 0.3
61.0 0.4
65.4
66.7 + 0.8 63.8 - 0.7
GM
47.9 + 0.5
42.5 + 0.5
52.8 : 0.4
63.2 1 0.5
PM
27.6 0.3
47.6 0.5
46.0 0.2
56.2 0.4
63.2 0.5
63.5 0.6
100
GM
45.1 #: 0.3
41.3
52.3 = 0.4
62.1 0.5
62.1 0.6
PM
22.2
41.6 1.8
40.1 0.3
51.1 0.6
56.1 + 0.6
56.5 0.5
200
GM
41.5 1.6
37.4 0.3
19.2 0.5
55.5 - 0.6
55.7 + 0.7
Table 5: Comparison of Personalized FL. Methods on CTFAR100
9.1 Setup
We evaluate approach in different settings: number of clients K {50, 100, 200). We followed the same model architectures as the prior work (Zhu et al. 2023 A simple 2-convolution layered-based model is used for CIFAR-10, while a deeper model having 6 convolution layers is used for the CIFAR-100 dataset. We assess both a personalized model (PM) and a global model (GM) at the server. The PMs are evaluated using test data that matches the labels (for label distribution skew) or subclasses (for label concept drift) specific to each client, while the GM is evaluated on the intire test set. All experiments are repeated 3 times, using the same of 3 random seeds for data generation, parameter initialization, and client sampling. For each client, the Ivon hyperparameters used are given in Table 7
9.2 Results
Table 4 and 5 presents results on CIFAR-10 and CIFAR-100 datascts respectively, which are used to simulato different types of data heterogeneity federated learning. Here local method refers to client only training without communication to server. CIFAR-10 represents class distribution skew, where each client has data from a limited of classes, while CIFAR-100 100 represents class concept drift, where each client has data from distinct subclasses within superclasses. For both datasets, evaluate client's average accuracy (personalized model) and server accuracy (global model) across varying client counts (50, 100, and 200). FedIvon uses 64 Monte Carlo samples to perform Monte Carlo averaging. On the other hand, sedlvon@@ean uses a point estimate using mode of the posterior.
On CIFAR-10, FedTvon achieves similar client accuracy pFedVEM, indicating both methods perform well under class distribution skew for individual clients. However, server accuracy, FodIvon shows a notable improvement over pFedVEM and other methods, highlighting FedIvon's strength in aggregating data from heterogeneous clients into accurate global model.
On CIFAR-100, which represents class concept drift, FodIvon demonstrates significant improvements over all other moth in both client's average accuracy and server accuracy. This performance advantage in both personalized and global evaluations suggests that FedIvon is well-suited to handling concept drift, achieving higher accuracy for individual clients and in the global model. Overall, FodIvon consistently outperforms other methods in both accuracy metries on CIFAR-10 and CIFAR-100 underscoring its robustness across different data heterogeneity scenarios.
11

---

Published in Transactions on Machine Learning Research (12/2025)
10 Limitations
One limitation of FodIvon is that approximates the global posterior with a diagonal Gaussian. If the true global posterior is not well approximated by single diagonal Gaussian, for example, if it exhibits strong parameter correlations, pronounced non-Gaussiaiity, or multimodal structure then this variational family cannot faithfully capture the geometry of the posterior. In such cases, the resulting approximation may distort epistemic uncertainty by misallocating variance across parameters and ignoring important covariance structure, which can in turn lead to miscalibrated predictive uncertainties, especially in regimes with strong client disagreement or significant distribution shift.
FedIvon also introduces additional resource overhead compared to FedArg. On the client side, IVON requires storing extra state (Hessian), which increases memory by small constant factor in our experiments but would be more substantial for very large models. Communication per round is roughly doubled, since each client sends and receives both mean and diagonal Hessian vectors instead of a single parameter vector. Although this overhead remains practical for the moderate-size CNNs and vision benchmarks considered here it would be more significant in largo-scalc scttings such as LLM training.
11 Conclusion
We presented a new Bayesian Federated Learning (FL) method that reduces the computational and commu- nication overhead typically associated with Bayesian approaches. Our method uses an efficient second-order optimization technique for variational inference, achieving computational efficiency similar to first-order methods like Adam while still providing the benefits of Bayosian FL., such uncertainty estimation and model personalization. We showed that our approach improves predictive accuracy and uncertainty estimates compared to both Bayesian and non- -on-Bayesian FL methods. Additionally, our method naturally supports personalized FL by allowing clients to use the server's posterior as a prior for learning their own models.
Acknowledgment P.R. acknowledges support from Google Research and Satya & Rao Remala Foundation. SP and SS acknowledge support from Visvesvaraya fellowship.
Broader Impact
FedIvon is a Bayesian personalized federated learning method, and its deployment raises several considera- tions that briefly discuss here.
Fairness across clients. By design, personalized federated learning can yield helerogeneous model qual- ity across clients. Clients with more or higher-quality data may obtain substantially better personalized posteriors than those with limited or noisy data. In scttings where model performance has direct socictal or economic consequences, this may amplify existing inequalities between users or institutions. FedIvon does not explicitly address such fairness concerns; incorporating fairness-aware objectives or constraints into Bayesian personalized federated learning is an important direction for future work.
Reliability of uncertainty estimates. A key motivation for FedIvon is improved uncertainty calibration, and we observe benefits on the benchmarks considered. However, miscalibrated uncertainty can be harmful in high-stakes applications if used to drive decisions. In particular, the diagonal posterior approximation and limited per-client data may lead to overconfidence on out-of diw-----iittbtttinnn inputs or on minority groups under- represented the participating clients. Our experiments are restricted to vision benchmarks in rolled settings, and we do not claim that FodIvon provides robust uncortainty guarantecs safety-critical domains. Stress-testing under distribution shift and demographic imbalance remains necessary before deployment in such contexts.
Privacy and security. Although federated lcarning keeps raw data on-device, transmitting model updates or posterior parameters can still leak information about local data. FedIvon shares client-specific posterior means and variances, which may, in principle, be exploited by a malicious server or adversary to infer properties of client data. We do not incorporate differential privacy, sccure aggregation, or other formal privacy mechanisms, and therefore do not claim any formal privacy guarantees. Combining FodIvon with established privacy-preserving techniques is an important prerequisite for in sensitive applications.
12

---

Published in Transactions on Machine Learning Research (12/2025)
References
Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advonces in Neural Information Processing Systems,
2021.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference Learning Representations, 2020.
Elaine Angelino, Matthew James Johnson, Ryan P Adams, et al. Patterns of scalable bayesian inference. Foundations and Trends@ in Machine Learning, 9(2-3)1119 247, 2016.
Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Federated learning with uncertainty via distilled predictive distributions, 2023. URL https://arxiv.org/ htwpgg/ab//222..ooo 07562
Sebastian Bischoff, Stephan Günnemann, Martin Jaggi, and Sebastian U Stich. On second-orde optimization methods for federated learning. arXiv preprint arXiv:2109.02: 02388, 2021a.
Sebastian Bischoff, Stephan Günnemann, Martin Jaggi, and Sebastian U. Stich. On second-order optimiza- tion methods for federated learning, 2021b. URL httptss///arww..orr/aaaasss// wttp:orrggaa///////// 02388
Christopher M Bishop. Pattern recognition and machine learning, volume 4. Springer, 2006.
Liwei Che, Jiaqi Wang, Yao Zhou, and Fenglong Ma. Multimodal federated learning: A survey. Sensors
(Basel, Swilzerland), 23, 2023. URL https://api. http://pp.s.mmann.e www.ssemmiiiseee cws.maatiiccccla.......mm. www.Corpss///..ccc...o.. 260693566
Hong- You Chen and Wei-Lun Chao. Feddistill: Making bayesian model ensemble applicable to federated
learning. CoRR, Abs/2009.01997 2020. URL
Gregory Cohen, Saeed Afshar, Jonathan C. Tapson, and André van Schaik. Emnist: an extension of mnist to
handwritten letters. ArXiv, abs/1702.05373, 2017. URL httpt::ttppppp.../ http:p///.ap..eemaiicccccooo... ctw.oooarrr.gg/rppusssi..
12507257
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International conference on machine learning, pp. 2089- 2099- PMT.R, 2021.
Nico Daheim, Thomas Möllenhoff, Edoardo Maria Ponti, Iryna Gurevych, and Mohammad Emtiyaz Khan. Model merging by uncertainty-basee gradient matching. arXiv preprint arXiv:2310. 128008 2023.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantoes: model-agnostic meta-learning approach. Advances in neural information processing systems, 53:3557-35663 2020.
Chelsoa Finn, Picter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126 1135. PMLR, 2017.
John Fischer, Marko Orescanin, Justin Loomis, and Patrick McClure. Federated bayesian deep learning: The application of statistical aggregation methods bayesian models, arXiv preprint orXiv:2403.1550 2024.
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning variational inference: scalable expectation propagation approach. In The Eleventh International Conference Learning Representations, 2023.
Rahif Kassab and Osvaldo Simcone. Federated generalized bayesian learning via distributed stein variational gradient descent. IEFE Transactions Signal Processing, 70:2180 2192, 2022.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Rayesian deep learning by weight-perturbation in Adam. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings mlr.
of Machine Learning Research, pp. 2611-2620. PMLR, 10-15 Jul 2018. URL https:t/pp/::eeeeeeee
hatps//v00/khan8aa. html
Mohammad Emtiyaz
Khan and Wu Lin. Conjugale-compulatiin variational inference variational inference in non-conjugate models to inferences in conjugate models, 2017.
Converting URL https:
7/arxiv. wwrggabs/1703..oo 04265
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. URL https://apl
semanticscholast senaatttccccoll.....r///ooooopDD
Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Aucrbach, David Brent, Ruslan Salakhut- dinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv e-prints, pp. arXiv- 2001, 2020.
Tao Lin, Lingjng Kong, Sebastian Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
Florian Linsner, Linara Adilova, Sina Däubener, Michael Kamp, and Asja Fischer. Approaches to uncer- tainty quantification in federated deep learning. In ECML PKDD Workshop on Parallel, Distributed, and Federated Learning, pp. 125-145. Springer, 2021.
Bingyan Liu, Nuoyan Lv, Yuanchun Guo, and Yawen Li. Recent advances federated learning: A system- alic survey. Neurocompuling, 597:128019, 2023. URL http:///pp..emmanicccc.oooooo/// corgocooooaarooooorrross.ss 255415857
Liangxi Liu, Xi Jiang, Feng Zheng, Hong Chen, Guo-Jun Qi, Heng Huang, and Ling Shao. A bayesian federated learning framework with online laplace approximation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(1):1 16, January 2024. ISSN 1939-3539. doi: 10.1109/tpamn 2003mi.202333222744 URL http:///x:: http://xxxdoo..oo org.o..o/ooooooooooo org/ 10. 1119/TPAMA.T 2023. 33227743
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicttion-cffcti learning of deep networks from decentralized data. In Artificial intelligence and slatistics, pp.
PMLR, 2017.
Yuval Netzer, Tao Wang, Adam Coates, A. Bissacco, Bo Wu, and A. Ng. Reading digits in natural im- ages with unsupervised feature learning. 2011. URL https://api semanticscholar wwwpolllarrrr//oooeeoe org/CorpusID;
16852518
Kilian Pfeiffer, Martin Rapp, Ramin Khalili, and Jörg Henkel. Federated learning for computationally constrained heterogencous devices: A survey. ACM Computing Surveys, 55:1 - 27, 2023. URL https:
//api. / seaaaniccmntcc Dotpa//iiiiicco...ooororoooooooa 258590978
M. H. Salaryan, Rustem Islamov, Xun Qian, and Peter Richtárik. Fedul: Making newlon-type mot
applicable to federated learning. ArXiv, hbbs/2206.02960 2021. URL stte::://ap..smmaaiicchht
Pog/CorppsID:23535555 235358296
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Tryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Möllenhoff. Variational learning is effective for large deep networks, 2024. URL https://aaxii.... abwpg//bs/240....m 17641
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in neural information processing systems, 2221994-21400 2020.
14

---

Published in Transactions on Machine Learning Research (12/2025)
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. 1EEE transactions on neural networks and learning systems, 34(12):9587-9600 2022.
Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized federated learning via variational bayesian inference. In International Conference on Machine Learning, pp. 26293-26310 PMLR, 2022.
Yifei Zhang, Dun Zeng, Jinglong Luo, Zenglin Xu, and Irwin King. A survey of trustworthy federated learning with perspectives on security, robustness and privacy. Companion Proceedings of the ACM Web Conference 2028, 2023. URL https://apl semanticscholar D 200 257000689
Junyi Zhu, Xingchen Ma, and Matthew Blaschko. Confidence-aware personalized federated learning via variational expectation maximization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2023.
A Theoretical insight and stability considerations
We do not provide full convergence proof for Fedlvon in the federated, non-convex setting. Instead, we sketch simple conditions under which the IVON-based client updates and the FedIvon aggregation lead to qualitatively stable fixed-point iteration. Single-client IVON Algorithm
maint ains diagonal-Gaussian posterior
(:()) (m,, A,
A₁ A (h+ 5),
the damping used in the update of where À [D], he the diagonal Hcssian estimate and At each step,
0: m
9. t- VE(O,),
hu - 91 1
of
and h, is updated as:
h4t1 (1 (1- phhe p he P (h+ n2,22
2 he so/
for some (0, and 80 > Shen et al. all t. Assuming the curvature statistics are bounded, lhul < Hmax, the EMA implics
2024 show that if ho > 0, then this update guarantees h; > 0 for
hmin S he S hmax oc.
With 8 O, the diagonal precision
A₂ = A(he + 6)
satisfies
A6 A. < A(hmax + 8) Vt,
so Ar is uniformly positive and bounded. The mean update is
m me a4 9r + 8m
h4
where G is the bias-corrected gradient momentum. This can be written as a preconditioned gradient step
Mi-1 Ou. A ghs
for some effective gradient of. Because A- is also bounded,
0 < X(hmax + 8) S SAT < < 0o.
A8
15

---

Published in Transactions on Machine Learning Research (12/2025)
Under standard smoothness assumptions and small dg, this yields stable preconditioned SGD iteration for the single-client IVON update. FedIvon aggregation with multiple clients. In Fedlvon, client returns
(x(0) N A
As = > + 8):
with the same IVON bounds 0 < Amin S h4 hmax. Therefore
d S Ak S (hmmax + +0) Vk.
The server aggregates postcriors natural-paramoton space:
K A
K Wk AR mi
(0) server
m (t) server
(ever))))
)-1
where Wi = =Dan//E,,,, |Dyl and K 2 2. By convexity we obtain uniform bounds
Amin == min(Aedd <
(6) server S ma + +5) =: Amax:
for all communication rounds t. Hence the global precision is also strictly positive and bounded, and the global covariance never degencrates explodes. In the next round, each client uses Qserver as its prior and applies a finite number IVON steps. With a small number of local epochs E, each client performs only a single approximate Bayesian update per round, so that all Aş remain within the same bounded interval. The combined client server mapping on (m, A) is therefore a stable fixed-point iteration over a compact subset of the Gaagonal-Gaussann family. A full convergence analysis for non-iid, multi-client, non-convex models is beyond the scope of this work, but these boundedness properties explain why FedIvon remains numerically stable in our experiments.
B Proof of claim 1
Let K clients hold disjoint datasets Di,..., DK, and assume common prior p(0). Define the global dataset D t Dr. By Bayes' rule and conditional independence, the global posterior is
() p(.)) (x)))))))))) p(0) Il
O))
(8)
km
Each local posterior can be expressed as
p(0 x ((0))((0))))))))))
(9)
Taking the product over all clients yields
K II p()))) ( II ((D
(10)
which implics the exact relationship
I (( D)
(11)
p(0)k- K-1
Equation 11 gives the ezact global posterior (up to normalization). However, it assigns equal importance to each local postcrior, regardless of dataset size. In practice, client datasots may be imbalanced ([Dxl varies), so we instead adopted weighted approzimation, following prior works (Liu al., 2024), to better reflect the contribution of each client.
16

---

Published in Transactions on Machine Learning Research (12/2025)
Weighted Approximation.
We define normalized weights wa
[P and approximate the global poste-
rior as
D) I ((DD).
(12)
where (() = p())Da) denotes the local approximate posterior. Equation 12 scrves as an approximation to the global posterior Eq. 8 The weighting by wy ensures that clients with larger datasets contribute proportionally more and implicitly prevents the overcounting of the prior.
Each local approximation can be written as
1 (() ~ p(0)p(Da)))
Z
which implies
(O ~ p(0) th p ())
Substituting this into Eq. 12
we obtain
Il {(() s 1 II
Rm
k=
Since Ek =k=1 W/ this simplifies to
K (())a) x p(0) P(D
K
II
(13)
k=1
Equation 13 has the same form the exact global posterior in Eq. 8 except that each likelihood term is raised to the power Wk. Thus, Eq. 112 provides an approzimation to the true posterior, where each client's contribution is balanced by the fraction of its local data.
Gaussian Case. If each local posterior is Gaussian, g(0)D.)) = N Ex), then the weighted product in Eq. 2 is also Gaussian Bishoo
2006
II k=1
N(u,, Ex) x N Es),
with parameters
K
Σ51
Wkk 1-
(14)
k=1
μs Es
Σ \==1
w Di pb
(15)
C More details on IVON
Computing exact gradients in equation 4 and 5 is difficult due to the expectation term in Cx(g). A naive way to optimize is to use stochastic gradient estimators. However, these approaches are not very scalable due the high variance in the gradient estimates. Using natural gradients, Khan Lin (2017) gave improved gradient based update equations for the variational parameters and they call this approach Natural Gradient
17

---

Published in Transactions on Machine Learning Research (12/2025)
Local Epoch E 1
60
Test Accuracy
50
30
Fedlvon
20
Ledlaplace
FedAvg
10
250
500
750
1000
1250
1500
1750
2000
Communication Rounds
Local Epoch E=2
Test Accuracy
50
40
30
Fedlvon
20
Laplace
FedAvg
250
500
750
1000
1250
1500
1750
2000
Communication Rounds
Figure 4: Convergence of all the methods on CIFAR-10 dalaset with varying local epochs
VI (NGVT). The major difference between NVGI and original update equations is that learning rate is now adapted by the variance of .++ which makes these updates similar to Adam.
NVGI:
m (+1 m + 38 2++1 2 V Cx(9)]
o
21+1
-2 o 1 V3'(VozLu(4)) Cx()
Further, Khan et al. (2018) showed that the NVGI update equations be written in terms of scholastic gradient and Hessian of where of = + The vector h& contains an online estimate of diagonal Hessian. This approach called Variational Online Newton (VON) is similar to NGVT except that. it does not require the gradients of the variational objective.
ge )(0)))) Amm 11 1+1 + A h h+t1 (1 3h)hi I 3'diag diagiVa (00)
VON: m 1+1
=m
In the update of VON for non-convex objective functions. the Hessian can be negative which might make of negative, and break VON. To mitigate this issue Khan et al (2018) used a Generalized Gauss-Newton (GGN) approximation of Hessian which is always positive. This method is called VOGN.
1 M
[V, e (00)² == h))
go ((0)) + Am
VOGN: m 1+1
m
h
h ++1 = (1 nt + (0))
VOGN Khan al. (2018) improves these equations where Gauss Newton estimation used instead of Hessian which gives similar update equations the Adam optimizer. However, it still uses per-sample
18

---

Published in Transactions on Machine Learning Research (12/2025)
squaring which is costly as compared to Adam.
IVON: hg (0) m
k
(hg - lig))
h ++1 = -1- h + phg + "0 (hg + 80/ 80/A)
Further, Shen et al. (2024) improved these update equations and provided much more efficient update equations similar to Adam optimizer, which is essentially the improved variational online Newton (IVON)
algorithm Shen et al. 2024
D Ivon Hyperparameter
params
SVHN
EMNIST
CIFAR-10
initial learning rate
0.1
0.1
0.1
final learning rate
0.01
0.01
0.01
weight decay
2e-4
2e-4
2e-4
batch size
32
32
32
ESS (A)
5000
5000
5000
initial hessian (ha)
2.0
5.0
1.0
MC sample while training
1
1
MC samples while test
500
500
500
Table 6: Tvon Hyperparameters for FL experiments
E Reliability diagrams for FL experiments
Figures and show the reliability diagrams for CIFAR-10 and EMNIST experiments, respectively. The diagrams indicate that Fedivon has better-calibrated predictions compared to FedAvg and FedLaplace, shown by its lower Expected Calibration Error (ECE).
F Client data distribution in FL experiments
Figure illustrates the data distribution among clients used in the FL experiments. Each client has highly imbalanced dataset, with the number of samples per client ranging from to 32. Additionally, each client's 's dataset is limited to only subset of classes, further emphasizing the non-llI) nature of the data. This experimental setup poses significant. challenges for training a robust global server model, as the limited and
params
CIFAR-10
CIFAR-100
initial learning rate final learning rate
0.1
0.1
0.001
0.001
weight decay
1e-3
1e-3
batch size
32
32
expected sample size (A)
10000
10000
initial hessian (hp)
1.0
1.0
MC sample while training
1
MC samples while test
64
64
Table 7: Ivon Hyperparameters for personalized FL experiments
19

---

Published in Transactions on Machine Learning Research (12/2025)
FedAkg with EMNSST
Fedeaplace with EMNIST
Fedivan with EKNIST
3.0
Emptr kul
Enghiial
Erptical
Gap
Cup
0.6
Trequency
0.5
Precuency
0
0.8
Trqquency
08
2.4
0A-
0.
6.4
ECS + C/9381
SCb UULES
6.0
0.2
0.0
Figure Reliability diagrams for EMNIST experiments (Jeft: FedAvg, center: Fedlaplace, right: Fedlvon).
Clex Wirelttie
Figure Client data distribution for CIFAR-10, EMNIST, and SVHN dataset used in FL experiments.
biased data from individual clients must. be aggregated effectively to learn model capable of generalizing across all classes. This scenario highlights the complexities and practical relevance of federated learning in real-world applications.
G Client data distribution in pFL
Figure illustrates the distribution of data points across classes and clients in three pFL experimental setups with 50, 100, and 200 clients. The number of data points per client varies significantly, with some clients having over 1,000 data points and others fower than 5, indicating a high dogree of imbalance. Despite this, every client retains examples from most classes, which is crucial for training personalized models that adapt the unique data distribution of each client. This setup highlights the challenge of learning effective personalized models in pFL. Similarly, Figure shows the data distribution for the CIFAR-101 dataset..
20

---

Published in Transactions on Machine Learning Research (12/2025)
Distbriin
Cleze deritates
15
Figure 7: Client data distribution for CIFAR-10 dataset used in pFL experiments (left: 50 clients, right: 100 clients, bottom: 200 clients).
Cheet Damtbrtin
Clere Dertbates
25
Figure 8: Client data distribution for CIFAR-100 dataset used in pFL experiments (left: 50 clients, right: 100 clients, bottom: 200 clients).
21


Reviews:
[{"review_id": "2w8XTzoi4g", "reviewer": "Reviewer_UPt5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "### Summary of Contributions\n\nThe paper addresses the challenges of federated learning (FL), where data cannot be shared across clients and communication is limited. Given $K$ clients with private datasets $\\mathcal{D}_k$, the global goal in standard FL is to solve:\n\n$$\n\\theta^* = \\arg \\min_\\theta \\sum_{k=1}^{K} -\\log p(\\mathcal{D}_k \\mid \\theta).\n$$\n\nHowever, directly solving this is difficult due to data localization, and point-estimate solutions like FedAvg do not quantify uncertainty. Computing the full posterior $p(\\theta \\mid \\mathcal{D})$ would be more informative as it captures model uncertainty, but exact Bayesian inference is computationally intractable, and common approximate methods (e.g., MCMC, Laplace approximation, VI) are too expensive for FL.\n\nTo overcome these limitations, this paper proposes **FedIvon**, a Bayesian FL framework that uses **Improved Variational Online Newton (IVON)** to efficiently estimate a posterior distribution over model parameters at each client. Instead of learning a single global model, FedIvon learns client-wise posterior distributions characterized by a mean and diagonal covariance. This enables:\n- principled **uncertainty quantification**, and  \n- **natural personalization** via hierarchical Bayesian modeling, where each client adapts from a shared prior.\n\nThe key contribution is demonstrating that Bayesian FL can be performed with **near-first-order computational and communication cost**, avoiding heavy matrix operations or full posterior transmissions. FedIvon brings uncertainty estimation and personalization without the typical overhead of MCMC or Laplace-based approaches, while improving accuracy and calibration over existing FL baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper provides strong empirical support for its claims regarding predictive accuracy, calibration, and uncertainty estimation. The experiments are well-designed, cover multiple datasets, and include both standard and personalized FL settings.\n\nHowever, a key contribution of the paper is the claim that FedIvon achieves computational efficiency comparable to first-order optimizers while being more efficient than other Bayesian FL methods. This is only argued qualitatively. I could not find quantitative evidence such as runtime, FLOPs, memory usage, or per-round communication cost. As a result, while the empirical results convincingly support the performance and uncertainty-related claims, the efficiency-related claims are not fully substantiated by measurable evidence.", "audience_yn": "Yes", "audience_explanation": "The paper addresses Bayesian federated learning, uncertainty quantification, and personalized FL; topics that are highly relevant to the TMLR community. Researchers working on federated learning, Bayesian deep learning, or efficient optimization methods would likely find these contributions valuable.", "requested_changes": "$\\mathbf{1.}$ The paper claims that FedIvon has computational cost comparable to first-order optimizers like Adam and is more efficient than existing Bayesian FL methods. However, this argument is presented qualitatively. I was unable to find quantitative evidence such as runtime, FLOPs, memory usage, or communication cost comparisons. Providing such empirical results would make this contribution more convincing and further strengthen the paper.\n\n$\\mathbf{2.}$ Clarify the limitations of the diagonal covariance assumption.  \nThe method approximates posterior distributions using diagonal Gaussians. A brief discussion on when this assumption may break down or how it might affect uncertainty calibration would improve the transparency and rigor of the work.\n\n$\\mathbf{3.}$ Discuss the choice of only using 2 local epochs.  \nIn most federated learning setups, the number of local training epochs is typically greater than 2 (e.g., 5 or 10), especially when communication is costly. It would strengthen the paper to either justify this choice or include results for higher values of local epochs to better understand convergence, stability, and communication–computation trade-offs.", "broader_impact_concerns": "I do not see any ethical concerns.", "additional_comments": "My recommendation is to accept after minor revisions. The work is not highly novel from a theoretical standpoint, but the empirical results and integration are solid and align with TMLR’s standards, which emphasize usefulness and clarity over pure novelty.", "extra_scores": {}, "extra_text": {}}, {"review_id": "jHiwtCEv19", "reviewer": "Reviewer_Kt8z", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposed a novel Bayesian federated learning (FL) method, where each client performs Ivon and updates its local weight and Hessian, and the server updates the weight and Hessian with weighted update of each client. The method is significant as a second-order optimization with efficiency of first-order method. Besides, the FedIvon framework also enables client personalization. The significance of the method is supported by empirical results on various datasets. \n\nStrength:\n- The paper is well structured and well written. The background is clearly introduced, the challenges and state-of-the-art methods are summarized. \n- The paper adapts the Ivon algorithm and deploys it in FL setting. \n- The empirical results are strong, the FedIvon approach achieves better metrics in data heterogeneous settings compared with the baselines.\n\nWeaknesses:\n- The paper does not invent Ivon algorithm, therefore partially limiting its originality. \n- Since the local update subsumes Ivon algorithm, whichever limitations Ivon has will also be inherited in the FedIvon framework. \n- More baselines of Bayesian FL are mentioned in related work but not compared with in experiments. \n- To show the efficiency, maybe experiments on runtime can be added, e.g., compare with FedAvg and Fedlaplace like Figure 4, but the x-axis being runtime.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The FedIvon framework compares with FedAvg and FedLaplace in terms of Bayesian FL, and clearly exhibits better results on EMNIST, SVHN and CIFAR datasets in terms of convergence, accuracy, loss, and uncertainty. Besides, FedIvon is compared with pFedME and pFedBayes baselines, and shows better results on CIFAR datasets.", "audience_yn": "Yes", "audience_explanation": "This work is novel and enough contribution to FL community.", "requested_changes": "- While this work shows significant improvement with Ivon algorithm on clients, limitation of the FedIvon framework should be more discussed. \n- Maybe add more comparisons with Bayesian FL methods.\n- Maybe add experiments showcasing the efficiency and computational overhead of FedIvon compared with FedAvg and FedLaplace.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7TibuWwiiR", "reviewer": "Reviewer_Q1Kh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces FedIvon, a novel federated learning (FL) algorithm that brings Bayesian uncertainty estimation and personalization into FL while remaining highly efficient. The core idea is to perform variational Bayesian inference on each client using an Improved Variational Online Newton (IVON) optimizer, which is a second-order (Newton-like) method that approximates the Hessian implicitly. By leveraging IVON, each client maintains a Gaussian posterior with a diagonal covariance (uncertainty) over the model parameters, but the computational cost per communication round is kept comparable to standard first-order methods like Adam. This means FedIvon can capture curvature information for faster convergence and uncertainty quantification without the heavy overhead typically associated with second-order methods or exact Bayesian updates. Another key contribution is that FedIvon naturally supports personalized models for each client in a hierarchical Bayesian manner: the server aggregates client posteriors into a global posterior which then serves as a prior for subsequent client updates. This hierarchy allows each client to personalize its model to its local data while still benefiting from global knowledge sharing. In summary, the paper’s main technical contributions are:\n\n- Efficient Bayesian Federated Optimization: An FL algorithm that uses second-order variational optimization (via IVON) to approximate each client’s posterior. It achieves per-round computation similar to Adam but yields better predictive performance and better-calibrated uncertainty estimates than both classic FL (point estimate) and prior Bayesian FL approaches. The method implicitly computes a diagonal Hessian approximation on-the-fly, avoiding expensive explicit Hessian calculations as required by methods like FedLaplace. This is a strength because it bridges the gap between Bayesian FL’s benefits and practical efficiency/scalability in resource-constrained settings.\n\n- Uncertainty Quantification and Personalization: By maintaining a Gaussian posterior for each client model, FedIvon provides predictive uncertainty estimates (via the variances) in addition to point predictions. The uncertainty information is used for tasks like model calibration and out-of-distribution detection. Moreover, in a hierarchical Bayes fashion, the global model’s posterior (mean and covariance) acts as a prior for clients, enabling client-specific personalized models that account for local data heterogeneity. This personalization is seamlessly integrated into the federated updates (clients receive the server’s posterior and update their own posteriors accordingly), which is an elegant way to handle non-i.i.d. data.\n\nOverall, the paper’s contributions tackle an important problem that is combining Bayesian FL with efficiency. A notable strength is the comprehensive experimentation demonstrating improved accuracy and calibration. One potential weakness in the contributions is that the approach builds upon an existing optimizer (IVON from Shen et al., 2024) and known Bayesian personalization frameworks, so the novelty is more in the integration and efficiency gains in FL context rather than an entirely new theory. Another limitation worth noting upfront is that FedIvon assumes a diagonal covariance for tractability. This is pragmatic, but it means uncertainty modeling is limited to diagonal (no correlations between parameters). This could restrict the expressiveness of the posterior, though it keeps computation feasible.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes, for the most part the claims are supported by convincing evidence, through arguments and experiments. \n\nThe submission’s central claims are: (1) FedIvon achieves superior predictive performance and uncertainty calibration compared to state-of-the-art baselines, (2) it incurs minimal overhead (computation and communication) similar to first-order methods, and (3) it enables effective personalization for heterogeneous clients. The evidence provided to back these claims is generally solid, with a few gaps as discussed below:\n\n- Technical Correctness and Clarity: The algorithm is described with sufficient clarity. The paper includes pseudocode for the client update (IVON) and the federated aggregation (Algorithm 1 and 2), which makes it easy to verify what exactly is implemented. The derivation of how client posteriors are combined into a global posterior is grounded in known results (product-of-Gaussians formula) and is clearly explained. The method for personalization is also described clearly. These descriptions suggest the approach is technically sound. Important details like how the Hessian is estimated on clients (via a reparametrization trick yielding an unbiased diagonal Hessian estimate) are given, showing that IVON indeed provides second-order information without high cost. The authors make a convincing argument that there is no explicit costly Hessian inversion or large matrix communication; instead, only diagonal variance vectors are maintained, which is efficient. This lends credibility to the claim of “efficient second-order optimization.” However, one concern is that the paper does not formally prove convergence or provide a theoretical complexity analysis as it relies on references and intuition (e.g. IVON’s known efficiency) to support these points. In the absence of theoretical proofs, the clarity of the algorithm and the empirical evidence become even more crucial, and fortunately the empirical results do back up the efficiency and performance claims.\n\n- Empirical Support for Performance Claims: The experimental evaluation is thorough and directly addresses the claims:\n\n  - *Accuracy and Calibration:* In standard federated learning experiments (classification on EMNIST, CIFAR-10, SVHN), FedIvon consistently outperforms the baselines in terms of test accuracy and yields much better calibrated predictions. Table 1 compares FedIvon against a representative strong FL baseline (FedAvg) and a recent Bayesian FL method (FedLaplace) on multiple metrics. FedIvon achieves the highest accuracy on all three datasets and dramatically lower calibration error (ECE) and negative log-likelihood, indicating its probabilistic outputs are more trustworthy. These results support the authors’ claim that FedIvon improves predictive performance and uncertainty quantification over both non-Bayesian and Bayesian competitors. \n\n  - *Out-of-Distribution (OOD) Detection:* The authors also test whether the improved uncertainty translates to better OOD detection. Using the entropy of predictions to detect when an input is outside the training distribution, FedIvon achieved the highest or equally-high AUROC on most dataset pairs. In Table 2, FedIvon’s OOD detection AUROC is slightly better than or on par with FedAvg and clearly better than FedLaplace in two of three cases. (It’s worth noting one nuance: on CIFAR-10 as in-distribution, FedAvg had a slightly higher OOD AUROC than FedIvon (0.7896 vs 0.7662). The authors describe FedIvon’s OOD performance as “better or competitive,” which is fair. FedIvon was best on EMNIST and SVHN scenarios, and a close second on CIFAR-10. This small discrepancy doesn’t undermine the overall claim, but it suggests FedIvon’s uncertainty is at least not worse than FedAvg’s in any case, and usually better.)\n\n  - *Personalization (Heterogeneous data):* A significant claim is that FedIvon can personalize effectively for each client thanks to the hierarchical Bayesian approach. The evidence for this comes from experiments simulating non-i.i.d. scenarios on CIFAR-10 and CIFAR-100 (with different types of data heterogeneity). The paper compares FedIvon to strong personalized FL baselines including pFedME (a state-of-the-art optimization-based pFL method), pFedBayes (a recent Bayesian personalized FL), and pFedVEM (variational expectation maximization approach), across varying numbers of clients. The results show that FedIvon achieves equal or higher personal-model accuracy than these methods and significantly higher global-model accuracy. \n\n  - *Efficiency and Scalability:* The paper’s efficiency claims are mostly supported qualitatively and through algorithmic analysis rather than direct runtime plots. The authors argue that per-round computation of FedIvon “matches first-order optimizers (e.g., Adam)” and that it does not incur the usual overheads of second-order methods in FL. This is plausible given that IVON’s updates require similar operations to Adam (plus simple variance updates) and avoid expensive steps like Hessian inversion or MCMC sampling. While the experiments report using up to 2000 communication rounds and reasonably large numbers of clients (up to 200) without issues, the paper does not explicitly graph runtime or communication costs. The submission could be stronger by quantitatively showing the runtime per round or total training time for each method, but the qualitative description and the known properties of IVON give confidence that the method is indeed efficient. Moreover, the authors explicitly discuss why FedIvon’s initial convergence can be slightly slower (due to stochastic weight sampling for gradients) and show an ablation varying local training epochs; this kind of analysis indicates they have a good understanding of the method’s behavior and that any efficiency trade-offs (like a slower start) are minor and well-managed. \n\n  - *Adequacy of Baselines:* The baselines chosen are mostly appropriate and help validate the claims, though there is room for a couple more comparisons. For standard FL (non-personalized), the paper benchmarks against FedAvg (the classic method) and FedLaplace. FedAvg is a natural baseline for accuracy (point estimate, no uncertainty), and FedLaplace (Liu et al., 2024) is a recent Bayesian FL method that also approximates client posteriors as Gaussians via Laplace’s method. FedLaplace is a strong baseline because it represents the idea of Bayesian FL with second-order info (though it computes Hessians explicitly at the clients). The results showed FedIvon outperforming both, which supports the claim of being state-of-the-art. One might question if other baselines like FedProx (for non-iid data) or FedPAshould have been included. Overall, the evidence provided is sufficient and the baselines are adequate to support the main claims. Only minor additions (like one more baseline in the global setting, or perhaps comparing to an ensemble approach such as FedBE) could further strengthen the empirical support, but their absence does not invalidate the claims.", "audience_yn": "Yes", "audience_explanation": "Yes. The topics and findings of this paper are likely to interest a segment of TMLR’s audience, particularly those working in federated learning, Bayesian deep learning, and uncertainty in machine learning. Federated learning is a well-established research area, and one of its current frontiers is how to handle heterogeneous data across clients (non-i.i.d.) and how to provide trustworthy models with uncertainty estimates. This paper squarely targets both issues by proposing a Bayesian personalization approach. The idea of using hierarchical Bayesian models in FL (i.e., a global prior and client-specific posteriors) has been gaining traction recently, and this work provides an efficient way to do it, which would be of interest to researchers who found previous Bayesian FL methods too cumbersome or impractical.", "requested_changes": "- Include Additional Baselines or Justify Their Exclusion: To solidify the SOTA claims, the authors should include at least one additional baseline in standard FL experiments: e.g., FedPA (posterior averaging), FedBE (Bayesian ensemble distillation), or FedProx (a widely used FL method for non-i.i.d. data). If these methods are omitted, please explicitly justify. A brief discussion of how FedIvon differs from and improves upon FedPA in particular would clarify novelty.\n\n- Add Theoretical Insight or Stability Guarantees: The paper currently lacks theoretical support. Please either (a) sketch convergence behavior using results from the IVON literature, or (b) include a qualitative explanation of why FedIvon remains stable in the federated setting (e.g., guaranteed positive diagonal Hessian, bounded variance updates, etc.). \n\n- Clarify Communication and Computational Overhead: Please describe the communication cost per client per round (e.g., sending mean and variance vectors = 2× model size) and clarify why this remains practical. Also discuss computational complexity—IVON appears to require two forward-backward passes per batch, similar to Adam. If available, please provide runtime estimates or FLOP comparisons to support the claim of efficiency.\n\n- Discuss Limitations of the Method: A short limitations paragraph would improve balance. Points to mention include: the diagonal covariance assumption (no parameter correlation), added memory cost from tracking variances, client-side compute overhead, and the fact that experiments were conducted on relatively small CNNs in the vision domain. Acknowledging these would help set realistic expectations.", "broader_impact_concerns": "The paper’s broader impact discussion (if any) appears to be minimal, so here I raise a few concerns and considerations that should be acknowledged:\n\n- Fairness Across Clients: Personalized federated learning can inadvertently introduce fairness concerns. Since each client learns its own model, differences in data quantity or quality across clients might lead to some clients obtaining much better models than others. \n\n- Reliability of Uncertainty Estimates: While FedIvon improves calibration overall, incorrect uncertainty estimates can be harmful if relied upon in high-stakes settings. A concern is whether the model might be overconfident on out-of-distribution or minority group data that it did not see (especially if a client’s training data is limited or biased).\n\n- Privacy and Security: Federated learning is often praised for on-device privacy, but sharing model updates (means and variances in this case) can still leak information about clients’ data. The paper does not discuss privacy-preserving enhancements (like differential privacy or secure aggregation). A potential impact risk is that malicious servers or other clients might infer personal information from the transmitted posterior parameters.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)