Connecting Federated ADMM to Bayes
Authors: Siddharth Swaroop
Mohammad Emtiyaz Khan
Finale Doshi-Velez
Venue: ICLR 2025 Poster (2025)
Keywords: federated learning
bayesian
variational inference
admm
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We provide connections between federated ADMM methods and Variational Bayes federated learning methods, providing ways to improve them

Abstract: We provide new connections between two distinct federated learning approaches based on (i) ADMM and (ii) Variational Bayes (VB), and propose new variants by combining their complementary strengths. Specifically, we show that the dual variables in ADMM naturally emerge through the "site" parameters used in VB with isotropic Gaussian covariances. Using this, we derive two versions of ADMM from VB that use flexible covariances and functional regularisation, respectively. Through numerical experiments, we validate the improvements obtained in performance. The work shows connection between two fields that are believed to be fundamentally different and combines them to improve federated learning.

Published as a conference paper at ICLR 2025
CONNECTING FEDERATED ADMM TO BAYES
Siddharth Swaroop Harvard University, US
Mohammad Emtiyaz Khan RIKEN Center for Al Project, Japan
s siddhatthseas hww.aad..edarereeaa
Finale Doshi- Velez Harvard University, US
ABSTRACT
We provide new connections between two distinct federated learning approaches based on (i) ADMM and (ii) Variational Bayes (VB), and propose new variants by combining their complementary strengths. Specifically, variables in ADMM naturally emerge through the "site" parameters used in VB with isotropic Gaussian covariances. Using this, we derive two versions of ADMM from VB that use flexible covariances and functional regularisation, respectively. Through numerical experiments, we validate the improvements obtained in per- formance. The work shows connection between two fields that are believed to be fundamentally different and combines them to improve federated learning.
show that the dual
INTRODUCTION
The goal of federated learning is to train a global model in the central server by using the data distributed over many local clients (McMahan al.) 2016). Such distributed learning improves privacy, security, and robustness, but is challenging due to frequent communication needed to synchronise training among nodes. This is especially true when the data quality differs drastically from client to client and needs to be appropriately weighted. Designing new methods to deal with such challenges an active area of research in federated learning.
We focus on two distinct federated-learning approaches based on the Alternating Direction Method of Multipliers (ADMM) and Variational Bayes (VB), respectively. The ADMM approach synchronises the global and local models by using constrained optimisation and updates both primal and dual variables simultaneously. This includes methods like FedPD (Zhang et al. (021). FedADMM
(Gong et al. 2022 Wang et al., 2022 Zhou and Li 2023) and FedDyn (Acar et al. (222T). The VB approach, on the other hand, uses local posterior distributions as messages and multiplies them to compute an accurate estimate of the global posterior (Ashman al. (022). This follows the more general Bayesian framework which has a long history in distributed and decentralised computations (Mutambara] 1998 Tresp 2000 Durrant-Whyee (001). Despite decades of work, the two fields remain disconnected and are thought to be fundamentally different: there is no notion of duality in VB, and posterior distributions are entirely absent in ADMM. Little has been done connect the two at a fundamental level and our goal here is to address this gap.
In this paper, we provide new connections between ADMM and VB-based approaches for federated learning. Our main result shows that the duaalvariable used in ADMM naturally emerge through the "site" parameters in VB. We show this for specific case where isotropic-covate covariance Gaussian approximations are used. For this case, we get a close line-by-line correspondence between the ADMM and VB updates. The expectations of the local likelihood approximations, also known as the sites, yield the dual terms used in the local updates of ADMM.
We use this connection to derive new federated learning algorithms. First, we extend the isotropic case to learn full-covariances through the federated VB approach; see Sec. 3.3 This leads to an ADMM-like approach with an additional dual variable to learn covariances, which are then used as preconditioners for the global update. Second, building upon functional regularisation approaches from Bayesian literature, we propose to use similar extensions to handle unlabelled data, providing a
relationship to federated distillation methods (Seo al. 2020 Li and Fedmd 2019 Lin al. 2020a

---

Published as a conference paper at ICLR 2025
Sco al. 2020; Wu et al., 2021), well work in function-space continual learning (Buzzega
et al. 2020 Kirkpatrick et al. 2017 Pan et al. 2020; TTtssaas et al.] 2020); see SeefSSe. 3.4 Our empirical
results in Sec. 4 s/show that the algorithms lead to improved performance (improved convergence)
across a range of different models, datasets, numbers of clients, and dataset heterogeneity levels.
2 FEDERATED LEARNING WITH ADMM AND BAYES
We start by introducing federated ADMM and Bayesian federated learning approaches, and then, in Secs. 3 and 4 we will provide new connections along with new algorithms and empirical results. The goal of federated learning is to train a global model (parameters denoted by w.g) by using the data distributed across K clients. Ideally, we want recover the solution that minimises the loss over all the aggregated data, for instance, we may aim for
w = arg miinn E
Σ ((ww) + 8R(w))
where lg is the loss defined the k 'tt client using its own local data Dj. and R is a regulariser with 2 as a scalar (set to 0 when no explicit regularisation is used). However, because the data is distributed, the server needs query the clients get access to the losses &.. Then the goal is to estimate W while also minimising the communication cost. Typically, instead of communicating la or its gradients, local clients train their own models (denoted by wi for the k'th client) and communicate them (or changes to them) to the server. The server then updates w and sends it back to the clients. This process is iterated with the aim of eventually recovering w
The Alternating Direction Method of Multipliers (ADMM) is popular framework for distributed and federated learning, where the local and global optimisation are performed with constraints that they must converge to the same model. That is, we use constraints W; = Wg enforced through the dual variables, denoted by vg.- The triple (wk; 1k: wgg) are updated as follows,
Client updates:
WR f argmin (ww) + ug w + Jall w
Uk Uk a(wk wg),
for all k,
(2)
K
1 K
1 wk
Server update: Wy
2k
a
Here, a > is a scalar and we denote & N where Ng is the size of the data the k'th client. Upon convergence, we get wi = =w and 2xz = (ww for all k.
Such updates are instances of primal-dual algorithms and many variants have been proposed for federated learning (Zhang et al. 2021 Gong et al. 2022 Wang et al. 2022 Zhou and Li] (023). For example, FedDyn (Acar et al. (022) >> is perhaps the best-performing variant which uses a client- specific ag a/Ng. It also adds an additional hyperparameter through local veight-decay regulariser added to each local client's 's loss. There are also simpler versions (not based on ADMM), for instance, FedProx (Li et al. 2020) where Uy is not used (or equivalently fixed to 0 in Eq. 2), and even simpler Federated Averaging or FedAvg (McMahan et al. 2016), where the proximal constraint w also removed. Such methods do not ensure that w " at convergence. The use of dual-variables is a unique feature of ADMM to synchronise the server and clients.
The Bayesian framework for federated learning takes a very different approach to ADMM. Instead of a point estimate wn" the goal for Bayes is to target global posterior Pg (un) = ppwr[DD, D2,..., Dx). Instead of loss functions for each Dr, we use likelihoods denoted by p(Dx|uu (w). The solution w in [Bq. can be seen as the mode of the posterior Pg whenever there exists likelihood such that log ww) = (t( and prior log po(w)) = -RR(w), where both equations need to be true only up to a constant. Therefore, targeting the posterior Py Py(w) (w) also recovers w Often, we compute local posteriors p& px(w)= p(w) DDx) at clients and combine them the servers, for instance, Bayesian committee machines (Tresp 2000) use the following update,
K
Py α Po
t
(3)
2

---

Published as a conference paper at ICLR 2025
where (x()) p& pw(W)/Po(w)) (w) is essentially equivalent to the likelihhod p(Dk/w) over the data DR at the client k. However, the data is never communicated and only the local posterior ps is passed. Unlike ADMM, the update is not iterative because the posteriors Pk: are assumed to be exact, giving rise to a one-step closed-form solution for Pg- Earlier work on distributed and decentralised Bayes focused on such cases (Mutambara, 1998 Durrant- Whyte, 200T] [Soott et al., (022). This is perhaps the main reason why Bayesian updates can seem disconnected from the ADMM-style updates.
We will derive connection by using approximate Bayesian approaches, which are required to scale Bayes to large complex models. Specifically, we will use Variational Bayes (VB), although there also exist work on Monte Carlo methods (Al-Shedivat al. 2021] Deng et al. 2024 Liang et al. 2024). VB finds an approximate posterior, denoted by 9g (), by minimising the Kullback-Leibe
(KL)) divergence TD Kg ( Pg]- For exponential-fane famiily distributions, the global og (w) has a similar form to Eq. but & (w) are replaced by likelihood approximations bg (w), also known as "sites". For instance, when the prior po(w) has the same exponential-fomm as qg (w), the optimal 0g (w) is
K
II
gg α Po
R=
The sites fg (w) in VB are parameterised through natural gradients as reported by Khan and Lin 2017
Eq. 11); see exact expression in Khan and Nielsen (2018 Eq. 18) or Khan and Rue (2023 Sec. 5.4). In this paper, we will show that a distributed estimation of (w) leads to a natural emergence of dual variables and ADMM-liK- algorithm. Specifically, we will rely on the Partitioned Variational Inference (PVI) procedure of' Ashman et al. (2022) and show that, for a Gaussian g, (t)). we get a line-by-line correspondence ADMM. We note that the other prominent approach called Expectation Propagation (EP) (Minka 2001 Guo et al.) (202)) also obtains site parameters by optimising the forward KL, I1 KL Pg 4gl but there no known connection to the dual variables.
CONNECTING VARIATIONAL BAYES TO ADMM
We precisely connect Variational Bayes, specifically Partitioned Variational Inference (PVD) (Ashman eet al. 2022), ADMM-style methods for federated learning. We start with PVI, noting it has similar components to ADMM. We then make approximations to derive a method (FedLap) that has close line-by-line correspondence ADMM. We use this connection derive new variants of ADMM to improve it, by (i) using full covariance information, and (ii) including function-space information.
3.1 CONNECTING PARTITIONED VARIATIONAL INFERENCE (PVI) TO ADMM
Partitioned Variational Inference (PVI) (Ashman et al. 2022 Bui et al. 2018) aims to find the best approximation 4g Q where Q is set of candidate distributions (for example, set of Gaussian distributions). The goal is to get the VB solution in Eg. 4 but by using an iterative message passing algorithm where local approximations qu (w) send the sites (w) as messages to improve 4, (w). The method updates the following triple (gk, tk. qg) as shown below,
log ((Dx)w) Ex (10)
Client updates:
4he f arg max Ey DEQ
DD DKrll ggl.
Ea x- i
qk
(5)
Qg
Server update:
4g
Po
f
The update of 4g (w) is exactly the same as Eq. 3 but uses ig (w)) which in turn is obtained by using the ratio gu (ww)// /4 (w) in the second line. The lg (w) essentially uses the discrepancy between the global and local distributions. This is then used in the first line to modify the local qu (x(w)) (w). Similarly to ADMM, at convergence
have (w) of (w). The site parameters are related to natural gradients (Khan and Lin 2017: Khan and Nielsen, 2018), therefore we may expect them to be related to dual
variables oy in ADMM which estimate gradients. In what follows, we will derive this precisely.
3

---

Published as a conference paper at ICLR 2025
The PVI update has line-by-ling correspondence ADMM. We replace log ((pp/ by -k(ww) and turn the max in the first line into a min, and write the second and third lines in log-space,
Client updates:
9k arg By((x( + F,g log fx] D qEQ
log & log & + p (log th log Qg)
(6)
Server update: log 4g + log Po + Σ log f + const,
where we have added a damping factor 0 < < as can slow down the rate of change of l (w). which can be important especially in heterogenous settings (Ashman al. [022) although the theory suggests that = should be ideal. The three updates have a similar form to the ADMM updates in Eq. The first line uses additional expectation over ((w) with the Fuclidean distance replaced by the KL divergence. The role of (x(w) appears be similar to that of the dual term og in ADMM. The update of (wr) in the second line is also very similar, while the update of qg (w) has some major differences. Next, we will choose the distributions PVT accordingly to get even closer to the ADMM update.
3.2 FEDLAP: A LAPLACE VERSION OF PVI
We now choose a specific form of the distributions in PVI to make even closer to the ADMM update, deriving new method which we call FedL.ap In Sncs. 3.3 and 3.44 we will use this connection to derive new, improved, variants of ADMM. Here, we set the family Q to be the set of isotropic Gaussian distributions N(w; m, I/5) where the mean m needs be estimated while the covariance fixed to I/ó where >> 0 is scalar. We also set the prior po(w) to zero mean Gaussian with the same precision as q(w). These choices are shown below,
(y(w) x N N(w; wg. I/8),
I (xww) x N(w; wk. I/8),
(o(w) x A N(w; 0, 1/6).
(7)
These choices imply that & (ww) takes a Gaussian form where only the linear term needs to be estimated, that is, we need find Us such that,
(e(w) 60%
(8)
This form ensured due to the form of the optimal solution in Eq. 4 (and can be shown more formally by using natural gradients). Roughly speaking. this is because both qg (w) and Po(w) have the exact same Gaussian form, therefore also appears in the expression of fz. (w) as well.
Plugging these in Eq. 6 and making a delta approximation (Khan and Rue, 2023 App. C.I), we get Fedl Fedlaap, ap,, Laplace variant of the PVI (we first provide the update, and then its derivation),
Client updates:
+ arg min la (w) + dug w + 26llw wyll"
UR f- + p(w; wg)
K
Server update: W Σn
Uk.
k=1
The first line is obtained by making the delta approximation, that is, Bgp.W(mmm..///lll so g(m).
Eg (e (w)) + E. [log (w)) + E., log q(w)
Ex (m) + dv, m + 5 1 + const.
4g (w)
Then, rewriting w, we get the first line. The second line follows due to the Gaussian form,
gog w + T w + (Sw)) w -8 w)
f + (w wyy)
The last line follows in the same fashion where we update the mean w, of qu (w).
The derivation shows clearly that the (w) (w)) terms used the KL minimisation PVI gives rise to the term V. in the ADMM update in Eq. Similarly to ADMM and PVI, convergence have w = w.. We can make the FedLap update look even more similar to the ADMM update by change of variable, absorbing into Ug (see Eq. 18 in App. A). We also note three subtle differences.
4

---

Published as a conference paper at ICLR 2025
1. First, instead of using to update both wy and Uk. here The Bayesian update suggests using the prior parameter 8 in the local client update, or the overall weight-decay. For p. we follow Ashman al. (2022) and set it to 1/K.
two separate parameters.
2. Second, by dividing the first line by N₁ we can recover Eg used in ADMM (see the first line in Eq. 2). This suggests that FedLap scales the ||w , term by 6/Ng, matching with a practical choice in FedDyn's 's codebase which seems to work well empirically.
3. Third, the update of wy in PVI does not have un in it or division by K. This is due to updates that follow the form in Eq. 3 where expect different l to already be automatically weighted: the Bayesian update does not suggest any additional weighting. Despite this difference, FedLap performs well the best-performing ADMM method (FedDyn) in our experiments.
FydDyn and LedLap have the same computation and communication cost. The communication cost P from clients to server (and server to clients), where P the number of parameters in the model. Each client k's computation cost is dominated by ((NxP)) where E is the number of epochs during optimisation. Computation cost at the global server is negligible, requiring simply adding together vectors of length P.
3.3 FEDLAP-COV: A NEW ADMM VARIANT WITH FLEXIBLE COVARIANCES
Our connection gives direct way to design new ADMM variants to improve it: use different candidate Q, that is, choose different posterior forms for the candidates q(w). Here, we demon- strate this for multivariate Gaussian q( () N(w; m. Σ) where, unlike the previous section, we aim also estimate the covariance matrix E. This extension leads to two dual variables, where the second dual variable acts as a preconditioner similar to the Newton update. We will often write the precision matrix $ = E because it is directly connected to the Hessian (denoted by H) which more natural for Newton-lik, like update. In practice, use diagonal matrix because it scales better to large models, but our derivation is more general.
We make the following choices to derive the new variant which we will call FedLap-Cov.
(ww) x N(w; Wgg Sy
(e(w)) x ((w; u. SF
(a(w) x (ww; 0, I/6).
(10)
Similar to the isotropic Gaussian case, these choices imply that fy (w) also takes a Gaussian form where we need to find pair of dual variables, vector Us and a symmetric square matrix Vx.
ix (w) x dh w- w
Vew
(11)
Again, the existence of (nt VE ensured due to Eq. but we skip the details here.
Now, we can simply plug-in the forms of (yk, (x, qn) into Eq. 6 to get the update, The main differences are, first, the inclusion of V, in the client's 's dual term (highlighted with red) and, second, the preconditioning with Sg used in the server updates,
Client updates:
(WE, S&) f- arg min gEQ
[e([e(w)] + By [ w f/w" V.w) + D 1 ngl
Uy t Uk + p (S Sywg) and V/ - V + P (S, - -Sy),
K
K
Server updates: Wg t S
Uk
where Sg 51
Vz. (12)
Similarly to before, the updates can be simplified by using delta approximation to derive a Laplace variant. Because we want local Sy to also be updated, we can use second-order delta approximation for any q(w)= = N(w: m. S 1-1).
E2 14 (w)) E. (m)
(m) ) + (w- ())44()))))m)))m)))))))))))))))) + "H)(()( -m)]
(13)
= lg (m)) Tr [Hamm)
where Hm denotes the hessian of & at w = m. The approximation decouples the optimisation my and S₁. at client (we assume that H(m) does not depend on avoid requiring
5

---

Published as a conference paper at ICLR 2025
higher-order derivatives), allowing to write them two separate updates,
FedLap-Cov client update:
w - arg min & + 2 w 2-2200 V. 2 - wylls (14)
S - H+(wx)= - V& Sg
A full derivation is in App. B Here, || w Sw denotes the Mahalanobis distance. The two updates are performed concurrently and we have used (w( the most recent parameters wik- These two steps along with the update for Us. V1, w Sg in Eq. 12 constitute our new Laplace variant, which we call FedLap-Cov. our experiments we use diagonal matrices (V,, H&, S), which already shows improved performance.
FedLap-Cov Cov improves several aspects of ADMM and FedLap. First, the update of Us uses the parameters Skw,, and S, weighted by their respective precisions. As a result, the uncertainty is naturally incorporated when summing Vy in the global updates to get wi.. Second, the global updates preconditioning through S, which also incorporates uncertainty in the global variables. Third, S is an accumulation of all V2 which captures the difference between S₁ and S.. Finally, the Euclidean distance is changed to Mahalanobis distance. We expect these aspects to be useful when dealing with heterogeneity across clients. The covariances should help in automatic weighing of information, which can potentially improve both performance and convergence.
For valid algorithm, we need to ensure that S₁ and Sg remain positive-defintt throughout. This can be ensured by rewriting the V& update by substituting Sy from Eq. 14
Vx + V2 + p(S - S₂) V + (Ha(HH)) Vv.) = (1 - P) + (Hz(w)))
Therefore, if the Hessian is positive semi-definite all V2 will be too, which will also imply both S, and S₁ remain positive definite. We can enforce this by using a Generalised Gauss-Newton approximation (Schraudolph) 2002 Martens 2020) or by using recent variational algorithms that ensure positive definite covariances (Lin et al. 2020b Shen et al. (022).. We provide the final implemented FedL.ap- Cov update in [q. 21| Similarly to ADMM, PVI and FedLap, at convergence we have " = w and S = S
FedLap-Cor Cov has communication cost 2P from clients to server (and server to clients), we send both and V2, (this is for a diagonal matrix, as in all our experiments; for a full-covariance structure the theoretical cost is 3P/2 + p²/2), Client k's computation cost is O P(F F: + 1)) as we need an additional backward pass through the dataset to calculate V2 in our implementation, which is small additional cost. We could alternatively use more efficient implementations, such as online implementations to reduce this cost (Daxberger et al. 1021 Shen et al. 2024). Overall, FedLap-Cor adds new improvements to FedLap by using a more flexible posterior distribution, and we see this empirically in [Sec.
3.4 TEDLAP-FUNCE IMPROVING BY INCLUDING FUNCTION-SPACE INFORMATION
In the previous section, we derived FedLap-Cov which improves FedLap by including full covariance information. In this section, we derive FedLap-Func, which improves FedLap by including function- space information over unlabelled inputs. Specifically, we assume some inputs are available to both a local client and the global server, and send soft labels (predictions) over those inputs, along with the weights we send in FedLap. This additional information can be seen improving the gradient reconstruction of each client's s data compared to just a Gaussian weight-space approximation, thereby improving the quality of information transmitted between local clients and the global server (Khan) and Swaroop 2021; Daxberger et al. 2023).
We start by writing FedLap's 's update at the global server optimisation problem,
K
Wy = arg min 40
log &
(15)
Aum client I's contribution
where we note that the solution to this problem (followed by calculating the Laplace covariance at Wy) gives FedLap-Cov's server update from Eq. 12 and it therefore looks very similar to Eq. 6
Following ideas from continual learning and knowledge adaptation (Khan and Swaroop, 2021 we improve the contribution from client & by also sending information in function-space, instead of only

---

Published as a conference paper at ICLR 2025
sending in weight-space. We do this client, or shared). We derive a similar equation to Daxberger et al. (2023). but through a more general form which avoids Taylor series expansion, and allows us adapt the theory to federated learning. Each client contribution to the global server in Eq. 15 becomes (derivation in App. C.I).
of inputs Mg (where M3 can be different for each
Σ iEMM
E(De w)
(w (w-ww)) w).)
[ I(w wo) +
log Eu (u()))
(16)
new weight-space term
unchanged from before
space information
where E(V:. w) is the loss at parameters with label yı. Ús is the soft label (prediction) over an input M4 using local client weights wi, and H- the (approximate) Hessian of the loss over points in M₁ at weights uy.. We see that the new weight-space term is subtracted from the additional function-space information: this avoids double-counting the contribution from points i E Mr. Although using this in the optimisation at the global server in Eq. 13 seems to require sending additional weight-space information (w), and HM)), the equations simplify, allowing us to still send only V₂ and Uμ. except that compared to earlier, these remove the contribution from points i E MM. We additionally send ú; for i E MA. which is a very small cost when the same inputs M₁ available the clients and the server.
We can simplify the equations further by again approximating all the covariances to be equal to I/5, like in FedLap. In App. C.2 we show that this gives,
Client updates:
wk + arg min (xw)) dv w + fo w 11
Same as FodLap
K
T((ii, w) + Σ Σ
T Tl(Ding w)
iEML
KE
(17)
Additional function
Uk + Ub + p(wss wg),
K
Server update: wg + arg min
Σ LiEM
Tl(ūis w) - oug w + 16112011
w
where Oū is the prediction over input i E ME using the previous round's client model weights wy- Di, is the prediction using the global server weights wy, and we have introduced hyperparameter т that upweights the function-space contribution (we might want > when we have few points in Mg). BedLap-Fune keeps the properties of FedLap at a stationary point: w = w (at a stationary point, all the function-spen space terms contribute zero gradient). We also see nice property of the algorithm: FedLap-Fune recovers FedLap when M₁ 0. FedLapp-Fing can therefore be seen as improving Fedl Fedlap by introducing some points in funetion-space.
We note that our points in Ma do not have to be points from client's data De. and can be publicly available (unlabelled) data, like in Federated Distillation (Li and Fedmd 2019 Lin al. 2020aa [Soo let al. 2020 Wu et al. 2021I. Our algorithm is related to these knowledge distillation approaches in Federated Distillation, except we also send weights, thereby connecting to non-distillation approaches. We hope future work can explore these connections further.
The communication cost of FedLap-Func per round is almost the same as FedLap: FedLap-Fun has communication cost P + C², where C' is the number of classes (C < P), because we send predictions over points (one point per class in our experiments) and each prediction vector is of length C. On clients, the computation cost is now O((N: + PF), where M is the number of points memory (M < Ng. so the additional cost is small). On the server, the computation cost increases compared to FedLap: the server has to optimise objective function instead of simply adding vectors. It is now O(MPE,). where M is the number points in memory and E, is the number of optimisation steps. This is less than the computation cost at clients (because M < N₂) and will typically not problem (typically the global server large and has more compute than clients).
7

---

Published as a conference paper at ICLR 2025
4 EXPERIMENTS
We run experiments on a variety of benchmarks, (i) on tabular data and image data, (ii) using logistic
regression models, multi-layer perceptrons, and convolutional neural networks, (iii) at different client data heterogeneity levels, and (iv) with different numbers of clients. We focus on showing that (i) our VB-derived algorithm FedLap performs comparably to the best-performing ADMM algorithm FedDyn (despite having one fewer hyperparameter to tune: no local weight-decay), (ii) including full covariances in FedLap-Cov improves performance, and (iii) including function-spant -ppcce information in Fedl Fan-Fac improves performance. Hyperparameters and further details are in App. E
Models and datasets. We learn logistic regression model two (binary classification) datasets:
UCI Credit [Quinlan 1987) and Fl Flaamb-HHaar (Janosi et al. 1988 du Terrail et al.) 2022). Details on the datasets are in App. E] We use the same heterogeneous split UCI Credit as in previous work (Ashman et al. (022). splitting the data into 10 clients (results with a homogeneous split of UCI Credit are in Table 4ii App. F). FLamby-Heart a
split consisting of data from 4 hospitals (du Terrail et al. 2022). We also train a 2-hidden layer perceptron (as in previous work (Acar et al., 2021 McMahan et al. 2016)) on MNIST (LeCun et al. 1998) and Fashion MNIST (FMNIST) (Xiao et al. 2017) with K = 10 clients. We use random 10% of the data in FMNIST to simulate having less data, and use the full MNIST dataset. We consider both a homogeneous split and a heterogencous split. To test having more clients, also
heterogeneous split of (full) Fashion MNIST across 100 clients. Lastly, we train a convolutional neural network on CIFARIO (Krizhevsky and Hinton 2009), using the same CNN from previous work (Pan et al. 2020 Zenke
et (017). We split the data heterogeneously into K = 10 clients. For all heterogeneous splits, we sample Dirichlet distributions that decide how many points per class go into each client (details in
App. E. Our sampling procedure usually gives 2 clients 50% of all the data, and 6 clients have 90% of the data. Within the clients, 60-95% of client data is within 4 classes out of 10.
Methods and hyperparameters. We compare Fedl_ Lap. Fedl ap-CCv and Fedl -ap-Fune with three baselines: FedAvg (McMahan et al. (0166), FedProx (Li et al. (020)) and FedDyn (Acar et al. [021]. which is the best-performing federated ADMM-style method (FedDyn performs much better than FedADMM). We fix the local batch size and local learning rate (using Adam (Kingma and Ba 2015) to be equal for every algorithm, and sweep over number of local epochs, the d and a hyperparameters. for FedDyn the additional local weight-decay hyperparameter, and for FedLap- -une the additional hyperparameter. For Feunn, we assume one randomly-selected point per class per client available to the global server (two points for CIFAR10). This breaks the strictest requirement of not sharing any client data with the global server, however, this is very few points, and it might be reasonable to share a few random points having obtained prior permission. We do not report results of Fedl LedLaa-Func on FLamby- Heart because of the sensitive nature of medical data.
We summarise results in Table I showing the average accuracy (across 3 random seeds for all datasets) after certain number of communication rounds. We provide further results in more tables in mppp.. F where we also report the average number of communication rounds to specific accuracies. They all show similar conclusions.
FedLap performs at least as well as FedDyn across datasets and splits, and both are better than Avd and FedProx. We first compare FedLap with the other baselines. We see that FedLap performs at least as well as FedDyn on all datasets and heterogeneity levels, showing that FedLap similarly strong to FedDyn, despite having one fewer hyperparameter tune (see also results on homogeneous split of UCI Credit in Table 4 in App. F where FedLap performs better than FedDyn). We note that FedDyn's 's performance very sensitive to the value of this additional weight- decay hyperparameter (we provide an example of this hyperparameter's importance in App. I.I. Additionally, we find that Fedl aap performs well with the optimal global weight-decay 8 (we show an example of FedLap's 's improved performance over FedDyn due to this property in App. D.2).
RedLap-Cov significantly improves upon FedLap. We see that FedLap-Cov consistently improves upon FedLap (and other baseline methods including FedDyn). At the highest number of communica- tion rounds in Table FedLap-Cov 's accuracy is higher than FedDyn's 's by 1.7%-55.9% on four of our dataset settings. and is only marginally worse on one (0.2% worse on CIFAR-10. 100 which is within standard deviation; FedLap-Co Covin is also significantly better earlier in training). This empirically shows the benefit of including covariance information to improve upon ADMM-style methods.
8

---

Published as a conference paper at ICLR 2025
Dataset
Comm Round
FedAvg
FedProx
FedDyn
FedLap
FedLap
FedLap
-Cov
-Func
FLambyll
10
72.6(1.1)
77.5(1.5)
76.9(0.5)
77.5(1.1)
79.2(1.8) (12.3)
(heterog)
20
77.6(0.6)
77.7(0.5)
78.1(0.6)
77.7(0.5)
80.0(0.5) (11.)
UCT Credit
77.0(1.3)
73.6(2.2)
73.5(1.0)
77.1(2.3)
80.4(0.4) (17.1)
77.1(1.1)
(heterog)
25
76.3(3.5)
75.0(4.1)
79.9(2.6)
80.7(1.4)
84.2(2.0) (14.3)
83.6(2.7)
50
78.6(3.3)
79.3(0.9)
83.5(2.2)
83.5(1.8)
85,2(1.6) (11.7)
85.8(2.0)
MNIST
97.9(0.0)
97.9(0.1)
98.0(0.1)
98.0(0.0)
98.0(0.1) (0.0)
97.9(0.1)
(homog)
20
98.1(0.1)
98.2(0.1)
98.1(0.1)
98.2(0.0)
98,2(0.1) (10.1)
98.2(0.0)
MNIST
10
97,5(0.1)
97.5(0.1)
97.0(0.2)
97.4(0.1)
97.6(0.1) (10.6)
97.5(0.1)
(heterog)
25
97.9(0.2)
97.9(0.1)
97.8(0.0)
98.0(0.1)
98.0(0.2) (10.2)
98.0(0.1)
50
98.0(0.1)
98.0(0.1)
98.0(0.1)
98.2(0.1)
98.2(0.1) (10.2)
98.0(0.1)
FMNIST
10
72.3(0.4)
72.2(0.3)
75.3(0.8)
72.1(0.2)
75.0(0.6) (10.3)
73.7(0.7)
(homog)
25
77.7(0.3)
77.4(0.1)
77.5(0.8)
77.100.1)
79.8(0.4) (12.3)
77.9(0.3)
50
80.0(0.2)
80.3(0.1)
78.2(0.5)
80.2(0.1)
81.((.1) (12.6)
80.0(0.2)
FMNIST
70.4(0.9)
69.9(0.4)
73.0(0.6)
71.3(0.)
74.6(0.7) (116
72.2(0.9)
(heterog)
25
74.3(0.5)
74.7(0.6)
74.6(0.4)
74.3(0.4)
78.3(1.)
(13.7)
75.4(0.8)
50
76.0(0.7)
76.9(0.9)
74.6(0.5)
77.6(0.7)
80,5(0.6)
(15.9)
78.1(0.7)
FMNIST
10
73.9(0.3)
73.3(0.7)
75.7(0.4)
73.9(0.3)
(6.9(0.7) (112)
76.9(0.8)
(heterog)
25
78.7(0.3)
78.4(0.1)
81.4(0.2)
79.4(0.2)
81.3(0.2) (10.1)
79.7(0.4)
100 clients
50
818(0.3)
81.60.1)
82.2(0.3)
82.4(0.3)
83.0(0.1)
(10.8)
81.8(0.3)
CIFARTO
10
73.8(0.5)
73.8(1.5)
72.7(0.9)
74.8(1.))
75.1(1.1) (12.4)
76.0(1.1)
(heterog)
25
75.0(0.5)
75.0(0.9)
77.4(0.6)
78.0(1.2)
77.6(0.7) (10.2)
78.5(1.0)
50
7..1
75.4(0.8)
79.4(0.4)
79.5(1.4)
79.2(0.9) (10.2)
79.5(1.4)
Table 1: The average accuracy (standard deviation in parentheses) over three runs reported after fixed numbers of communication rounds. We see that FedLap performs at least as well as FedDyn in all settings, while FedLap-Cov significantly improves over FedLap and FedDyn (except for homogeneous MNIST where methods perform equally well): FedLap-Cov's improvement over FedDyn is reported next to FedLap-Cov's accuracy. FedLap-Func also improves performance (especially seen on CIFAR1O and heterogeneous FMNIST). For each run, we use the average accuracy over the previous 3 rounds to account for instabilities (maximum accuracy over previous 3 rounds in Table 2). and bold the top two performing algorithms (even their standard deviations overlap with others).
FedLLa-Fuune improves upon FedLap. We also that FedLap-Func improves upon FedLap, especially in the heterogeneous splits of Fashion MNIST (10 clients) and CIFAR10. FedLap-Func often reaches better accuracies earlier rounds, as uses function-space information while the model is still far from optimal. We note that, in our experiments, we send function spaae information over very few inputs (one randomly-chosen point per class per client, or two points with CIFARIO), Future work can look at using public unlabelled datasets to send function-space information over many more inputs, like in Federated Distillation. In App. E.5 we show that performance improves when we send function-space information over more datapoints (on FMNIST), as expected.
We see the same conclusions with 100 clients instead of 10 clients. We also run heterogenous setting with 100 clients (instead of 10) to simulate having many more clients (results in Taale T). We do this on FashionMNIST (full dataset), using the same method to heterogenously split data across clients as with 10 clients. We see the same conclusions with 10 clients: FedLap performs as well as FedDyn (and better than FedAvg and FedProx) despite having one fewer hyperparameter to tune, Lev significantly improves performance, and FedLap-Fumm improves performance particularly earlier in training.
5 CONCLUSIONS AND FUTURE WORK
In this paper, we provide new connections between two distinct and previously unrelated federated learning approaches based on ADMM and Variational Bayes (VB), respectively. Our key result shows that the dual variables in ADMM naturally emerge through the site parameters used in VB. We first show this for the isotropic Gaussian case and then extend it to multivariate Gaussians. The latter is used to derive new variant of ADMM where learned covariances incorporate uncertainty and

---

Published as a conference paper at ICLR 2025
used preconditiones. We also derive new functional regularisation extension. Numerical experiments show that these improve performance.
This work is the first to show such connections of this kind. No prior work has shown the emergence of dual variables while estimating posterior distributions. The result is important because it enables new ways to combine the complementary strengths of ADMM and Bayes. We believe this to be especially useful for non-convex problems, such as those arising in deep learning.
There are many avenues for future work. We believe this connection holds beyond the federated learning setting, and we hope to explore this in the future. This could give rise to a new set of algorithms that combine duality and uncertainty solve challenging problems in machine learning, optimisation, and other fields. There are also several extensions for the federated case. Our current experiments update all clients every round of communication, and future experiments could relax this. We expect that our methods will still work well, just like algorithmicall--eeattee algorithms such as FedDyn. Future work can also analyse convergence rates of FedLap, following similar theoretical assumptions to those for FedDyn and FedADMM. We could also explore differentially-priee versions of the algorithm, following bifferentially-privanl versions of PVI (Heikkilä et al. 2023). We also expect that our method will perform well in continual federated learning, using Bayesian continual learning techniques (Kirkpatrick et 2017] Pan et al., 2020 Titsias et al. 2022). Lastly, future work explore connections between Federated Distillation and FedLap-Func.
ACKNOWLEDGMENTS
This material is based upon work supported by the National Science Foundation under Grant No. IIS-2107391. Any opinions, findings, and conclusions or recommendations expressed this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation. MEK supported by the Jayes-Dualij project, JST CREST Grant Number JPMJCR2112.
AUTHOR CONTRIBUTIONS STATEMENT
List of authors: Siddharth Swaroop (SS), Mohammad Emtiyaz Khan (MEK), Finale Doshi-Veler (FDV).
SS led the project and conceived the original ideas with regular feedback from FDV and MEK. The emergence of dual variables through site parameter hypothesised by MEK, and confirmed by SS. SS coded and ran all experiments. SS wrote first draft of the paper with feedback from FDV, which was substantially revised by MEK with feedback from SS and FDV.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based dynamic regularization. In International Conference on Learning Representations, 2021. URL Id-B7v4QMR622W
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Confer- ence on Learning Representations, 2021. URL petttp:://opnzeviiee new forummirriio
GFsU8aUsGB
Matthew Ashman, Thang D Bui, Cuong V Nguyen, Stratis Markou, Adrian Weller, Siddharth Swaroop, and Richard E Turner. Partitioned variational inference: A framework for probabilistic federated learning. arXiv preprint urXiv:2202.12275, 2022.
Thang D. Bui, Cuong V. Nguyen, Siddharth Swaroop, and Richard E. Turner. Partitioned variational inference: A unified framework encompassing federated and continual learning. arXiv preprint 1811 1.1 1206, 2018.
Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, and Simone Calderara. Dark
experience for general continual learning: strong, simple baseline. In Advances in Neural Information Processing Systems, volume 33, pages 15920-15930, 2020.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Bayesian deep learning. In Advances in Neural Informa-
Philipp Hennig. Laplace fion Processing Systems, 2021.
Erik Daxberger, Siddharth Swaroop, Kazuki Osawa, Rio Yokota, Richard E Turner, José Miguel Hemández- Lobato, and Mohammad Emtiyaz Khan. Improving continual learning by accurate gradient reconstructions the past. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. htt.hhtss:// htttppss///peereeeeeeee www.oeeieeev...oo
net. forum?idee d=bb1 fpfCijal f
Wei Deng, Qian Zhang, Yian Ma, Zhao Song, and Guang Lin. On convergence of federated averaging langevin dynamics. In The 40th Conference on Uncertainty in Artificial Intelligence, 2024. URL
htttss://ppenrev
net
Jean Ogier du Terrail, Samy-Safwan Safwan Ayed, Edwige Cyffers, Felix Grimberg, Chaoyang He, Regis Loeb, Paul Mangold, Tanguy Marchand, Othmane Marfoq, Erum Mushtaq, Boris Muzellec, Constantin Philippenko, Santiago Silva, Maria Teleñczuk, Shadi Albarqouni, Salman Avestimehr, Aurélien Bellet, Aymeric Dieuleveut, Martin Jaggi, Sai Praneeth Karimireddy, Marco Lorenzi, Giovanni Neglia, Mare Tommasi, and Mathieu Andreux. Flamby: Datasets and benchmarks for cross-silo federated lcarning in realistic healthcare settings. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. URL https
//openrview.. net. / Forum? wtw.mmmdd=GGMDDDDGG AbbG22
Hugh Durrant-Whyte. Data fusion in decentralised sensing networks. In Fourth International Conference on Information Fusion, 2001, 2001.
Yonghai Gong, Yichuan Li, and Nikolaos M. Freris. Fedadmm: A robust federated deep learning framework with adaptivity to system heterogeneity. In 2022 IEEE 38th International Conference on Data Engineering (ICDE), pages 2575-2587, 2022. doi: 10. (0.1109/CDDE557422222.0000
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning as variational inference: scalable expectation propagation approach. In The Eleventh International Conference on Learning Representations, 2023. URL httttss////peeeeeiee
Mikko A. Heikkilä, Matthew Ashman, Siddharth Swaroop, Richard E Turner, and Antti Honkela. Dif- ferentially private partitioned variational inference. Transactions Machine Learning Research, 2023. ISSN 2835-8856.
Andras Janosi, William Steinbrunn, Matthias Pfisterer, and Robert Detrano. Heart Disease. UCI Machine Learning Repository, 1988. DOI: htps://oii.ng/1122333//54444
11

---

Published as a conference paper at ICLR 2025
Mohammad Emtiyaz Khan and Wu Lin. variational inference in non-conjugate models to inferences conjugate models. In International Conference Artificial Intelligence and Statistics, pages 878-887 2017.
variational inference: converting
Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural- natural-gradient descent for variational inference in complex models. In 2018 International Symposium Information Theory and Its Applications (ISITA), pages 31-35. IEEF, 2018.
Mohammad Emtiyaz Khan and Havard Rue. The Bayesian learning rule. Journal of Machine Learning Research, 24(281):1-46 2023.
Mohammad Emtiyaz Khan and Siddharth Swaroop. Knowledge-adaptett priors. In Advances in Neural Information Processing Systems, volume 34. Curran Associates, Inc., 2021.
Mohammad Emtiyaz Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into Gaussian processes. Advances in neural information processing systems, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference Learning Representations, 2015.
James Kirkpatriek, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114 (13):3521-3526, 2017.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Yann LeCun, Corinna Cortes, and Christopher JC. Burges. The MNIST database handwritten digits. 1998. URL htttht//////nnn. lecun. com/ wxwb/mnissmmi////ec
Daliang Li and Wang Fedmd. Heterogenous federated learning via model distillation, 2019.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization heterogeneous networks. I. Dhillon, D. Papailiopoulos, and V. Sze, editors, Proceedings of Machine Learning and Systems, volume 2, pages 429- 450, 2020. URL http://pss//pooceeeiingni..//i/p Fi File:/BBfee399333333237777777713
files Fites//pppee//22////
Jajuun Liang, Qian Zhang, Wei Deng, Qifan Song, and Guang Lin. Bayesian federated learning with Hamiltonian Monte Carlo: Algorithm and theory. Journal of Computational and Graphical Statistics, pages 1-19, 07 2024, doi: 10.1080/1061 T101///066600020200 8600.2024.2380055
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 2351-2366 Curran Associates, Inc., 2020a. URL
www.eddiii.nnnnnnnneeeee.......e ccpppppppccpp.....eeppppppppppe. cctpeeerrreeeeee pdd
Wu Lin, Mark Schmidt, and Mohammad Emtiyaz Khan. Handling the positive-definite constraint in the Bayesian learning rule. In Proceedings of the 37th International Conference Machine Learning, volume 119 Proceedings of Machine Learning Research, pages 6116-611 6126. PMLR, 13-18 Jul 2020b. URL htttps://proces
ngs. www.smmww....r.....mmmmmmmm.m... / 122 html
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
2020.
H. B. McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Ar- cas. Communication-efficin learning of deep networks from decentralized data. In Inter- national Conference on Artificial Intelligence and Statistics, 2016. URL attps://aap semant cschol wwwhhollrrrrg/g/// org.hoaa....ooooooopossa11111115
iemanticscholl
12

---

Published as a conference paper at ICLR 2025
T. Minka. Expectation propagation for approximate Bayesian inference. In Proceedings of the Conference on Uncertainty Artificial Intelligence, 2001.
Arthur G. O. Mutambara. Decentralized estimation and control for multisensor systems. Routledge, 1998.
Pingho Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard Turner, and Mo- hammad Emtiyaz Khan. Continual deep learning by functional regularisation of memorable past. In Advances in Neural Information Processing Systems, volume 33, pages 4453 4464. Curran Associates, Inc., 2020.
J. R. Quinlan. https//d6olor/110224422////
Credit Approval.
UCI Machine Learning Repository, 1987. DOI:
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7): (117))1723-1333 2002.
Steven L. Scott, Alexander Blocker, Fernando V Bonassi, Hugh Chipman, Edward I George, and Robert E McCulloch. Bayes and big data: The consensus Monte Carlo algorithm. In Big Data and Information Theory, pages 8-18. Routledge, 2022.
Hyowoon Seo, Jihong Park, Seungeun Oh, Mehdi Bennis, and Seong-Lyun Kim. Federated knowledge distillation, 2020.
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Möllenhoff. Variational learning effective for large deep networks. In Forty-first International Couference on Machine Learning. ICML 2024, Vienna, Austria. July 21-27, 2024. OpenReview.net, 2024. URL 7GKvk]
httpp:://ppenrett ew net
Michalis K. Titsias, Jonathan Schwarz, Alexander G. de G. Matthews, Razvan Pascanu, and Yee Whye Teh. Funetional regularisation for continual learning with gaussian processes. In International Conference on Machine Learning, 2020.
Volker Tresp. A Bayesian committee machine. Neural Computation, 12(11):271--2741 2000.
Han Wang, Siddartha Marella, and James Anderson. Fedadmm: A federated primal-dual algorithm allowing partial participation. 2022 IEEE 61st Conference on Decision and Control (CDC), pages 287-294, 2022. URL htt.::ttpp.......... app. semant wwt.amtncccchhla... wtt.aaarrrrr/ooouusoooaoooo http://ooou11:::278999009 247778987
Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Communication efficient federated learning via knowledge distillation. Nature Communications, 13, 2021. URL
https://api www.nticscootooooo.c.o wwt.chholl...../oooouuuseeee wttp:ooroo//...////33 1D:
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017. URL htttp://arxiv. org/ Dataset freely available at
abs/1708.07747 cite arxív: 700.0077770ommm
mnist Benchmark is available at http:/ffashion-miitt.s.
website.eu-centrall-1 wtw.://wc.....aaaaaaaaaaaao.aooo www.zoaww.ccom...............
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence.
In International Conference on Machine Learning, pages 3987-3995. PMLR, 2017.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: A federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055-6070, 2021. doi: 10. 10L.1099T2222..1111515
Shenglong Zhou and Geoffrey Ye Li. Federated learning via inexact ADMM. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(8):9699-9788 2023. doi: 10.1109/TPAMI.2022. 3243080.
13

---

Published as a conference paper at ICLR 2025
A MAKING FEDLAP MORE SIMILAR TO THE ADMM UPDATE
In this section, we re-write the FedLap update (... 9 look more similar to the ADMM update Eqq 2). Specifically, we divide by Nk, and make the substitution üy = ÓU,/Nk, giving,
Client updates:
Wk f- arg min ls (wr) + U w + (1//N)) w
w
bp t Ü2 + p(8/Ns) wg)
(18)
k
Server update:
Wg
(5/Nk)
o*.
B DERIVATION OF FEDLAP-COV UPDATES
Here, we give full derivation of the update given in Eq. 14 We use the definition of the KI. divergence for two multivariate Gaussians = Nww m. m. S) and 4g = NNww; g. Sy):
D [g = m wy 1|3, + Tr (S,5
log S + const.
We plug this and Eq. 13 in the first update given in Eq. 12 to get the following.
(Wk, Sk) < arg min (a(m) + Tr [H (ma S v Tr [v (mm² + + m.5
(19)
+ f ll|ا + (S 2 log Sys
The Vs dual term above (the fourth term) is derived using the following identity,
E w V.wl (r (Veww" Tr VVE., (ww T) Tr V& (mm² +E)]
With this, the updates over wa and S₁ decouple into two different updates:
Wi < arg min (m) + Ur m m Vam + 3Imm 1m--wyls.. wyllss m
(20)
S arg min 1TT (m S Tr [V.S Tr (S, log 1,,5
By replacing m with w, we recover the update for we Eq. 14 The update of S₁ has a closed form solution, as shown below by taking the derivative with respect to S-1 and setting it to 0,
(m - V + S, - - =0
S H4 (mala) V Sy.
In the main update, moldd to be the most recent value of my before updating it. This ensures that the Hessian is estimated at the most recent parameters and so is the best estimate possible.
Using this derivation and the additional details in Sec. 3.3
write final FedLap-Cov update as,
Client updates: Wk arg min Eg ( + 1h un Viw 1 1| w
S HA (ww) Vv + Sgy
Uk - + p ppSSpppSn (Saw: S,wa) and V4 + (1- - V2 + pH₂(wx),
K
Server updates:
wg S
Ubs
where
S = 51 +
Vz.
(21)
1=1
k=1
C DERIVATION OF FEDLAP-FUNG EQUATIONS
In this section derive the Fedl-ap-uune equations from Sec. 3.4
14

---

Published as a conference paper at ICLR 2025
C.1 DERIVATION OF EQ.
At the global server, we add in the true loss over inputs in Mr (temporarily pretending as we had true labels available), and then subtract the Gaussian approximate contribution over these points from Khan et al. (2019),
Σ Elyi, w) - A0, - J. w)
(22)
M.
where the second expression comes from Khan et al. (011)) They take Laplace approximation of posterior at wi, and show that this Laplace approximation is equal to the posterior distribution of the linear model In = ur Jn + 6n after observing data D {J,, De JED, where D: = _w J; A (hu,s - yi), en ~ NO,, A, Ar,s is the Hessian of the loss for point at parameters wK, and J; the Jacobian input and weights wp (this is equal to the input features a Generalised Linear Model).
Substituting in the values of Ht. and noting that the Generalised Gauss Newton approximation the Hessian is Ma 2 LiEM, Ai, gives us that,
Σ REM&
Elyi, w) AA,,a00 (/w)²=
Σ REML
Eli w)
f( wE)
w
weight- space
function-space information
(23)
where i are soft labels (predictions) over inputs E Mr using local client weights wi, and HMM is the (approximate) Hessian of the loss over points in Mr at weights wy. Adding this to Eq. 15 completes the derivation of Eq. 16
C.2 DERIVING THE FEDLAP-FUNC ALGORITHM
We derive the FedLap-Fund equations (Eq. 17) from Eqq. 15 a5]and We note that we can derive combined version of FedLap-Cov and FedLap-Funde but for simplicity we will focus the non-covariance case.
We start with our approximation that all covariances are set to őI, just like in FedLap: - S₁ SI. We note that normally, Gauss-Newton approximation). By making this approximation, information, setting all H = O. This means the contribution from client to the server (Eq. 16) becomes,
H + òI, where is the Hessian over points (eg the Generalised
are ignoring this second-ordes
Σ e(ir, w) M
+ 50y w,
(24)
function space information
where the new weight-space term has been removed as M. = 0. We have simply added the
function-space information, where ÿJ is the prediction over inputs E M, at model parameters ur
We include this function- space term in both the client update for wi and the global server update for wy- We illustrate this with the global server update. Using App. C.T. gives us,
(giigg w) (w wgg H w(w wp) + 1 (w TT,TTS,(C Sy(w wT) ÓI, (25)
KEIiEM.
new weight-space term
unchanged from before
function-space information
where 0i.g is the prediction over inputs i E M₁ using model parameters wy from the global server, and Ma is the Hessian over points in M2 at parameter Wy. In practice, we can approximate this Hessian with Gauss-Newton approximation, and can either recalculate it the global server before sending back to local clients, or approximate it as the Hessian sent from the client (calculated at wy)) M. xv M
15

---

Published as a conference paper at ICLR 2025
the two weight-space terms have the same (w wgy quadratic form. Therefore we can straightforwardly consider this as changing the global server precision S to remove the contribution from i MM, vk. When take the approximation that all S ÓI, again set the Hessians to all be O, and this simplifies the equation further. Plugging this into the first client update in [Bq. 9 and multiplying all function-space terms by hyperparameter т, gives us final FedLap-Fune algorithm.
Eq. 25 is simpler than Eq. 16
D BENEFITS OF FEDLAP OVER FEDDYN AND FEDADMM
In this section, we show two simple experiments showcasing why FedLap has better properties than FedDyn and FedADMM. First, in App. D.I we show on homogoneous MNIST split that FedDyyn requires an additional weight-decay term the local loss during training in order for it to perform well, unlike Fedlap. Second, in [App. 1..2 we show how Fedl lap targets a better global loss (with weight-decay incorporated already), and gets closer as the number communication rounds increases, unlike FedDyn and FedADMM, which target worse global optimum (with no weight-decay the global loss).
D.1 IMPORTANCE OF LOCAL WEIGITT-EECA IN FEDDYN
In this section, we show the importance of the additional weight-decay parameter to FedDyn's performance on simple example (homogeneous MNIST), hyperparameter that FedLap removes. Fig. shows the results. We first see that FedDyn (with optimal weight-decay setting) and FedLap perform similarly. However, FedDyn's performance is highly dependent on the setting of its weight- decay hyperparameter: changing by order of magnitude in either direction significantly degrades performance. Removing it entirely results in catastrophic performance after the first communication rounds. In contrast, changing the ó hyperparameter in FedLap (or FedDyn) by order of magnitude does not reduce performance significantly this setting.
98
Accuracy (96
Led
Dyn (optimal weight-decay)
Dyn (large weight-decay)
Dyd (small weight-decay)
Dyd (no weight-decay)
40
Communication rounds
Figure 1: FedDyn's performance is sensitive to the additional weight-decay hyperparameter that has compared to FedLap. When set an order of magnitude too small or too large, performance drops significantly. When there is weight-decay, performance gets worse rounds. Results mean and standard deviation over three seeds on homogeneous MNIST (10 clients).
D.2 FEDLAP TARGETS A BETTER GLOBAL LOSS THAN FEDDYN/FEDADMM
In this section, we present a simple setting showcasing the potential benefits of FedLap over FedDyn (and FedADMM). We take the UCI Credit dataset (a binary classification task), split the data equally (and homogeneously) into just two clients, and train a logistic regression classifier, ensuring we train until convergence both locally (full-batch, Adam learning rate 0.001, 1000 local epochs/iterations) and for a large number communication rounds (1000 communication rounds). FedADMM is only consider FedDyn. Note that in this
always slower and worse-performing than FedDyn,
16

---

Published as a conference paper at ICLR 2025
section, do not include weight-decay in the local optimisation of FedDyn, because FedDyn's theory does not include this, and this setting is very simple.
First, we note that ideal global performance (of w;)) is better with å 1 (giving test nll of 0.337) than = (test nll 0.420). Therefore the w that FedLap targets is already a better target than that of FedDyn (and FedADMM), which targets global optimum with d = 0.
Secondly, FedLap is quicker to target its ideal global parameter than FedDyn. FedLap reaches within 1% of its ideal w" train loss within just 2 communication rounds, and continues getting closer as we train for longer (up to the 1000 rounds). This happens even when setting different values of (meaning that Fedlap might be targeting unideal values of w.))
On the other hand, how close FedDyn gets to its own (worse) ideal train nll depends on the value of its hyperparameter a: the best value takes 78 rounds to get to within I 1% of its optimal train nll. Larger values of the hyperparameter do not reach within 1% within 1000 rounds: they lead to quick improvement earlier in training, but then performance (and train nll) gets worse over time.
E ADDITIONAL EXPERIMENTAL DETAILS
In this section we provide further details on experimental setup, including hyperparameter values.
For Fp-Fna, we set a datapoint-specificc value of r, making it depend on the number of datapoints from that class in the client: for datapoint i belonging to class c, we have the number of datapoints in the class e in client k; and MR.e is the number of datapoints in class in memory M, (usually 1). We then perform a hyperparameter sweep over Ty. In this way, we multiply the contribution of each datapoint by how many datapoints is representing from class in client k.
NE.e, where Ng.c is Tf Me...
F. HETEROGENEOUS SAMPLING PROCEDURE FOR MNIST, FASHION MNIST AND CIFARIO
Similar to previous work, we use two Dirichlet distributions to sample heterogeneous data splits. We first sample from Dir(aa)) where ar is of length number of clients K, to sample client sizes. We then sample a separate Dir(O:2) within each client, where ar is of length number of classes, to sample class distribution within a client. We then multiply the sampled class distribution within each client by the sampled client size. When assigning datapoints from a specific class to clients, we normalise these values across clients, so that all datapoints are assigned between the clients.
We use 0 and 22 = 0.5 in our heterogeneous splits on MNIST, Fashion MNIST and CIFAR10 Our sampling procedure usually gives 2 clients 50% of all the data, and 6 clients have 90% of the data. Within the clients, 60-95% of client data is within 4 classes, We use the same a values for the 100-client case.
E.2
UCI CREDIT
In UCI Credit, the binary classification task to predict whether individuals default on payments, and there are a total of 520 training points (with 45% positive labels), with an accuracy of 90% (when global weight-decay is not zero). We use the heterogeneous split from Ashman et al. (2022) in the main text: data is split into K 10 clients, 5 of which have 36 datapoints each (positive label rate of 6%), and the other 5 have 67 datapoints each (positive label rate of 66%). We provide results for homogeneous split over K- = 10 clients in Table 4 where we see similar conclusions (there is an even bigger improvement of Fedl.ap and Fedl ap-relate methods).
We set local Adam learning rate 10-3, and minibatch size (ensuring there is gradient noise always). All methods have a hyperparameter sweep over number of local epochs = 5, 10. 20). We perform a sweep over the a or hyperparameter for FedProx, FedLap, FedLap-Cov and FedLap-Func 10 10-51, $1, and local weight-decay sweep is over [10 ,,10 10005 10 10061 8). For FedLap and FedLap-Fuum Func, we set damping factor = N₂/N, and for FedLap-Coo p 1/K. For RedLap-Fune we 7 = 1, and run global server optimisation for 5000 iterations with a learning rate of 0.001. Each run took up to 4 minutes on a standard laptop CPU.
over [10, 1, 10 For FedDyn the sweep is over [10-3, -,0 10
17

---

Published as a conference paper at ICLR 2025
E.3 FLAMBY-HEARD
FLamby-Hees Heart is a learning benchmark, and consisting of data from hospitals to predict whether a patient has heart disease (binary classification). There are 486 training examples.
split, proposed by du Terrail et al. (202)) as a federated
We set batch size to 4 (like in du Terrail et al. (2022), but use the Adam optimiser locally. We perform hyperparameter sweeps over number of local epochs [1, 5, 10), and (Adam) local learning rate = [10- 1.. For FedProx, FedLap, and FedLap-Cov, we perform a sweep over (or a) = [100, 10, 1, 1001 10 10-21. For FedDyn, this is over = [10, 1, 10-1 10
100 and FedDyn's local weight-decay sweep is over [0, 10 10 10 *]. For FedLap, we set damping factor = NN/N, standard laptop CPU.
and for FedLap-Cov, P 1/K. Each run took less than I minute
Note that we do not report results of FedLap-Fuund Sun-Func&& as FLamby- heaart has sensitive data (medical data), where it is not likely to be reasonable share this data with a global server.
E.4 MNIST
We use the full MNIST dataset. Batch size is set to 32. We train a two-hidden layer multi-layer perceptron, with 200 hidden units in the first layer, and 100 units in the second layer (the model gets 98.3% accuracy on all data).
For all methods, perform hyperparameter sweeps over number of local epochs [1, 5, 10), and (Adam) local learning rate [10 10 10 For FedProx, FedLap. FedLap-Cov and FeddLap-Fnn, we perform sweep
8 (or a: [1, 10 10- For FedDYn, this is a [10-1 10 10-9) and FedDyn's 's local weight-decay sweep is over [10-4 10-5 10-0). J. For FedLap and FedLap-Fune. we set damping factor p / and for FedLap-Cov, p = 1/K. For FedLap-Fune, we sweep over Ty [1, 10 10- and run global server optimisation for 3000 iterations with a learning rate of 0.0005. Each run took up to 2 hours on standard laptop CPU.
E.5 FASHION MNIST
We use a random 10% of Fashion MNIST every random seed, and average across (the same) three random seeds. Batch size is set to 32. We train two-hidden layer multi-layer perceptron, with 200 hidden units in the first layer, and 100 units in the second layer (the model gets 85% accuracy). Local Adam learning rate is set 10-3
For all methods, we perform hyperparameter sweeps over number of local epochs = [1, 5, 10]. For FedProx, FedLap, FedLap-Cov and FedLap-Fune, we perform sweep over (or a) [1. 10000-11000 1000'-10-2 10- 10-3, 1003310-4)1 10-4]. For FedDyn, this is over a = (10-2 10-3, 1004 10 10-%) and FedDyn's 's local weight-decay sweep is over [10 110-4,10 10 10 10060 J. For FedLap and FedLap-Func, we set damping factor p N2. J//N, and for FedLap--ovv p = 1/K. For
we sweep over Ty = [1, 10_1, 10 10-3 10-1), and run global server optimisation for 5000 iterations with a learning rate of 0.001. Each run took up to I hour 20 minutes on a standard laptop CPU.
FedLap-Func: sending information over more datapoints. If sending information over 2 points per class per client (instead of 1), FedLap-Fuum improves performance. In the homogeneous setting, mean accuracies (compare with Table T are 74.0(0.7) at round 10. 77.8(0.4) at round 25, (0.1( at round 50. Max accuracies (compare with Table 2) are 74.4(0.6) at round 10, 78.2(0.2) at round 25, 80.3(0.4) at round 50. 75% accuracy is reached in 12(2) communication rounds, and 78% in 25(2) rounds (compare with Table 3). In the heterogeneous setting, mean accuracies (compare with Table 1) are 73.2(0.1) at round 1O, 76. (0.6) at round 25, 78.4(0.2) at round 50. Max accuracies (compare Table 2) are 73.7 73.7(0.3) at round 10, 76.7 (0.7( round 25, 78.8(0.2) at round 50. 75% accuracy is reached in 16(2) communication rounds, and 78% in 37(5) rounds (compare with Table 3).
Table
E.6 CIFAR-10
Batch size is set to 64. We set Adam learning rate 10-3 (note that 10-2 fails learn anything).
For all methods, we perform hyperparameter sweeps over number of local epochs = [5, 10, 20). For FedProx, perform a sweep over
[1, 10 10-2, 10-3). For Fedlap, Fedl FedLap-Cov
18

---

Published as a conference paper at ICLR 2025
and FedLap-Fuune perform sweep over 5 [1,10-1, 10- For FedDyn, this is over [10 3,-3,, 30000000 10 and FedDyn's 's local weight-decay sweep is over [10 9,100 ,010,1110,00 10.
For FedLap and FedLap-Fune, damping p N&/N for first 10 rounds of training, then 1/K after that, and gradient clipping was used to stabilise training. For FedLap-Func, we sweep over Ty = [10-2 10-1 10-%), and global server optimisation for 500 iterations with learning rate 0.0005. We assume that each client has 2 randomly-selected points per class shared with the global server (all other dataset/benchmark settings have I point per class per client). Each run took up to 10
hours on an A100 GPU.
F ADDITIONAL EXPERIMENTAL RESULTS
In this section we provide some more results from our experiments. We see the same conclusions / results as from Table in the main text.
Table reports the maximum accuracy over the previous three communication rounds (instead of the mean accuracy), and averages this across runs. We see the same results and conclusions as from Table 1 FedLap performs similarly to FedDyn across all dataset and heterogeneity settings, and FedLap-Cor Cov and FedLap-Fune improve upon this. [Table 3 shows the number of communication rounds to a desired accuracy. Again, we see the same takeaways/ / conclusions.
We provide additional results on homogeneous UCI Credit split over 10 clients in Table Here, Fedl.ap ap performs even better than FedDyn consistently, because it is targeting better global loss (with non-zero weight-decay).
In Table 5) we report average test negative log- likelihood results (instead of accuracy) at the end of training. We see that on average, FedLap performs at least as well as FedDyn (sometimes significantly outperforming FedDyn, such as on FMNIST (homog), FMNIST (heterog) and CIFARI0). FedLap- Cov and FedLap-Fing improve even further, as might be expected due to the Bayesian interpretation.
Dataset
Comm Round
FedAvg
FedProx
FedDyn
FedLap
FedLap
FedLap
-Cov
-Fune
MNIST homog
10
98.0(0.0)
98.0(0.1)
98.1(0.1)
98.0(0.0)
98.1(0.1)
98.0(0.1)
20
98.2(0.0)
98.3(0.1)
98.2(0.0)
98.2(0.0)
98.3(0.0)
98.2(0.0)
heterog
10
97.6(0.1)
97.6(0.1)
97.1(0.2)
97.5(0.1)
97.7(0.2)
97.7(0.1)
25
97.9(0.1)
97.9(0.1)
98.0(0.1)
98.0(0.1)
98.0(0.2)
98.1(0.1)
50
98.0(0.1)
98.1(0.1)
98.1(0.0)
98.2(0.1)
98.2(0.1)
98.1(0.1)
FMNIST homog
10
73.6(1.0)
73.4(0.3)
75.8(0.4)
73.1(0.2)
76.0(0.2)
74.5(0.5)
25
78.2(0.3)
77.9(0.1)
78.0(0.8)
77.5(0.2)
80.2(0.3)
78.6(0.5)
50
80.3(0.1)
80.5(0.1)
78.5(0.3)
80.4(0.1)
82.0(0.2)
80.4(0.1)
heterog
71.3(0.7)
71.2(0.7)
74.0(0.3)
72.2(0.6)
75.4(0.9)
72.7(0.9)
25
74.7(0.5)
75.1(0.3)
75.1(0.8)
74.9(0.8)
78.6(1.0)
75.7(0.8)
50
76.4(0.5)
77.4(0.8)
75.0(0.6)
78.1(0.6)
80.9(0.5)
78.5(0.7)
CIFARIO heterog
10
74.0(0.5)
74.2(1.4)
73.2(0.8)
75.3(1.))
75.4(1.))
76.6(1.5)
25
75.1(0.5)
75.1(0.8)
77.6(0.7)
78.2(1.1)
77.8(0.8)
78.6(1.0)
50
75.3(0.3)
75.6(0.7)
79.6(0.4)
79.7(1.3)
79.3(0.8)
79.6(1.3)
Table 2: Mean accuracy (standard deviation in parentheses) over three runs (except CIFAR-10), 10), after a fixed number of communication rounds. For each run, we use the maximum accuracy over the previous 3 rounds (mean accuracy over previous rounds in Table Tp. We bold the top two performing algorithms (even if their standard deviations overlap with others). We see that FedLap performs at least as good as FedDyn in all settings (with bigger difference when the global model performs better with a non-zero weight decay, like in UCI Credit). Fedl ap-Cov significantly improves upon Fedl.ap in most settings (except for homogeneous MNIST, where all methods perform equally well), and FedLap- Func also improves performance (especially seen on CIFAR10).
19

---

Published as a conference paper at ICLR 2025
Dataset
Acc
FedAvg
FedProx
FedDyn
FedLap
FedLap
FedLap
(%)
-Cov
-Func
MNIST homog
97.5
6(1)
6(1)
4(1)
5(1)
4(1)
6(1)
heterog
97.5
10(2)
10(2)
18(3)
11(1)
9(1)
9(1)
FMNIST homog
75
I5(1)
15(2)
9(1)
15(1)
9(1)
12(2)
78
25(1)
26(3)
17(3)
28(1)
15(1)
24(1)
heterog
75
26(4)
22(5)
24(7)
24(4)
11(3)
20(2)
78
22(6)
CIFARIO heterog
72
6(1)
6(3)
8(1)
7(1)
6(1)
5(1)
14(2)
10(3)
10(3)
8(2)
78
30(5)
32(11)
32(11)
22(10)
Table 3: Mean number of communication rounds (standard deviation in parentheses) over three runs, to reach desired accuracy. If every run does not reach desired accuracy within 50 communication rounds, number is reported. We bold the top two performing algorithms (even if their standard deviations overlap with others). We see that Fedl ap performs similarly to FedDyn in all settings. Fedl Fedl .ap-Cov significantly improves upon FedLap in most settings, and Fedl.a-- aunn-un also improves performance (especially seen on CIFAR10).
Dataset
Comm Round
FedAvg
FedProx
FedDyn
FedLap
FedLap
-Cov
UCT Credit
10
81.4(0.6)
80.5(
81.8(1.4)
83.6(0.4)
84.4(0.5)
(homog)
25
84.6(1.1)
84.2(0.4)
86,3(1.5)
87.0(0.4)
86.5(1.0)
50
86.8(0.6)
87.4(0.1)
86.9(0.5)
88.0(0.6)
87.7(0.7)
100
88.7(0.2)
88.5(0.6)
87.9(0.3)
88.2(0.2)
89.1(0.7)
Table Mean accuracy (standard deviation in parentheses) over three runs, after a fixed number of communication rounds, homogoneous split of UCI Credit into K = 10 clients. For each run, we use the average accuracy over the previous rounds to account for instabilities, and bold the top two performing algorithms (even if their standard deviations overlap with others). We see that FedLap and FedLap-Cov perform the best.
FedAvg
FedDyn
FedLap
FedLup .Cow
FedLap -Func
Round
MNIST
(boroog)
0.11(0.01)
0.07(0.04)
0.06(0.00)
0.09(0.00)
0.09(0.01)/10.03)
0.07(0.00)
MNIST
(heterog)
50
0.13(0.01)
0.12(0.01)
0.07(0.01)
0.06(0.00)
0.0770.00x0.00
0.96(0.00)
FMNIST
(homog)
50
0.56(0.01)
0.56(0.01)
0.63(0.01)
0.57(0.01)
0.56(0.00)(1..07)
0.57(0.00)
FMNIST
(heterng)
50
0.67(0.04)
0.66(0.03)
0.77(0.05)
0.65(0.05)
0.60K0.04)100.1))
0.62(0.02)
(beterog)
100 clients
0.53(0.01)
0.53(0.01)
0.51(0.01)
0.530.01)
0.4900.11)1000)
0.52(0.01)
CIFARIO
(belerog)
1.5(0.1)
1.5(0.))
0.7(0.0)
0.600.0)
0.600.0X10.1
0.6(0.0)
Table 5: The average validation negative log-likelihool (standard deviation in parentheses) over three runs reported after few fixed number of communication rounds. Lower is better. We see that on average, FedLap performs at least as well as FedDyn (sometimes significantly outperforming FedDyn, such FMNIST (homog), FMNIST (heterog) and CIFAR10). FedLap- Cov and FedLap-Func improve even further, as might be expected due to the Bayesian interpretation.
20


Reviews:
[{"review_id": "ozx4oC9ghc", "reviewer": "Reviewer_QNup", "summary": "The work provides a new link between two well-known federated learning approaches: i) Alternative Direction Methods of Multiplier (ADMM) and Partitioned Variational Inference (PVI). The driving idea is that PVI updates have a line-by-line correspondence to the ADMM algorithm, which is even more similar under the assumption of isotropic Gaussian distributions for the family $Q$ of variational densities. In this direction, the authors introduce FedLap, a Laplace-approximation based version of PVI. Such derivation suggests that there exists an almost identical duality between ADMM and the common Bayesian approach to the federated learning problem. Later, having established such a connection, two more variants of FedLap are introduced, the first one (FedLap-Cov) puts a bit more structure in the candidate set Q, such some improvement is observed given some increment in the computational cost. The second variant (FedLap-Func) extends ideas from (Khan & Swaroop, 2021) by allowing the FedLap update of the global server to be an optimization problem and assuming that some inputs are available between local and global servers. Empirical results confirm the hypotheses made and the improvement in the accuracy for the three new FedLap-based methods.", "questions": "Some minor questions that I thought about while reading and writing the review:\n\n**[Q1]** - I'm curious to know where the term ‘sites’ in L112 is coming from\n\n**[Q2]** - How does this delta approximation work in L185? How does this approximation affect the general problem?\n\n**[Q3]** - What is the dimension of the second auxiliary/dual variable? I'm always afraid of the use of preconditioning.. could the authors add a bit more details in that regard\n\n**[Q4]** - I didn't understand very well what is indicated in L309\n\n**[Q5]** - L342: M_k = null space, Isn’t this obvious?\n\n**[Q6]** - For Algorithm in Eq. 17. How many training steps per round and updates?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**[S1]** - The work is clear, well-written, and polished. I particularly liked the spirit of first building the connection between ADMM and PVI, for later proposing three new methods/variants of FedLap that incrementally improve the accuracy performance. For those familiar with the probabilistic perspective of federated learning, PVI, and Laplace approximation, the work and particularly the contributions are easy to follow and nice to be proven on the empirical results.\n\n**[S2]** - The drawing of correspondence between ADMM and PVI is elegant, and for sure is a great novel contribution if it has not been done before (I think it hasn't). Despite the introduction of the diagonal covariance in FedLap-Cov made sense to me, and the adaptation of updates seemed also correct. I liked the final technique to ensure the positive definiteness of local and global precision matrices in the paragraphs after (14), quite interesting and useful.\n\n**[S3]** - Despite I did not particularly see the advantages of FedLap-Func in Section 3.4, I see its utility and the spirit of even improving a bit more the performance of FedLap. Here, my experience with such methods is that the ratio improvement/computational-cost-increment is not as favorable as in the FedLap-Cov. I imagine that could have limited the experiments to 10-100 clients only.. I could be wrong on this aspect, though.", "weaknesses": "Despite the fact that I think the work has significant strengths in its current version for being accepted, I detect some unclear parts, or at least \"corners\" where there is not much light to understand what is going on. Some of these are:\n\n**[W1]** - Computational cost is for sure one. While reading the paper, the missing aspect of the computational cost always rings the bell of an average reader. I say this mainly because the connection between ADMM and PVI is fantastic, but once the improvement is made with FedLap and its variants, it is not explicitly indicated what are we paying to get that extra improvement in the performance. In this direction, I want to say that I am just missing the analysis and the clarity in that regard of the problem.\n\n**[W2]** - I'm fine with Eq. (15) and setting up the update as an optimization problem, but then I do not really know if sharing some soft labels or inputs between local clients and the central server is worth it.. It kind of modifies a bit the spirit of federated learning and introduces extra issues like privacy, etc, that are ignored imho.\n\n**[W3]** - It makes a lot of sense to put the work in the context of Laplace approximation methods, BNNs, etc where basically the parameter times parameter dimension of matrices is a huge problem for scalability. Some light-light indication of this is said around L224-225, but overall, I feel that all of this between Eq. 10 and Eq. 14 is ignored, or at least not being highlighted for the reader (the dimensionality problem, or an indication of what are we considering here). I am particularly concerned about V_k and the client updates.\n\n**[W4]** - Experiments are fine, but Table 1 is somehow limited. I basically missed some larger datasets, where there is not that much similarity or structure in the data as in MNIST or FMNIST (i.e. Imagenet or similar), and augmenting a bit the number of clients to see what is really the flexibility/resistance of new variants to such challenges.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2fKhNoHX28", "reviewer": "Reviewer_8nbj", "summary": "Federated learning is the problem of training a model where training data is distributed over clients and due to privacy must stay on clients. Given $K$ clients and a parameterzed model with paramters $w$, each client can optimize its parameters $w_k$ for $k=1,..,K$ and we need to merge the information to learn a global model $w_g$.\n\nMany methods have been proposed and within which Alternative Direction Methods of Multiplier is a method that ensures the clients and the \nglobal model converge to the same stationary point, the clients send a $v_k$ vector to the server that (loosly) shows the weight change from the global model $w_k - w_g$, the global model aggregates and updates and replies with the new global model $w_g$.\n\nSimilarly Bayesian methods learn an approximate posterior distribution over weights $q(w)$ where clients have their learn their own model $q_k()$ send the server an approximate likelihood function $\\hat{t}(w)$. The server aggregates likelihoods, updates the global approximate posterior $q_g()$ and sends it back to the clients, Partitioned Variational Inference (PVI) is a recent work that constructs a common framework covering many prior work methods.\n\nThe paper describes the above frameworks. Then the authors show that for a particular choice of $q_k()$, $q_g()$ in the PVI framework, notably where they are Gaussian with a priori fixed variance so only the mean is learnt, corresponds almost exactly line by line with ADMM, i.e. ADMM is a MAP estimate of some specific Gaussian distributions. The authors propose FedLap-Cov that simply generalizes this observation by allowing learnable covariance matrices in the Gaussian distributions. Finally FedLap-func is proposed that effectively allows a subtle data sharing between client and server hence the server can perform weight optimization as well.\n\nExperiments are performed across a range of standard benchmarks showing FedLap variants outperform both poor (FedAvg) and strong (FedDyn) baselines.", "questions": "- for the benchmarks, the weights used in to make predictions were just the MAP estimate? Was there any marginalising predictions over the distributions of weights?\n- Equation 13 uses gradients from a point in space that is not the centre of the Taylor approximation, is this principled or just a necessary hotfix?\n- Equation 16: can the authors provide some context as to why the new weight space term is subtracted?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "## Main Comments\n- __intuitive demonstration of equivalence__: models trained by minimizing a loss function and a regularizer are often easily shown to correspond to a log likelihood and a log prior hence correspond to a MAP estimate. This work shows the same connection in a less obvious setting.\n- __intuitive extension__: as corollary from above, generalizing the ADMM method for more general distributions makes perfect sense and allows a natural update equations, notably the precision matrix \"weighted sum\" of global and local models.\n- __well written__: the paper takes the reader on a well paced journey through prior works ADMM, PVI using very readable mathematical notation (more readable than the original works in my view). FedLap is then presented to specifically bridge the gap between the prior discussion and the new discussion of FedLap-Cov and FedLap-Func.\n- __extensive benchmarks__: UCI datasets, MNIST and Cifar10 experiments are provided and a deeper dive in to FedLap vs FedDyn is given in the appendix. Hyper-parameters were tuned for all methods.\n\n## Minor comments\n- I personally like that the authors avoid saying \"our method\" and opt for the more humble approach of using the method name \"FedLap\" throughout the paper.", "weaknesses": "I felt the paper was very good and only have minor comments.\n\n### Minor Comments\n- __limitations__ I may have missed this but limitations do not seem to be clearly discussed\n- __computational complexity__ on P6 only discusses the diagonal covariance case which has linear added complexity. While if I understand correctly, the FedLap-Cov has a quadratic additional parameters, upper triangle of a cov matrix. A few concerns come to mind (1) one goal of federated learning is handling unstable client connections by minimising communication rounds, but presumably also minimising traffic, FedLap-Cov seems to add significant traffic (2) the scalability of the FedLap-Cov to many parameters is limited, e.g. resnet-50, although a quick look at past works suggest small applications like UCI datasets, Cifar10 and MNIST are common benchmarks (and not imagenet).\n- __final model is point estimate__ I may have misunderstood, but I could not find explicit mention about what weights at inference time, I assume a single set of weights $w_g$ was used, there was no marginalising. I.e. FedLap methods learn a Bayesian distribution over weights, but the benchmarks have just used a point estimate.\n- __Equation 13__: I apologize if I am mistaken but I was rather confused, the first $m_k$ may be a typo? As I understand, in a delta approximation, normally, we would make a 2nd order Taylor approximation of loss $l_k(w)$ around $w=m$. But to avoid headaches with a hessian term with both $m$ and $S$ interacting as we optimize, we freeze the hessian to a value from a prior iteration $w=m_{old}$ artificially removing the dependence on $m$, this is arguably not an accurate Taylor approximation anymore, is this a common practice?\n\n\n### Possible Typos\n- Equation after Eqn 9: $\\mathbb{E}[\\hat{t}_k(w)]$ missing $\\log()$\n- Equation 12 square brackets on the second term\n- Equation 13: first $l_k(m_k)$ -> $l_k(m)$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ksVb4QkLU5", "reviewer": "Reviewer_GZSC", "summary": "The authors provide new insights connecting ADMM (Alternating Direction Method of Multipliers) and variational Bayesian (VB)-based approaches for federated learning. The main finding reveals that the dual variables utilized in ADMM naturally correspond to the 'site' parameters in VB.", "questions": "Missing important literatures in Bayes federated learning:\n\n[1]. Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms. ICLR. 2021.\n\n[2]. On Convergence of Federated Averaging Langevin Dynamics. UAI 2024.\n\n[3]. Bayesian Federated Learning with Hamiltonian Monte Carlo: Algorithm and Theory. JCGS. 2024", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. connections between ADMM and VB-based approaches may be interesting for federated learning.", "weaknesses": "1. Whether the federated locals can be used to recover the global solution is not justified both theoretically and numerically.\n\n2. Why do you want to connect ADMM to Bayesian? Bayesian learning is known to be the go-to method for conducting uncertainty quantification and non-convex optimization. The algorithm is only evaluated in the optimization perspective. The Bayes interpretation is only used as a preconditioner. If this is the only purpose, the connections to Bayes are not needed at all.\n\n3. Bayes seems to be an important component in this paper. However, the connections to Bayes are not fully evaluated empirically. I recommend the authors to see how the uncertainty experiments are evaluated in [1,2,3].\n\n[1]. Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms. ICLR. 2021.\n\n[2]. On Convergence of Federated Averaging Langevin Dynamics. UAI 2024.\n\n[3]. Bayesian Federated Learning with Hamiltonian Monte Carlo: Algorithm and Theory. JCGS. 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3P0J3Kqomq", "reviewer": "Reviewer_E4hs", "summary": "This paper provides a first connection between two federated learning approaches, one based on ADMM and another based on Variational Bayes (VB). Based on this connection, the paper proposes two variants that incorporate the full covariance information and functional regularization. The proposed methods improve the ADMM baseline in various benchmarks.", "questions": "see the weakness section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The connection between ADMM and VB is novel and interesting, which opens up new avenues for improving the algorithm from a variational / Bayesian perspective. \n\n- The overall empirical results are convincing, showing improvements of the proposed methods against the baseline.", "weaknesses": "The paper can be improved  by adding more discussions on complexity and limitations, as well as empirical results quantification. In particular, the key weaknesses are the following. \n\n\n- missing discussion on limitations. What is the complexity of the extended VB approaches against the baselines, in particular when full covariance estimates are considered. Additionally, FedLap-Func is used in the setting where some inputs are shared globally. Can you discuss the practical implication of that?\n\n- the empirical comparison needs more statistical rigorousness. Table 1 displays the average accuracy with standard deviation in parentheses, and top two performing algorithms are bolded. Can you elaborate how do you define the \"top two\"? A more standard practice in statistics is to highlight the top results within 2 *standard errors*, in that case I would imagine, for example, FedDyn and FedLap to be statistically equally performant on CIFAR10 with comm Round=50 (last row of the table 1).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r06mjJenWc", "reviewer": "Reviewer_kAKo", "summary": "This paper establishes the connection between federate learning and variational Bayes through the \"site\" parameters in partition variational inference (PVI). For particular choice of variational distribution in VB, PVI could (almost) be identified as an Alternative Direction Methods of Multiplier (ADMM), hence named FedLap. Base on such connection, the authors propose two improvements of FedLap (ADMM) by considering full covariances (FedLap-Cov) and function space (FedLap-Func) information coming from VB literature. Then multiple benchmark tests are carried out to demonstrate the superior performance of the proposed improvements compared with vanilla FedLap and competitive federate learning algorithms including FedAvg, FedProx, and FedDyn.", "questions": "* Could you please include a paper outline at the end of Section 1?\n\n* Could you bring the definition of $\\hat t^*_k=\\frac{q_k}{q_g}$ from (5) to (4) where it appears for the first time? I guessed $\\hat t^*_k=\\frac{q_k}{p_0}$ while at equation (4).\n\n* `FedLap` seemingly indicates the Laplace variant of the PVI. But it is still Gaussian used as the variational distribution. Why is it called \"Lap\"? Would you please elaborate?\n\n* FedLap-Func only provides marginal improvements over FedLap-Cov and sometimes even makes it worse. Any explanation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "A nice connection is drawn between federate learning and variational Bayes. Two algorithms improving upon ADMM are proposed based the newly discovered connection. There is good potential in proposed methods, both numerically and methodologically.", "weaknesses": "FedLap-Func only provides marginal improvements over FedLap-Cov and sometimes even makes it worse. Did you explore other problems outside the benchmarks to seek possibly more appropriate applications?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)