Federated ADMM from Bayesian Duality
Authors: Thomas Möllenhoff
Siddharth Swaroop
Finale Doshi-Velez
Mohammad Emtiyaz Khan
Venue: ICLR 2026 Poster (2026)
Keywords: bayesian deep learning
variational inference
variational learning
federated learning
convex optimization
splitting methods
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a new Bayesian approach to derive, extend and improve federated ADMM.

Abstract: We propose a new Bayesian approach to generalize the federated Alternating Direction Method of Multipliers (ADMM). We show that the solutions of variational-Bayesian (VB) objectives are associated with a duality structure that not only resembles the structure of ADMM's fixed-points but also generalizes it. For example, ADMM-like updates are recovered when the VB objective is optimized over the isotropic-Gaussian family, and new non-trivial extensions are obtained for other exponential-family distributions. These extensions include a Newton-like variant that converges in one step on quadratic objectives and an Adam-like variant that yields up to 7% accuracy boosts for deep heterogeneous cases. Our work opens a new Bayesian way to generalize ADMM and other primal-dual methods.

Published as a conference paper at ICLR 2026
FEDERATED ADMM FROM BAYESIAN DUALITY
Thomas Möllenhoff"
Siddharth Swaroop"
RIKEN Center for AI Project
University College London London, United Kingdom
Tokyo, Japan
thomas. .noellenhoff@riken ip
s s www.oop@uc..... ac. uk
Finale Doshi-
Mohammad Emtiyaz Khan RIKEN Center for AI Project Tokyo, Japan emt iyaz. khan@riken, jp
Harvard University
Cambridge, United States
finale@seas harvard.
ABSTRACT
We propose a new Bayesian approach to generalize the federated Alternating Di- rection Method of Multipliers (ADMM). We show that the solutions of variational- Bayesian (VB) objectives are associated with a duality structure that not only resembles the structure of ADMM's fixed-points but also generalizes it. For exam- ple, ADMM-like updates are recovered when the VB objective is optimized over the isotropic-Gaussiaa distributions. These extensions include Newton-like variant that converges in one step on quadratic objectives and an Adam-like variant that yields up to 7% accuracy boosts for deep heterogeneous cases. Our work opens new Bayesian way to generalize ADMM and other primal-dual methods.
Gaussian family, and new non-trive ooiviaivial extensions are obtained for other
exponential-famie
INTRODUCTION
The Alternating Direction Method of Multipliers (ADMM) forms the backbone of many federated learning algorithms (Acar et al., 2021; Gong et al., 2022; Mishchenko et al., 2022: Wang et al., 2022; Zhang et al., 2021: Zhou & Li, 2023). The goal of federated learning to train a global model at server without ever accessing the local data stored at the clients. In ADMM, this is achieved through communication between the server and clients, as shown in Fig. (left). The server first broadcasts the global parameter to the clients, and the clients use it to update their local parameters. The clients then send the updated parameters and their local gradients back to the global server, which combines them to update the global parameter. By repeating these steps, ADMM can recover the solution obtained by jointly training on all data, provably so in many cases.
ADMM was proposed back in the 1970s (Gabay & Mercier, 1976; Glowinski & Marroco, 1975) and its roots can be traced back to an early work by Douglas Rachford (1956). Yet, it continues to be used more or less in the same form it was originally proposed. The robustness of ADMM's algorithmic structure is intriguing, and we wonder whether there a more general formulation of this structure. Such formulations could be especially relevant tackle new issues that arise in federated deep learning. Our search for new generalizations follows a recent result by Swaroop et al. (2025) who connect a driiational-Byessii a close resemblance between ADMM and VB, but it falls short of deriving ADMM as a special case of VB. Our goal here to fill this gap by proposing new Bayesian way to generalize ADMM.
(VB) version of federated learning to ADMM. Their work shows
We propose a general VB framework that can be used to derive and extend federated ADMM. Our main result is to show that the solutions of the VB objective are associated with duality structure that not only resembles ADMM's fixed-point equations, but also naturally generalizes them. We call this structure Bayesian-duality and use it to propose a new algorithm called Bayesian-ADMM. Our generalization is obtained by making two changes to ADMM. First, we introduce distributions over parameters and, second, we replace gradients by natural gradients (Fig. 1). The use of natural gradients is crucial to fix the issue of Swaroop et al. (2025) and enables us to derive classical ADMM
"Equal contribution, CCoresponding author. Code available here.

---

Published as a conference paper at ICLR 2026
ADMM
BayesADMM
g
4g
Server
Client k
Server
Client k
(0 VC2)
(qk- VE(GA)
Figure 1: In ADMM, the server broadcasts the global parameter 0g to the clients who then update their local 02 and send them back to the server along with the gradients Vl2 of their loss functions. ManyssianAAMMM generalizes ADMM by using distributions over global and local parameters (denoted by 9g and qJ. respectivelly and replacing gradients by natural gradients (denoted by VE((2).
ResNet-20 on CIFAR-100 with 10 clients
MLP on MNIST with 100 clients
50-
IVON-ADMM
IVON-ADMM
Test Accuracy
45-
Test Accuracy
90-
FedAvg
40-
35-
FedAvg
FedLap FedProx
85
FydDyn
30
20
40
60
100
Communication Round
Time (min)
We derive an Adam-like extension of ADMM, called IVON-ADMM, which gives up to 7% accuracy boost over existing federated deep-learning methods (left) without any increase the cost and overall runtime (right). Details of the experiment are in Sec. 4.
Figure 2:
as a special case of Bayesian-ADMi that employs isotropic Gaussian posteriors. New non-trivial generalizations also automatically obtained by using other exponential-famil distributions.
We derive two new extensions of ADMM using Bayesian- ADMM. The first extension a Newton- like variant obtained by using multivariate Gaussian distributions. Unlike classical ADMM, this new variant converges in one communication round when applied to quadratic objectives. The second extension is an Adam-like variant obtained by restricting the covariance to be diagonal matrix. This variant can be efficiently implemented by using the IVON method of Shen et al. (2024) and yields up to 7% accuracy boost for deep heterogeneous cases without increasing the cost and overall runtime (Fig. 2). Ultimately, our work opens new Bayesian way to generalize ADMM and other primal-dual methods. This was not possible before, even though lot of work has been done recently on Bayesian methods for federated deep learning (Al-Shedivat et al., 2021; Guo et al., 2023; Kotelevskii et al., 2022; Louizos et al., 2021; Pal et al., 2024; Yurochkin et al., 2019).
FEDERATED LEARNING VIA ADMM AND VARIATIONAL BAYES
Federated learning aims to train a global model with parameter 0, at a central server by communicat- ing with K clients and without ever gaining access to their local data. The clients can access their data through local loss functions denoted by (0) for the k'th client. The goal is to solve for
K
0 arg min Σ (O),
where (0) denotes the regularizer. However, since the server not allowed direct access to lg, is forced to solve this problem by communicating with clients. The main idea is perform distributed
2

---

Published as a conference paper at ICLR 2026
ADMM's Duality
Bayesian Duality
Server
Clients
Server
Clients
0%
0+
A K
K
V VCk k
VE[CK]
v*
1.a v
2
1 K
Figure 3: The left side shows the dual structure associated with ADMM's fixed-point equation (Eq. 4). The right side shows our new -uualian-DDaaliyy structure associated with VB's optimality condition (Eq. 13). In ADMM, 0r and 0% are the primal variables, while v and vî are dual variables. Analogously, in Bayesian Duality, 2 and μ are primal, while X and Mi are dual.
optimization by exploiting the structure of the solution. For instance, for Eo(0, 1 01 the optimality condition of the problem can be written as sum over local gradients (proof in App. A),
K
(lx(O)))
This structure suggests distributing the computation of local gradient V4 across the clients, and gathering those results to estimate 8- ADMM provides framework to perform such distributed opti- mization. It introduces local parameters 02 and aims to solve an equivalent constrained optimization problem:
K Σ ( 0, O
min
+ (o,,), such that 0, = 0g for all = 1,,2.... K.
The problem can be solved by formulating a Lagrangian with multipliers, denoted by Vą for the ''th constraint. The stationarity condition of the Lagrangian, shown below, then provides a way to distribute the computations (proof in App. A),
K
0 0%
V = - (()
0 Vzg
The first condition says that we set 0 = 0 to satisfy the constraint. The second condition sets the optimal multiplier VI to the negative of the local gradient. The third condition gathers all the multipliers into global variable vg. and finally the fourth condition uses it to get the global 0g back.
The optimality condition can be drawn in the form of dual structure, similarly to that used by Rockafellar (1967, Fig. 2). The parameters 0 and 0% are the primal variables defined in the space of valid parameters, while v and v are dual variables defined in the space of valid gradients. The duality structure, shown in the left panel of Fig. 3, summarizes the flow of information. First, the server broadcasts its global 0 to the clients and the clients set 0% 0-. Then, the gradient V&& is evaluated and assigned to the dual vi. All the dual variables are then gathered using a sum and assigned Vy. which then used to obtain 0-
The Lagrangian also yields an algorithm to perform optimization that preserves this flow of infor- mation. Essentially, we first perform local optimizations at all clients obtain Qs, then we update the dual ve, and finally these results are communicated to the server which updates Og. A detailed derivation is in App. A and the resulting updates are summarized below:
Client updates:
OK arg min Eu (0x) +v 0x
2 -110= 0
(5)
V& - Vk + p( -0,),
(6)
K
Server update:
0g f- a.-argmy min (0)),) 6,
T v
0
+
0
02
(7)
1=1
1=1
3

---

Published as a conference paper at ICLR 2026
Here, p > is the (inverse) step-size, added to handle the quadratic proximal terms for each update. An important property of these updates, as shown in App. A, is that after any client's 's updates, the dual vectors are simply equal to their local gradients, that is,
VA (0,))
(8)
As result, v, slowly approaches the optimal i. This can be more clearly seen for the special case when ()g) = [1011 where the update for og simplifies nicely (see Eq. 27 in App. A).
The form of the ADMM updates has remained more less the same since its original proposal in the 1970s. For instance, variants to accelerate ADMM's convergence only introduce additional variables, but do not change the form of the algorithm; this includes methods that overrelaxation (Eckstein & Bertsekas, 1992) or momentum (Chambolle & Pock, 2016). The same is true of variants that replace quadratic proximal terms by scaled (Mahalanobis) norms (Goldstein al., 2013; Pock & Chambolle, 2011; I: Ye et al., 2020) or Bregman divergences (Ma et al., 2025; Wang et al., 2014; Wang & Banerjee, 2014). The robustness the ADMM's algorithmic structure is intriguing and wonder whether there is more general formulation to go beyond it. Such generalizations could be useful to handle new issues arising in federated deep learning due to ieter-ehteroggeeeity and missing data.
Our approach here is inspired by the work of Swaroop et al. (2025) who show connection between federated ADMM and distributed variational Bayesian (VB) method called Partitioned Variational Inference (PVT) (Ashman et al., 2022). They find line-by-line correspondence between the ADMM and PVI updates, but fall short of establishing an exact connection. In what follows, we will fix this issue by using duality structure associated with the solutions of the VB objective.
BAYESIAN DUALITY TO GENERALIZE ADMM
We will now present our new Bayesian- duality structure to generalize ADMM. We start by introducing the dual structure associated with the fixed point equation of VB, and then use it to introduce Bayesian duality and Bayesian-ADMM,
3.1 VARIATIONAL BAYES AND ITS SOLUTIONS
We will use a VB reformulation of Eq. which 'lifts' the original problem to define a new problem where we optimize with respect to probability distributions g in set Q. The goal then is to solve for
K
9 argmin 9
+KL(q || mp).
Q
The second term is the KL divergence between the distribution q and prior To x exp(-4o). When & corresponds to proper likelihoods and Q contains all possible distributions, then g is the standard posterior distribution. Restricting Q to a smaller set, say a set of Gaussian distributions, yields an approximation to the posterior. An important property of the VB formulation is that the solution 0 of Eq. can be recovered from 9 when Q is set to isotropic Gaussians (Khan Rue, 2023, App. C.I). We will use this technique later to derive ADMM from federated algorithm for the VB objective.
The first element needed to generalize ADMM to set q E Q to be of exponential-family (EF) form. An EF has log-linear form with respect to a sufficient statistic, denoted by T(Ø). shown below,
(0) (0))) exp ((A, T(O) A())).
(10)
Here, À denotes the natural parameter, h(0) denotes the base measure, and A(X) is the (convex) log-partition function. For readers unfamiliar with EFs, a brief introduction is included in App. B.I.
The second element is to define natural gradients that scale the gradients by using the inverse Fisher- information matrix of q. For EF distributions, the natural gradient can be conveniently obtained by using the expectation parameter, defined fe VA(A) Eg T()).. We can show that gradients with respect to μi are equal to natural gradients with respect to À (Khan & Rue, 2023, Eq. 4):
(pL( = F(2) (VX--VVX((( ( V)C(m(2))
(11)
4

---

Published as a conference paper at ICLR 2026
where (A) and 4 denote the Fisher information matrix and natural gradient with respect X, respectively. This equation holds because the (λ, μi) pair constitutes a dual-map (Amari, 2016) through the convex function A(X) and its Fenchel conjugate A (H),.. Such maps common in information geometry and convex optimization (Malagò et al., 201 I; Raskutti & Mukherjee, 2015). In the rest of the paper, will drop the subseript i and write the natural gradient simply (L(U).
Khan & Rue (2023, Eq. 5) used this property to express the solution Zg in terms of natural gradients. Specifically, they show that the natural parameter X of og is given by
K II (0),
1 Σ (P)
( Z 1mm0
(12)
where we define the optimal "site" functions as E (0) exp (-(VCx T(O))). and Z is the partition function. A proof is included in App. B.2. The two equations were originally proposed in Khan & Lin (2017, Eq. 11) and Khan & Nielsen (2018, Eq. 18), respectively.
These equations have a strikingly similar form Eq. 2. There, f is equal to the sum of the local gradients, while here, X is equal to the sum of local natural gradients and o is equal to the product of the local site functions. The existence of these equations is also mentioned Swaroop al. (2025, Eq. 4), but they do not make use of natural gradients. We will now show that, by using these equations, we can directly generalize ADMM's dual structure.
3.2 BAYESIAN DUALITY
We will now derive dual structure for the VB solutions by drawing analogy to Eq. 4. There, introduce local 0% and their corresponding dual variables Vž which are set to the (negative) local gradients (0.) The Vī are then added together obtain v and 0r For the VB fixed-point, we follow the same process to define local variables and their corresponding dual variables.
We introduce local distributions of with pair (Xx+ (x)) and denote their corresponding dual variables by ni. Using these variables, we can express the optimality condition in Eq. 12 as the following equivalent set of conditions:
K
nh (pE)
m
15 = AA(A))
(13)
Here, 14, and 2e are the primal variables in the expectation-parain parameter space, while ni and X are the dual variables in the space of valid natural gradients. The first three conditions extension of those in Eq. 4. The last condition connects the pair (X (x.)) through the dual map of the EF This duality of EFs provides foundation build the duality of VB, which refer the Bayesian-duality structure. The dual structure is visualized in the right panel in Fig. 3.
direct
This duality is missed in the work of Swaroop et al. (2025) who connect VB to ADMM without using natural gradients. They show that the parameters of sites l correspond to the dual variables v& in ADMM, but fall short of providing an exact correspondence. This can be fixed by using Bayesian duality to write dual structure in the space of distributions and sites, as shown below:
K
ch = G
(0) exp (-( [e.)) T(0))
II
T
9
(14)
2
k=
This shows that the sites are locally constructed using natural gradients. These are then multiplied together and normalized to get the global 9 which is also equal the local gr The normalization in the last step plays the same role as EF's dual map in Eq. 13.
3.3 BAYESIAN-ADMM
We will now present an optimization algorithm that follows the flow of information suggested by the Bayesian Duality structure. The algorithm closely follows the ADMM updates given in Eqs. to 7, but makes two important modifications. First. the quadratic proximal terms are replaced by the KL divergence which more natural choice for the EFs. Second, the update of the dual ny is modified
5

---

Published as a conference paper at ICLR 2026
to ensure similar condition Eq. 8 holds, that is, make sure that the duals are set to the latest local natural gradients after every local update.
The final algorithm, which refer to Bayesian-ADMM, is shown below: Client updates:
Ms arg min ((( + (M m&) PKL( (gg)
(15)
Me
Th - 7 + p()= 2g))
(16)
Server update:
n c arg min KT (T) - E ( + Σ PKT 11))). By
(17)
k=1
The first and third line both be seen as modified VB problems that server respectively. The modification in the second line deserves a bit more explanation. If we follow the update in Eq. 6, then the dual update should be 7/, 7) p(11 - (1))) However, this update does not ensure that 7h are equal to (px( Instead, if we use the difference between the natural parameters A Ag. this issue resolved; see a proof in App. C.I. Intuitively, subtraction among natural parameters corresponds to division in EFs, while the same for expectation parameters may not always make sense. We further discuss this choice in App. C.2, along with its relationship other methods, such as, Bregman ADMM (Wang& & Banerjee, 2014) and the Bayesian learning rule (Khan Rue, 2023).
at the clients and
The algorithm can also be written in distributional form by defining sites 4 = exp((n)) T(0))).
Client updates:
qle f- arg min Egg [6x + loo tx] pKL(gg ll (g)
(18)
9hi
qk
to x to
(19)
4g
K Σ Eg log tel Σ PKL lk).
Server update:
gg f arg min (g To)
(20)
4a
This algorithm closely resembles the PVI algorithm used by Swaroop et al. (2025). Our derivation here highlights an additional fact that the site parameters 77; are equal to the local natural gradients after every local update. The main differences PVI are highlighted in red, which include the use of step-size P and addition of the KL term in the global update. In fact, PVI is more similar to an algorithm by Tseng (1991) called the Alternating Minimization Algorithm (AMA). The importance of step-sizes are well-established for AMA (see Tseng (1991, Eq. 3.4d)), and this connection can be used to improve the convergence PVI as well. We show such results later in Fig. 5b. More details regarding the differences between Bayesian-ADMM and PVI are given in App. C.3.
We note that our derivation here is to follow the ADMM update closely so as to preserve the information flow suggested by the dual structure. It is also possible to derive these updates by using a Lagrangian, similarly to the ADMM case. This derivation does not yield the update of n, we used. Neither does it connect to the update used in PVI, which is inspired by similar updates used in expectationn-poopagation (Minka, 2001). Our modification here is important to connect such updates used in variational methods to those used in ADMM. These details are further discussed in App. C.4. We also remark here that Lagrangian formulations of VB have been used for dual optimization of Gaussian latent models (Khan, 2014; Khan al., 2013), but they do not exploit the duality of exponential families, instead they use aeaarcovvre parametrizations of Gaussians.
3.4 DERIVING FEDERATED ADMM FROM BAYESIAN DUALITY
We now derive the original ADMM as a special case of Bayesian-ADMM and then propose new extensions of it. All of the derivations that follow rely on simple two-step procedure:
1. Choose an EF and identify its X, μs and T(Ø), for example, using Nielsen & Garcia (2009).
2. Derive the form of the natural gradient, for example, using Khan & Rue (2023).
To derive ADMM as a special case, we choose the EF to be isotropic Gaussian ( = (( m with mean m. For the two steps, we can write the necessary quantities as follows:
m. p m.
T(0)= 0, V
E,g [v(x]
(21)

---

Published as a conference paper at ICLR 2026
Algorithm 1 (IVON-ADMM) Adam-like variant of ADMM implemented using IVON (Shen et al., 2024). Additional steps or modifications when compared to regular federated ADMM are highlighted in red. A derivation is in App. D.3 and details on IVON in App. D.4.
Hy Hyperparameters: 8 > 0. p 0, 0.
Initialize: V* + 0, U2 < mg f- 0, Sy + a. pK).
I: while not converged do
2: Broadcast my and Sy to all clients.
3: for each client 1,..., K in parallel do
4: Form the loss {()) + (fx (0) o vk 1 1 diag uu )x)
5: Get (m,, Sk.) by training on (()) using IVON with prior N( m. diag(s,)
6:
Vk V& + 7 ms s,
7:
Ug < us + (Sk= 8g)
8:
end for
9:
Gather m V2 and Sg. un from all clients.
0:) Mean(Si.x) + a [6 + Sum(u1.)) mg < [(1 a) (aaanss:kmm::d + a Sum (v1vk)) //g
8g
11:
12:
end while
The last equality due Bonnet's theorem (Rezende et al., 2014). Plugging these equations in
Eqs. 15 to 17, we can write a simplified version that closely resembles ADMM. A full derivation is given in App. D.1, and here we illustrate simplification of the client updates. We start with Eq. 15 where, after renaming 7h to the update simplifies to
m f argmin Ex(0) m. m) (0) + v 0 + m -
-0,1-mmll myl1
(22)
m.
The simplification is obtained by using the definition of the KL divergence for isotropic Gaussians, followed by some rearrangement. The update takes similar form to ADMM where linear term v 0 is added to the loss. The main difference is that there is an expectation over the isotropic Gaussian distribution, which is similar to those used in sharpness- awaree aware minimization (Foret et al., 2021; Möllenhoff & Khan, 2023). Similarly, by using the definitions in Eq. 21, we can simplify Eq. 16 to Vı + VA + p(m; mg). This update has similar property to Eq. 8 where, after every client update, the dual V3 EE.. [ve] is the expected gradient.
To recover ADMM as special case, we employ the delta method (Khan & Rue, 2023, Sec. 1.3.1) where we set [E (0) = (m,). It easy to see that, after this approximation, the update is identical to Eq. 5 if we rename m, and my to 02 and 0g. respectively. The simplification of other lines follows in a similar manner and detailed in App. D. We note that the method of Swaroop et al. (2025) does not recover ADMM mainly because it lacks the last KL term in Eq. 20.
3.5 NEW NEWTON-LIKE AND ADAM-LIKE VARIANTS OF FEDERATED ADMM
We now present two new extensions of federated ADMM. The extensions are obtained by spe- cializing Bayesian-ADMMM to Gaussians with full covariance and diagonal covariance respectively. We first explain the Newton-like extension which uses full-covariance Gaussians for both server 9g (0) = (0 my: S and clients qu (0) A mg S ). The derivation exactly follows the two step procedure described earlier and is detailed in App. D.2. Below, we give brief summary.
When using full-Gaussian, the Bayesian client step in Eq. 15 essentially adds quadratic term (highlighted in red),
(m:, Sk)= + argmin Ex (0) + + 0 V + pKL ll).)
(23)
This happens because the sufficient statistics now T(0) = (0, 000) giving rise to an additional dual variable, denoted by the matrix Vt. As shown in App. D.2, this dual variable is related to the Hessian, that is. V₁ Egg. [V after every local update. Due to this, the local and global updates can be implemented in the form of a Newton step. Similarly to Newton's method, the new method converges in a single-step on quadratic objectives, which we show in App. D.2 and also later verify in the experiments section.
7

---

Published as a conference paper at ICLR 2026
Server
Client 1
FedADMM
Outlier
Outlier
Bayesian- ADMM
Outlier
Outlier
Figure 4: single outlier that slows down ADMM (top row) poses no issues for our new Bayesian version Bayesian-ADMM (bottom row). The server (left column) with ADMM takes 5 iterations while with Bayesian-ADMi.
only needs iterations (decision boundaries are numbered with iteration number). Client (right column) is the source of the issue which takes 5 iterations to ignore the outlier, while with Bayes, is much faster due to the use uncertainty (gray lines). Light points indicate data on client 1, dark points indicate data client 2 (not shown).
We also derive a scalable Adam-like variant of ADMM by using Gaussians with diagonal covariances. The resulting method is shown in Alg. and makes of the Improved Variational Online Newton (IVON) optimizer of Shen et al. (2024) to solve the client subproblem. For this reason, we call IVON-ADMN A detailed derivation of this method is given in App. D.3. Implementation of IVON-ADMM is nearly identical to the standard ADMM and costs are nearly identical too. Lines 4 and in Alg. solve the client subproblem in PayessianADMer but with diagonal covariances. In practice, this is a standard application of the IVON optimizer, and more details on how exactly used are App. D.4. Lines and implement the dual update in Eq. 16 in Bayesian-ADMM and involve just simple pointwise arithmetic operations. We have introduced a different step-size y in the dual, as we found this to improve the practical performance. Finally, lines 9-11 implement the server step in Eq. 17 in Bayesian-ADMMM
which takes a similar form to Eq. 23
All changes are simple pointwise operations and easy to implement, and the IVON optimizer has no overhead over training with Adam (Shen et al., 2024). Therefore the complexity is the same as federated ADMM algorithms such as FedDyn (Acar et al., 2021). Communication cost is doubled compared to ADMM as we send both mean and diagonal variance vector, as in other preconditioned or Bayesian methods (Swaroop et al., 2025). This is expected to help in cases with heterogeneity.
4 NUMERICAL EXPERIMENTS
Illustrative examples. We show how the Newton-like variant of Bayesian-ADMM better deal with heterogeneous and noisy data in Fig. 4. We compare ADMM to the Newton-like variant in linear classification; data is split over two clients, and client contains an outlier. Federated ADMM takes many (five) communication rounds to resolve the issue due to the outlier. The Newton-like Bayesian-A ADMM assigns high uncertainty to the outlier and quickly adapts in single round, showing the benefits of using posterior uncertainties. Details are in App. E.
In Fig. 5a, we show linear regression problem with full Gaussian posteriors (details in App. E.2), where Poyesian-ADMMr with full covariances converges in single round. This is similar to Newton- methods on quadratic functions, but there are no previous generalizations of such results to ADMM. In Fig. 5b compare the Newton-like Bayesian-ADMM ADMM method with full covariances to PVI and BDMM on logistic regression (details are in App. E.3). We see that PVI without damping diverges, whereas PVI with damping converges, but slower than Bayesian-DDMM.
Evaluation on Federated Deep Learning Benchmarks. We follow similar setup to Swaroop al. (2025) and train neural nets image datasets. We consider different levels of heterogeneity on the clients and different numbers of clients. We focus on comparing IVON-ADMM (Alg. 1) to (i) the distributed Bayesian algorithms Fedl Lap and Fedl Fedl ap- (Swaroop al., 2025), (ii) FedDyn
8

---

Published as a conference paper at ICLR 2026
20000
Global Loss Client
Federated-ADMM
Variational Objective
+
PVI (no dumping)
Bregman-ADMe Bayesian-ADDe
PVT (with damping)
40-
16000
Bregman- ADMM
Accuracy
30-
Rest non-Ruyesin Method
210-
Bayesian-A
12000
20
180-
IVON-PVI
8000
10-
IVON-ADMM
10
25
100
Communication Round
Communication Round
Communication Round
(a) Ridge Regression (Quadratic)
(b) Logistic Regression
(c) ResNeet-20 (CIFAR-I00)
Figure 5: (a) Bayesian-ADMM converges in a single round for quadratics, whereas other methods need many steps. (b) PVI can diverge on logistic regression (MNIST). Damping used by Ashman et al. (2022) improves this, but is slower than Bayesian-ADDMM, also for CIFAR- in plot (c).
Table I: IVON-ADMM outperforms all baselines in terms of accuracy and NLL. Averaging over the posterior in IVON-ADMM (as opposed IVON-ADMM @ m) often improves performance.
Test accuracy ( larger better)
Test NLL (L smaller is better)
Scenario
Method
10 rounds 25 rounds
50 rounds
rounds
25 rounds
50 rounds
MLP. 100 clients
FedAAv/FedPrrox
73.9+0.3
78.7+0.3
81,8++0.3
0.73++0.01
0.60+0.01
0.53+00001
heterog.
FedDyn
75.7=10.4
81.4:002
82.2410.3
0.66:10.01
0.53+0.02
0.51:0.01
FMNIST
FedLap
73.9=0.3
79.4+0.2
82.4++0.3
0.68++0.01
0.59=0.00
0.53=0001
Fedl ap-Cov
76.9+0.7
81.3+0.2
83.0++0.1
0.64++0.01
0,54+0.00
0.49+0.01
IVO-ADDMemm,
78.7:0.11
82.2.0011
82.91:02
0.62:10.00
0.50
0.00
IVON-ADMM
78.7=0.1
82.3+0.11
83.0++0.2
0.62++0.00
0.50
0.48=0.00
FedAvg
72.3:10.4
77.7+00.3
80.0-10.2
0.70:10.00
0.61
0.56:1001
FedProx
72.2-0.3
77.410.1
80.3100.1
0.711:0000
0.61
0.56:10.01
ML.,, 10 clients
FodDyn
75.3=0.8
77.5
78.240.5
0.67
0.63
0.63=0.02
homogeneous
FedLap
72.1:0.2
77.1-0.1
80,2400.1
0.71
0.62:10.01
0.57:10.11
FMNIST(>)
FedLap-Cov
75.020.6
79.810.4
818810.1
0.67110.01
0.5910.01
0.56:1001
IVON-DMM@m,
80,4-0.2
83.1
83.410.1
0.56110.00
0.53
0.60=0.01
IVON-ADMM
80.6:1.2
83.5
84.110.2
0.55
0,50
0.53
CNN
FedAvg/FedProx
73.821.5
75.010.9
75.410.88
0.940.0
1210011
1.510.1
(Zenke et al., 2017)
FodDyn
72.7=0.9
77.4=0.6
79.440.4
1.0+0.0
0.8+0.0
0,7+0.0
10 clients
FedLap
1.3
78.0112
19.511.4
0.7110.0
0.6:0.0
0.6:0.0
heterog. CIFAR-10
FodLap-Cov
75.1=1.1
77.610.7
79.21:0.9
0.71:0.0
0.7
0.620.0
IVON-ADMMeN my
71.5=118
78,9=0.8
80,3+0.6
0.8+0.00
0.6=0.0
0.6=0.0
IVON-ADMM
71.5=1.8
79.0
80.310.6
0,810.0
0.620.0
0.6=0.0
CNN
FodAvg/FedProx
64.3:22.0
65.9=1.6
66.3+114
1.740.1
2.240.1
2.6=0.1
(Acar et al., 2021)
FodDyn
63.6=
64.7+0.6
65.4++10
2.0+0.6
3.6+1.
2.6+0.11
10 clients
Fedi ap
60.2+22.4
66.4+1.1
66.5+112
1.140.0
1.3+0.1
1.8+0.1
heterog. CIFAR-10
FedLap-Cov
58.4=2.4
65.4+1.1
67.5+11.2
1.240.0
10000.0
1020.0
ADMM@M,
63.8+114
69.5+0.8
70.2++0.7
1.0+0.0
1.+0..0
1.5+0.0
IVON-ADMM
63.8
69.5+0.8
70.3++0.6
1.0+0.0
1.0+0.0
1.4+0.0
(Acar al., 2021), which the best-performing federated ADMM-stylc algorithm, and (iii) standard baseline algorithms for federated learning, FedAvg (McMahan et al., 2016) and FedProx (Li et al., 2020). Further hyperparameters and details are in App.
Overall, we find that IVON-ADMM outperforms baselines in most scenarios; IVON-ADMM has lower test loss than baselines, showing the benefits of Bayesian method; and IVON-ADMM is significantly computationally cheaper than the second-best method predictions over the posterior at test-time usually improves the performance.
Ensembling
Scenarios and splits. We consider fully connected neural networks on the MNIST and FashionM- NIST datasets, and convolutional networks on CIFAR-IO and 100, following the setup in Swaroop et al. (2025). We use homogeneous and heterogeneous splits on K 10 and K = 100 clients. On CIFAR- 100 we use a ResNet-20 architecture. For all heterogeneous splits, we follow previous work and sample Dirichlet distributions that decide how many points per class go into each client (Swaroop et al., 2025). For the hehhly-heerrogneouss MNIST split on MNIST (K 100 clients), we assign data from only two random classes per client (McMahan al., 2016).

---

Published as a conference paper at ICLR 2026
Table 2: Test accuracy and NLL after 25, 50 and 100 rounds, with mean and standard deviations over 3 runs. We see that IVON-ADMM outperforms all baselines in both scenarios.
Test uecuracy larger is betfer)
Test NLL (L smaller is better)
Scenario
Method
25 rounds
50 rounds
100 rounds
25 rounds
50 monds 100 rounds
MLP, 100 clients
FodAvg/FodProx
89.620.4
91.8:10.5
94.0=0.4
0.59::0.11
0.36::0.01
0.22::0.11
highly heterog.
FodDyn
89.8+0.1
92.6+:0.2
94.9=0.1
0.60+:0001
0.33-40.01
0.18++0.00
MNIST
FedLap
89.710.3
92.0100.4
94.4=0.3
0.63::0.01
0.381:0.11
0.22+:0.01
FedLap-Cov
91.0=:0.4
92,941044
94.9=0.3
0.47::0.01
0.29:10.01
0.18:00.01
VON-ADMM@m;
91.0110.1
94.2:10.1
96.0=0.1
0.311:0.02
0.19:10.00
0.13+10.00
IVON-ADMM
91.0++0.1
94.2410.1
96.0=0.1
0.31+:0.03
0.19:10.01
0.13+:0.00
ResNet-20,
FedAvg/FedPrro
37,9+0.4
40.4+099
39.8+0.8
2.4+0.00
2.6+0.11
3,4+0.1
clients,
FedDyn
38.4+0.55
39.24:.88
39.6=004
3.54:0.3
3,240.22
3.2+0.11
beterog. CIFAR- 100
FedLap
40.1#:1.1
40.4+1.0
39.7=00.7
2.900.1
3.0+:0.0
3,2+0.00
IVON-ADIE
MON-ADMM@ny
40.0
46.21:0.4
46.5=0.6
23400
2.11:000
2.310.1
40.0:11.0
46.24:0.4
46.6=:066
2.3+0.00
2,110.0
2.2:0.11
Our results are summarized in Tables 1, 2 and 4, as well as Fig. 2. For IVON-ADMM@ m. evaluate using the server's mg, while for IVON-ADMM, we ensemble predictions across 32 samples from qg. We average accuracy and test loss (NLL) over three seeds after 10, 25, 50 and 100 rounds.
IVON-ADMM. is overall the best method. We compare IVON-AD both 10 and 100-client splits FMNIST and CIFAR-10 (Table 1). It is much better than all five baselines on FMNIST in terms of test NLL and accuracy, at all rounds. In particular, we (i) improve upon FedDyn (which is Federated ADMM), showing the importance of including posterior covariances in IVON-ADME -DM;; (ii) improve upon FedLap and FedLap-Cov, showing the importance of the Bayesian-ADMM update equations and of using IVON instead of a covariance obtained via Laplace approximations.
to other baselines for
For longer runs, IVON-ADMM significantly outperforms all other methods. We compare IVON-ADMM to baselines for 100-client highly heterogeneous MNIST split, and a 10-client heterogeneous CIFAR-100 split, in Table 2. We see that IVON-ADMM significantly outperforms all baselines, showing the benefits with more clients and over more communication rounds. After 100 rounds on CIFAR-100, accuracy is improved by 6.7% and test NLL by 0.9.
IYON-ADMM has less overhead than FedLap-Cov. IVON-ADIA has similar cost and runtime to FedAvg, FedProx, FedDyn and Fedlap, despite estimating a (diagonal) covariance matrix. FedLap- Cov requires a Laplace approximation to the covariance matrix, which is slow and memory-intensive: FedLap-Cov is times slower than FedLap on MNIST, times slower on CIFAR- 10, and prohibitively slow on CIFAR-100. even though use LaplaceRedux (Daxberger et al., 2021).
Faster convergence than IVON-PVL. We introduce new method, IVON-PVI, which is a special that the additional step-sizes
case of our IVON-ADMM (with 1, p 1). In Fig. 5c in IVON-ADMM means that converges quicker than IVON-PVI, though the final solution of comparable quality. Additional results in App. G.I.
Sensitivity to Hyperparameters. We perform an ablation study over the step size (p) and the temperature (r) of IVON-ADMM. The results discussed in App. G.2, where the main findings are: (1) Larger p leads to more stable dynamics but lower accuracy, (2) The optimal temperature is at around 0.1 0.2 and deviating too much from can cause degradations in performance.
DISCUSSION
We introduced Bayesian duality, from which we naturally extended ADMM to optimize over distributions. For Gaussians with fixed variance, we recover regular ADMM, and general Gaussians give Newton-like methods and IVON-ADMM. These show good performance when compared to recent baselines. Other approximating distributions may lead new interesting splitting algorithms, and more generally, our work opens up new research paths to extend and improve primal-dum algorithms using Bayesian ideas.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work is supported by JST CREST Grant Number JP-MJCR2112. This work used computational resources on the TSUBAME400 supercomputer provided by Institute of Science Tokyo. This material based upon work supported by the National Science Foundation under Grant No. IIS-2107391. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations (ICLR), 2021.
Vincent Adam, Paul R. Chang, Mohammad Emtiyaz Khan, and Arno Solin. Dual parameterization of sparse variational Gaussian processes. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In International Conference on Learning Representations (ICLR), 2021.
Shun-Ichi Amari. Information Geometry and its Applications. Springer. 2016
Matthew Ashman, Thang D Bui, Cuong Nguyen, Stratis Markou, Adrian Weller, Siddharth Swaroop, and Richard E Turner. Partitioned variational inference: A framework for probabilistic federated learning. arXiv:2202.12273: 2022.
Behnam Bbbagholami-Mohhmaddaaaa Sejong Yoon, and Vladimir Pavlovic. D-MFVI: Distributed mean field variational inference using Bregman ADMM. ht0::/507.0022 2015.
Dimitri P Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014.
Stephen P. Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. Distributed opti- mization and statistical learning via the alternating direction method of multipliers. Found. Trends Mach. Learn., 1-12, 2011.
Antonin Chambolle and Thomas Pock. On the ergodic convergence rates of a first-order primal-dual algorithm. Math. Program., (5((1)))2253 287, 2016.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace ceftx-effforless Bayesian deep learning. In Advances in Neural Informa- tion Processing Systems (NeurIPS), 2021.
Jim Douglas and Henry H Rachford. the numerical solution of heat conduction problems two and three space variables. Transactions of the American Mathematical Society, 82(2):421 1956.
Jonathan Eckstein and Dimitri P Bertsekas. On the Douglas-Rachforl Rachford splitting method and the proximal point algorithm for maximal monotone operators. Math. Program., 55(1):293-318 1992.
Jonathan Eckstein and Michael C. Ferris. Smooth methods of multipliers for complementarity problems. Math. Program., 86(1):65-90 90. 1999.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimization for efficiently improving generalization. In International Conference on Learning Representations (ICLR), 2021.
Daniel Gabay and Bertrand Mercier. A dual algorithm for the solution of nonlinear variational problems via finite element approximation. Computers & Mathematics with Applications, 2(1): 17-40, 1976.
11

---

Published as a conference paper at ICLR 2026
Roland Glowinski and Americo Marroco. Sur I'approximatinn par éléments finis d' d'ordre un, et la résolution, par pénalisation-dualin d' une classe de problèmes de Dirichlet non linéaires. Revue française d' automatique,, informatique, recherche opérationnelle. Analyse numérique, 9(22):41-76, 1975.
Tom Goldstein, Min Li, Xiaoming Yuan, Ernie Esser, and Richard Baraniuk. Adaptive primal-dual hybrid gradient methods for saddle-point problems. arXiv:/305.05:: 2013.
Yonghai Gong, Yichuan Li, and Nikolaos M. Freris. FedADMM: A robust federated deep learn- ing framework with adaptivity to system heterogeneity. In International Conference on Data Engineering (ICDE), 2022.
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing. Federated learning as variational inference: scalable expectation propagation approach. In International Conference on Learning Representations (ICLR), 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Edwin Jaynes. Information theory and statistical mechanics. Phys. Rev., 106(4):620, 1957.
Mohammad Emtiyaz Khan. Decoupled variational Gaussian inference. In Advances in Neural Information Processing Systems (NeurIPS), 2014.
Mohammad Emtiyaz Khan and Wu Lin.
variational inference: Converting variational inference in non-conjugate models to inferences conjugate models. In International Conference Artificial Intelligence and Statistics (AISTATS), 2017.
Mohammad Emtiyaz Khan and Didrik Niclsen. Fast yet simple natural -gradint descent for varia- tional inference in complex models. In International Symposium Information Theory and Its Applications (ISITA), 2018.
Mohammad Emtiyaz Khan and Haavard Rue. The Bayesian learning rule. J. Mach. Learn. Res. (JMLR), 24(281):1-46, 2023.
Mohammad Emtiyaz Khan, Aleksandr Aravkin, Michael Friedlander, and Matthias Seeger. Fast dual variational inference for non-conjugate latent Gaussian models. In International Conference on Machine Learning (ICML), 2013.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.
Nikita Kotelevskii, Maxime Vono, Alain Durmus, and Eric Moulines. FedPop: A Bayesian approach for personalised federated learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, 2020.
Christos Louizos, Matthias Reisser, Joseph Soriaga, and Max Welling, An expectation-maxiiizaiion perspective on federated learning. 41X1v:2111.10119, 2021.
Javier Luque. The nonlinear proximal point algorithm and multiplier methods. Laboratory for Information and Decision Systems, MIT, 1986.
Shiqian Ma. Lin Xiao, and Renbo Zhao. arXiv:2509.08739, 2025.
Bregman Rauhlos-Rachford splitting method.
Luigi Malagò, Matteo Matteucci, and Giovanni Pistone. Towards the geometry of estimation of distribution algorithms based on the exponential family. Foundations of Genetic Algorithms XI, FOGA, 2011.
12

---

Published as a conference paper at ICLR 2026
H. B. McMahan, Eider Moore, Daniel Ramage. Seth Hampson, and Blaise Agüera y Arcas. conmunication-eeficcnn learning of deep networks from decentralized data. International Conference Artificial Intelligence and Statistics (AISTATS), 2016.
Thomas P. Minka. Expectation propagation for approximate Bayesian inference. In Conference Uncertainty in Artificial Intelligence (UAI), 2001.
Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtárik. ProxSkip: Yes! Local gradient steps provably lead to communication acceleration! Finally! In International Conference on Machine Learning (ICML), 2022.
Thomas Möllenhoff and Mohammad Emtiyaz Khan. SAM as an optimal relaxation of Bayes. In International Conference on Learning Representations (ICLR), 2023.
Frank Nielsen and Vincent Garcia. Statistical exponential families: A digest with flash cards.
arXiv:0911.4863, 2009.
Iman Nodozi and Abhishek Halder. Wasserstein consensus ADMM. arXiv:2309.07351, 2023.
Konstantinos A Oikonomidis, Alexander Bodard, Emanuel Laude, and Panagiotis Patrinos. Global convergence analysis of the power proximal point and augmented Lagrangian method. arXiv:2312./2205, 2023.
Manfred Opper and Ole Winther. Expectation consistent approximate inference. Journal of Machine Learning Research, 6(12), 2005.
Shivam Pal, Aishwarya Gupta, Saqib Sarwar, and Piyush Rai. Simple and scalable federated learning with uncertainty via Improved Variational Online Newton. In NeurlPS Workshop Optimization for Machine Learning (OPT), 2024.
Pock and A. Chambolle. Diagonal preconditioning for first order primal-dual algorithms convex optimization. In International Conference Computer Vision (ICCV), 2011.
Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Trans. Inf. Theory, 613)):1551-1557 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning (ICML), 2014. 2014.
Ralph Tyrrell Rockafellar. Duality and stability in extremum problems involving convex functions. Pacific J. Math., 21(1):167-187, 1967.
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 is effective for large deep networks. In International Conference on Machine Learning (ICML), 2024.
Siddharth Swaroop. Mohammad Emtiyaz Khan, and Finale Doshi-Velez. Connecting federated ADMM to Bayes. In International Conference on Learning Representations (ICLR), 2025.
Paul Tseng. Applications of a splitting algorithm decomposition in convex programming and variational inequalities. SIAM J. Control Optim., Opti 19-3)111--13,, 1991.
Martin Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2):1-305, 2008
Fenghui Wang, Zongben Xu, and Hong-Kun Xu. Convergence Bregman alternating direction method with multipliers for nonconvex composite problems, arXiv:1410.862) 2014.
Han Wang, Siddartha Marella, and James Anderson. FedADMM: A federated primal-dual algorithm allowing partial participation. Conference on Decision and Control (CDC), 2022.
Huahua Wang and Arindam Banerjee. Bregman alternating direction method of multipliers. Advances in Neural Information Processing Systems (NeurIPS). 2014.
13

---

Published as a conference paper at ICLR 2026
John Winn, Christopher M Bishop, and Tommi Jaakkola. Variational message passing. J. Mach. Learn. Res. (JMLR), 6(4), 2005.
Zhenzhang Ye, Thomas Möllenhoff, Tao Wu, and Daniel Cremers. Optimization of graph total variation via active-set-based combinatorial reconditioning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference Machine Learning (ICML), 2019.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International Conference on Machine Learning (ICML), 2017.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A federated learning framework with adaptivity to non-iid data. IEEE Trans. Signal Process., 69:6055-6070) 2021.
Shenglong Zhou and Geoffrey Li. Federated learning via inexact ADMM. IEEE Trans. Patlern Anal. Mach. Intell. (PAMI), 45(8):9699-9700 2023.
14

---

Published as a conference paper at ICLR 2026
A ADMM DERIVATION FOR EQS. 2 AND 4 TO 8
To derive the optimality condition given in Eq. 2, we note that lo(0g) = =2110 Therefore, taking the derivative of the objective with respect to gives us
K
K
( + +
0 = Ven (0,)
A Lagrangian can be formed as follows:
min max Oy. OLI VLK
Σ v (0x-0,)] ( +
+ (0x ko()))
(24)
For the Lagrangian, the derivative with respect to 02 yields
(0,) + (vi) 0 (v =0
V ((x
which is the second condition. Similarly, derivatives with respect to 0g gives us
(o(0') [(v))
9 =0
0
Then defining the sum vg as V we get the third and the fourth condition. Finally, taking the derivative with respect Vk, we get the first condition 0f 0%
We now provide short derivation of federated ADMM as shown in Eqs. to 7. The method is standard application of ADMM the consensus problems with regularization, see Boyd et al. (201 1, Section 7.1.1). First, we collect from Eq. 24 all the terms depending on 0₁. and add quadratic proximity term to optimize the objective follows,
Oy f- arg min l (0x) + vI 0x + 0 0gl1?
(25)
0e
with p as the (inverse) step-size. We locally optimize this to update Os. which is shown Eq. 5. The update of va follows in similar fashion,
1 2p
V& + arg max V (0, 02) +
lv
vall
Vk + V old + p(0 0,).
(26)
This is the update given in Eq. 6. The third line the algorithm similarly optimizes,
K +
θ + arg min (0,))
v
e
0
0
0,
2
k=1
For the special case when (0,) a sum over Vk. as shown below by taking the fixed point of the above update:
0 12 the server update only requires average over 0g and
K
K
0
P
O
+
pke,
k=
k=
1 a)
0 +
VB
+ (1
Σ
0
(27)
K
where we define a = 1/(pK + 1).
Finally, to see why the dual variables correspond to gradients in Eq. consider the optimality condition of the client update in Eq. 25:
V (0x) + v + p(0, 0,) -
v p 0g) vl+(Ox).
(28)
Plugging this into the dual update the right in Eq. 26, immediately get v (Ox.
15

---

Published as a conference paper at ICLR 2026
Table 3: A summary of exponential-famie famiiiis used in the paper showing natural parameters sufficient statistics T(Ø), and expectation parameters μ. reproduced from Khan & Rue (2023, Table 2).
Distribution
T(0)
/4
Resulting Method
N(O m.
ADMM
Preconditioned ADMM
N( m. ",)) fixed S
Sm
0
m
N(O m '). S diag(s)
sim
0²
m²
m
+ 1/s
Alg.
N(O m. m-S)
- S
Sm
00
0
mm
m
S
App. D.2
B BACKGROUND ON VARIATIONAL BAYES
B. EXPONENTIAL FAMILIES
In this appendix. provide short self-contained introduction to exponential families. Many commonly used distributions such as Gaussian, categorical or gamma distributions belong to an exponential family. Viewing them such provides us with rich dual structure which exploit in this paper. Table provides an overview over all exponential families used this paper.
Here, give short introduction and refer the interested reader Wainwright al. (2008, Sec. 3) for more details and a rigorous treatment. Similar to Jaynes (1957), we introduce exponential families maximum-entropy solutions of constrained optimization problem. We assume base measure which has density h()) with respect to the Lebesgue measure and then search for density g()) which has maximum entropy (relative to h) and whose expectation under sufficient statistics T T(Ø) matches given expectations or moments moments μ., but only one has the maximum entropy, see Fig. 6a for an example with uniform density h(0)) = 1. The maximum entropy density is given as the solution to the following problem:
There are many distributions which share the same
4μ argmin, KL(q h) s.t. E. ((0)) μ. (()) 2 0, ((0)
(29)
Similar as Jaynes (1957, Sec. 2), introducing Lagrange multipliers for the moment constraints and another Lagrange multiplier a for the normalization constraint, we get the Lagrangian,
1 (1- (1-f ((())a0))
C(g. A, a) KL(g (A, E,(T))) μ4) +
(30)
Taking the first variation of the Lagrangian with respect to q and setting zero, we arrive at,
4f (0) h((0) exp((A (A, T(0)) a).
(31)
which for a A()) log exp((2 (A, T(O)) (d) the definition of the exponential family as provided in the main text in Eq. 10. We also see that the Lagrange multipliers X correspond to the natural parameters the exponential family.
The Lagrangian we used in the main text to derive our Bayesian-ADMIN method closely follows this dual structure of the exponential family. The above derivation provides an additional motivation why one should use the natural parameters À in the dual update of our method in Eq. 16. They live in the correct space Lagrange multipliers to constraints on the expectation parameters.
As mentioned in the main text, Lagrange multipliers X and moments fa form dual coordinate system, where 44 = (A()) = E.g [T()). The map VA: Ω + M maps from the space valid natural parameters Ω {A A(X) 30) the space of valid moments M {E, [(@)] q E Q). The inverse is given by V :M + M where A denotes the convex conjugate A. For an illustration of the two spaces Ω and M the example of Gaussian distributions, see Fig. 6b.
16

---

Published as a conference paper at ICLR 2026
VA
== (0,
M VA
(a) Maximum Entropy Distribution
(b) Dual Coordinates
Figure (a) There are many distributions (two possible ones shown in dashed gray lines) which share the same moments Distributions which have maximum entropy among all possible candidates belong to an exponential family (here: the Gaussian shown in solid black). (b) There a one-to-one correspondence between the moments μ E M and the natural parameters À E Ω which are Lagrange multipliers the moment constraints the maximum entropy variational problem.
B.2 DERIVATION OF THE FIXED-POINT EQUATION EQ. 12
We start by writing the natural gradient of the KL divergence as follows,
KL || mo) = B, log To. = À V (yl.
The last equality follows from Khan & Rue (2023, Eq. 64) and also noting that mo oc exp(-Lo). Writing the second term in terms of Co. we can use the above to write the gradient of Eq. 9 as
K
(L) + +
k=
Setting this n. we get the first equation in Eq. 12. Then, substituting the first equation in Eq. 10 and simplifying, we get the second equation.
C DERIVATIONS AND DETAILS ON BAYESIAN-ADMM
C.I DUAL VARIABLES ARE NATURAL GRADIENTS
The result follows from the optimality condition of the client step in Eq. 15,
(uCx(MA) + m. + p(>x - Ag) =0 O,
(32)
which is directly obtained by taking the derivative in μ and setting to zero. The above the fact that V (L (gg A Ag. see Khan Rue (2023, Eqs. 21, 23). Inserting this expression into the dual update in Eq. 16, get,
T
new f- 7h + p(>: X,) = (pp))
(33)
Thus, the dual variables in Bayesian-ADMM correspond to the natural gradients of the client losses.
C.2 RELATIONSHIP OF BAYESIAN-ADMA TO METHODS FROM LITERATURE
Here, we discuss connections Moyesian-ADMM to various methods from the literature.
Bregman ADMM and augmented Lagrangian methods. Many variants of ADMM have been proposed since its inception. However, we believe that the dual update in Bayesian-ADMM in Eq. 16,
7s 7) p(>s Xg),
a fundamental departure from other generalizations of ADMM used in the optimization litera- ture (Babagholami- Maami--Mhamadaaadi et al., 2015; Nodozi & Halder, 2023; Wang et al., 2014; Wang & Banerjee, 2014), which use the primal variables in the dual update. our setting, this would amount to using μ4 H4 in the dual update, which is unnatural from Bayesian viewpoint. Arithmetic in A- c-oordinatss corresponds to multiplication and division of distributions, but no such simple interpre- tation exists in 4-ccordinaatss. Moreover, have seen in the previous section of the appendix, the
17

---

Published as a conference paper at ICLR 2026
choice of natural parameters in the dual update ensures the natural condition ny - (ppg)) at every step of the algorithm, which is essential to our proposal. Nonlinear augmented Lagrangian and Rachford splitting approaches (Bertsekas, 2014; Eckstein & Ferris, 1999; Luque,
Bregman Douglas-Rachford 1986; Ma et al., 2025; Oikonomidis et al., 2023) come with modified but different dual updates.
Bayesian Learning Rule. Bayesian- ADMM is also closely related to the Bayesian learning rule (BLR) (Khan & Rue, 2023), natural-gradien descent algorithm for VB optimization in the non- federated case. We can show that Bayesian-ADi AMMM mimics the BLR steps when the client updates iterated until convergence before doing server update, stated below.
Proposition C.I. If the client steps and 2 in Eqs. 15 and 16 iterated until convergence for given Ag, then the server step that follows in Eq. 17 will be equivalent to the Bayesian learning rule (BLR) of Khan & Rue (2023).
Proof. To show this, we note that we run the client steps iteratively until convergence, then Ng M and X Xg. and the dual vector 7h is also equal the natural gradient at Mg. Then, write out the server step of Bayesian-ADMM more explicitly in closed form:
Xg (1 a)
A
7k-
(34)
R=
derivation for this form is given around Eq. 51. The average over A is equal Xg. and the update becomes the BLR as written in Khan & Rue (2023, Eq. 6) since 7) are the natural gradients.
C.3 PARTITIONED VARIATIONAL INFERENCE (PVI)
We first rewrite PVI as introduced in Ashman et al. (2022, Alg. 1), bringing it into the form used by Swaroop et al. (2025). Then, we will discuss differences to Bayesian-ADMM. The first step of Ashman et al. (2022, Alg. 1) is a local update of the distribution 9u which takes the following form,
qs argmax,ee / ((0) log
Lg (O)p(( (0) do
(35)
argmin gE@ F2, (4 (0) + log (0)) + (T.(4Il 4g):
(36)
where for the equality step, we switched from maximization to minimization and replaced the negative log-likelihood with a loss &x.
44 (0) are site-functions which aim to approximate the likelihood, but they also play a similar rule to Lagrange multipliers in ADMM. Ashman et al. (2022, Alg. 1) then update tę as follows:
tu (0) + (0)
tu (0)
(0) log tu (0) + (0) + (log gx (0) - log ( (0)). (37)
g (0)
Finally, the server update in Ashman et al. (2022, Alg. 1) is given by,
(0)
(0)
II
x (o II (0).
(38)
where the last step holds when & (0) = 1 and q3 (0) (0) To (0), which is the initialization condition in Ashman et al. (2022, Alg. 1). The last step can be seen by recursively applying the definition of 4g
and noticing that successive terms in the product cancel out in a telescoping fashion.
We now discuss the difference between Bayesian-ADMM and PVI (Eqs. 36 to 38) in more detail. Unlike PVI, Bayesian- ADMM uses the learning rate in all three updates, whereas PVI does not use any learning rates in its raw form (Ashman et al., 2022). The learning rate is essential for convergence and it is well-known that some splitting algorithms can diverge if learning carting-rates rates are not carefully chosen. To get method that works well in practice, Swaroop et al. (2025) used a learning rate p in the dual update with a justification of damping (similarly to Ashman al. (2022)) but they did not use in the update of qi-
18

---

Published as a conference paper at ICLR 2026
The server update of PVI in Eq. 38 makes the method more similar to the Alteraating Minimization Algorithm (AMA) of Tseng (1991) than ADMM. AMA does not use a proximal-term on the server, which is similar to the PVI update. However, AMA also uses step-size in the primal and dual updates, whereas PVI only uses a step-size in the dual, which is referred to as a damping parameter. From the convergence proof of AMA, it is clear that step-sizes required to obtain overall convergent method, see Tseng (1991, Eq. 3.4d).
Similar to Bayesian-ADMM,, PVI with damping can also be connected to the Bayesian learning rule. Taking the optimality condition of Eq. 36 for an exponential-famin q(0) and inserting it into the dual step Eq. 37 with damping, we get:
nh - (1 mm + pV pu [6x] p=
(39)
where we used & (0) ex T())). This is a reformulation of the Bayesian learning rule (Khan Rue, 2023) terms of local parameters, see Khan & Rue (2023, Eq. 60); see also Khan Lin (2017), except that the natural gradients are evaluated at 2 instead of Ja.
C.4 DERIVATION OF BAYESIAN DUALITY AS STATIONARITY CONDITION OF A LAGRANGIAN
Here, we show how suddle-points of Lagrangian recover the Bayesian duality structure of Fig. (right) and Eq. 13. Consider the following,
K
Lagrangian (14g 72 (9)) = Σ (((()))) + (1 12 1gy1 + KT(gg mO)
(40)
At a saddle-point, all the derivatives with respect to the primal- and dual variables vanish. Taking the derivative with respect to 77 and setting it to zero we immediately get f4 = M4g from Eq. 13 due to linearity of the inner product. Taking the derivative with respect to μ4, and setting it to zero directly leads to m = [ Taking the gradient with respect to 24 and setting it to zero, we get
K
K
V VC((gg)) ")
k=
0
km0
which follows as in App. B.2. The last part in Eq. 13, p1" = VA(Ag" "), simply follows from the EF's dual map.
The Lagrangian is a direct extension of that used for Eq. I, where (0z, VA: 0g) are replaced by (Mg 775 14g). One non-trivial change is the KL term, which makes sense given its presence in the VB objective. The linear terms stem from the constraint H4 = μ4g in the p-space. This choice of coordinates is natural and used in other works as well (Adam al., 2021; Opper & Winther, 2005). Exponential families are also derived using such constraints, see App. B.1.
D DERIVATION OF THE ADMM EXTENSIONS IN SEC. 3.4 AND 3.5
D. RECOVERING FEDERATED ADMM
As mentioned in Sec. 3.4, recover federated ADMM when using Gaussians with fixed covariance. These form an exponential family as described in Eq. 21, see also Table 3. From Eq. 21, we see that the natural gradients are simply an expectation of the regular gradients. Since T(Ø) O, and the natural gradients are just gradients as in regular ADMM, the dual variable 7) has the same size in ADMM, that is, we set 7) Vk. As dual variables are the natural gradients, we also get Vk = - Egg [+)]
The KL divergence between two Gaussians with fixed covariance is simply the squared Euclidean norm (up to a constant), and with μ m, the client update in Bayesian- ADMM in Eq. 15 becomes
m argminm, E T) fl (0)) + V m 2
P [m my
(41)
Using the definition of the expectation parameter and T(Ø) 0. we can write v m. Ey [v] Ø) and using linearity of the expectation, Eq. 41 then recovers Eq. 22 from the main text. The dual update of ADMM is recovered from Baye by plugging in 7 Vk. d my, Zg my-
19

---

Published as a conference paper at ICLR 2026
The same arguments as for Eq. 41 allow to write Bayesian-ADMM server step as,
K
K
mg < argminm, B T) [o(
my
m 11?
(42)
2
k=
Using the delta method, we see that this recovers the server update of regular ADMM.
D.2 DERIVATION OF NEWTON-LIKE ADMM VIA FULL COVARIANCES
To derive the Newton-like method, we use the following forms for the posterior distributions
dg (0) NN mg: STt (0) N m t
(43)
These fit our setup (see Table 3) through,
A = (Sm, -),, μ m, mm +
T(Ø) = (0,00
(44)
Following Khan & Rue (2023, Eq.10-11) the natural gradients for Gaussians with full covariance
V Ey [&x] = (g. H.m. Hba).
7% = (V.. Vu),
(45)
where H₂ = Egg [ is now the full expected Hessian which determines the form the second dual variable to be matrix. Since 77 is the negative natural gradient, get that V H₃ to be the Hessian. To rewrite the client problem in Eq. 15, we use linearity of expectation and definition of the expectation parameter to get (Ms Mx) Eng (7 T(0)) and expand (n,, T(0):
([VE, - - VR), (0, 00 = vi 0 10 V arrive at:
mc S + argmin my S. Eg (0) + v T V + pKL(gg ll)
(46)
where the KL-divergence is given as,
KL(gx (g) (u((s S,)) + log (Sa) + || 100m--mpll 111 const.
(47)
Substituting the form of m7 and D as well as Ag. the dual update in ADMM (Eq. 16) is,
Vk Ve + p(S Symg),
(48)
V Vv + (S Sg))
(49)
Finally, to write the update of Bayesian-ADMM (Eq. 17), first notice that it has simple closed form in natural parameters. To see this, we take the derivative in H and set it to zero to get,
K
K
Ag no
7
+
(xg A&) = 0,
(50)
k=1
k=1
where no VuLo(Hyy) Rearranging and denoting a = 1/(1 (1 + pK), this is equivalently written as
> K
A + Σ n
(51)
k=
Using the form of X from Eq. 44, the form of m) from Eq. 45 and simple isotropic Gaussian prior with precision this can be simplified as:
Sg f- (1 a)) Mean (Si.)) + a [8I + Sum (V
(52)
m, S [(1 Me.
aSum (1:am()
(53)
We use this full Newton-like method for the experiments Figs. 4, and 5b where the variational subproblem in Eg. 46 is solved using VON (Khan & Rue, 2023, Eq. 12).
20

---

Published as a conference paper at ICLR 2026
Convergence in single communication round. Similar to message passing algorithms (Koller & Friedman, 2009; Winn et al., 2005), Bayesian- ADMMM can converge in a single round of communica- tion, which is unlike regular ADMM and BregmanADMM (Wang & Banerjee, 2014).
Proposition D.1. (0) = (ca. T(0)) (for example, ex (x) 1 Ae + b 0 in the case full Gaussian distributions), step-size p = 1 K and initializing the server at the prior Xg = Th then Bayesian-ADMM (Eqs. 15 to 17) converges to the solution of Eq. 9 after one communication round.
Proof. From the optimality condition of Eq. 15, we get that A = >g + 1 c as we initialize 7, = 0. The server update in Bayesian-ADMN. takes a simpler form, as derived in Eq. 51,
K
X = (1 a)
Σ
n
(54)
k=1
k=
Noticing that 7 C₁: after the dual update, the server update is:
K
K
K
1
2 = (1
no
+
ck P
+a
No +
Σ
Ck
70
+
Ck
(55)
K
k=1
1=1
where in the second equality we used == 1/K. This is the optimality condition of Eq. 9 for a conjugate prior 7h- In the second client step, the Lagrange multiplier term cancels out with the loss (since 7) C2 and (x = - -g T())). and minimizing the KL leads to A = Xg. and therefore all clients and the server reach the optimal solution after one communication round.
D.3 DERIVATION OF ADAM-LIKE ALG. (IVON-ADMM)
To derive the Adam-like method shown in Alg. 1, we use a Gaussian posterior approximation with diagonal covariance. Denoting (0) = (O m, diag(s) where s is the precision vector, similarly to the previous section, we get the following setup:
À (sm, - }
H = [m, m + 1/s),
T(Ø) = (0, 00,
(56)
For the natural-gradien we get:
1 hs) V F. [6x] = (g.- hg,,; 2
(57)
Here, all operations (such as sm) are element-wise The last equation is from Khan & Rue (2023, Eq. 10-11) where we denote g = E,g, V(x) and h2 = E,,, [diag( (x)
Substituting the definitions of diagonal Gaussian an EF (see Table 3 and Eq. 56) into the BMMM updates straightforward mechanical calculation similar to the previous section. We denote the dual variables 77 (Vx, - match the structure of the natural parameter A = (s - = which is also a Lagrange multiplier (see App. B.I). Unlike A which has to be valid natural parameter, 77 be arbitrary (and also zero), so we do not use V, but rather just Vk for the first argument.
We expand the linear term in Bayesian-ADMin get,
(, T())) ((Vx, uk), (0, 0²) v 0
u (0) vī - 10 diag(u: ).
(58)
Inserting this into Eq. 15, we arrive at:
mR S < argminm..et ss. Equ. (4 vx 0 10 diag 00| + pKI((u 1l
(59)
where the KL-divergence is given by:
KL(qk (g) = Σ
log s;
+ ||m myll
+ const.
(60)
Alg. uses IVON (Shen al., 2024) to minimize Eq. 59. We slightly generalize IVON handle a general prior and the two Lagrange multipliers. The resulting method is shown in Alg. 2 with modifications over Shen al. (2024) highlighted in red.
The IVON method in Alg. 2 minimizes the following objective function,
XE, [())) - diag(u)o) ( (L
(61)
where ((0) is a generic loss and Ve(()) denotes a stochastic gradient. The objective Eq. 61 matches the subproblem in Eq. 59. We provide details on how to use Alg. in the following App. D.4.
21

---

Published as a conference paper at ICLR 2026
Algorithm 2 IVON for minimizing Eq. 61. Hyperparameters ho, or, B1, 82 are chosen following Shen et al. (2024, Appendix A). highlight in red the extra terms due to Lagrange multipliers.
Inputs: Loss L, prior p(0) = (0 mp: o loss scaling X, Lagrange multipliers
Initialization:
mg. h hgg g 0, 1/(2o²))
I: for = 1, 1, do
2: g - V((0), where 0 ~ (m, x²)
3: h g. 80 m)//² -u
4: g g+ šig + +1 (1-31))
5: h 3,h + (1- 82) 32) h/2//hh//
(h)
6: m m ag +V - um + 6(mm= mp (hpp)//h+0
7: o + Ahh+++
8: end for
9: return (m, 1/02)
D.4 DETAILS ON USING IVON (ALG. 2) TO IMPLEMENT IVON-ADMM. (ALG. 1)
Here, we provide additional details on how the IVON step in Alg. is implemented in our experiments. As commonly done, we consider a tempered version of the Priational-Bayestie problem in Eq. 9,
g argmin EQ
Eu, [0) + rKL(gg To).
(62)
where - > 0 is a temperature parameter. Dividing the whole objective by t does not change the minimizer and gives rescaled client losses ls (0)T T the original objective Eq. 9.
For the tempered problem, the local client optimisation problem in IVON-ADMM reads as follows,
argmin
Egg
(0)/T++ 100= diag(u:)
+ pKLL(q 1g),
(63)
which when dividing by P corresponds to the loss in Alg. where per-example loss functions and N₂ is the number of data examples on the client k. We now bring this problem into form that allows us to directly apply Alg. 2. Dividing by p we get:
1. e denote the
NR Egg argminge pt
O
VE
10 diag(u;)) ) + KL(qx 9g). N₂
(64)
N
N
Matching the forms of the above equation with Eq. 61 gives us Ne/(p7), v (T (T/N+)vis
(7/Nx) For each client subproblem in line Alg. 1 we are calling Alg. subroutine I
with the following inputs: lx, (m,. )) No. pr No Vk. and No Ug. Alg. then returns (m4 s4) which a minimizer of the client variational objective in Bayesian-ADMM.
EXPERIMENTAL DETAILS FOR THE ILLUSTRATIVE EXAMPLES
E. TOY EXAMPLE FIG. 4
For the toy example in Fig. 4 we run the standard federated ADMM algorithm (with tuned quadratic è regularizer 2 |[)l? and Aoyesian-ADMM with full covariances from App. D.2. The dataset consists of the shown points, where we append single dimension to the features to have a bias present in the linear classifier. We use binary cross entropy loss (logistic regression) and train with 0.2. The step-sizes are set to p = 0.2 for both methods.
E.2 ONE-STEP CONVERGENCE ON RIDGE REGRESSION IN FIG. 5A
The one step convergence plot in Fig. 5a uses the same data as in App. E.3, but we consider a linear re- gression loss l (0) | y| and regularizer (o) / |
2 Note that this fits the assumption
22

---

Published as a conference paper at ICLR 2026
of the one-step convergence result in Prop. D.1 for sufficient statisties T(Ø) (0. 00 We again compare the standard ADMM method to our Bayesian- with full Gaussian covariances as well as direct application of BMMM (Wang & Banerjee, 2014) which uses fu-coordinates in the dual update of Te.-
E.3 LOGISTIC REGRESSION CONVERGENCE PLOT IN FIG. 5B
The setup is multiclass Bayesian logistic regression problem on MNIST. We have 10 classes, split onto K 5 clients where the data is split up in a heterogeneous way: (00.1, (2.3), (4,5), (6,7). (8,9). We run PVI, PVI with damping (Ashman et al., 2022; Swaroop et al., 2025), BrMM (Wang Banerjee, 2014) and Bayesian-ADMe ADMM for full Gaussian posteriors. BMM imple- mented similar to BDMM but using the fu-coordinates in the dual update. BregManADMM applicable to our setting, since the KL-divergence between two exponential families a Bregman divergence. The client subproblem is solved by running the Variational Online Newton method (Khan Rue, 2023, Eq. 12) until convergence. PVI with damping uses p = 1/K on the dual update (larger a did not converge, smaller were slower), whereas for method set p 1.
F HYPERPARAMETERS, ARCHITECTURES AND DATASETS
We provide further details about the hyperparameters, model architectures and datasets used in the experiments in Tables 1 and 2.
For the hyperparameters in IVON-ADMM we perform a coarse grid search over the step-sizes and y in Alg. the prior precision 6 and the temperature т. Compared to FedProx, there are only two additional hyperparameters (the dual step-size and the temperature). We found the choice 0.1 to work well for both and we did further tune them. The subproblem is solved using the IVON optimizer outlined in Alg. 2 whose hyperparameters we set according to the recommendations given in Shen et al. (2024, Appendix A).
For all baseline methods, we follow the hyperparameter tuning procedure from Swaroop et al. (2025). We also ensure that the random dataset splits for our experiments match those from the reported results in Swaroop et al. (2025).
Details on heterogencous splits. We follow the heterogeneous sampling procedure from Swaroop et al. (2025, App E1) to generate heterogeneous splits of MNIST, FashionMNIST and CIFAR-10. When there are 10 clients, 90% of all data is usually within clients, with 2 clients having 50%. Within each client, usually 60-95% of client data belongs to just classes.
F.1 MNIST
We train the full MNIST dataset of 60000 examples. All methods batch-size 32. The model is fully connected neural network with sigmoid activation functions and two hidden layers with 200 and 100 neurons, respectively.
Results for baselines are taken from Swaroop et al. (2025) as our setup is identical to theirs.
For the highly-heterogencous 100-client setting, we use the same split in McMahan al. (2016), where each client has data from 2 classes only (300 examples from each class in each client). This particularly difficult setting as each client only has data from 2 classes. We perform hyperparameter sweeps over baselines following Swaroop al. (2025, App E4), except we also allow each hyperparameter to be one order of magnitude smaller (as there are more clients).
F.2 FASHIONMNIST
For the 10 client splits, we train on 10% of the FMNIST dataset. The 100 client split uses the full dataset. All methods use batch-size 32, and the model is again the same fully connected neural network as for MNIST.
Results for baselines taken from Swaroop et al. (2025) setup is identical theirs.
23

---

Published as a conference paper at ICLR 2026
Table 4: Additional results showing test accuracy and test NLL for 10, 25 and 50 rounds, with mean and standard deviations over 3 runs. IVON-ADMM significantly outperforms all baselines. Averaging over the posterior IVON-ADMM (as opposed IVON-ADMM m) often improves performance.
Test accuracy (1 larger is better)
Test NLL
smaller better)
Scenario
Method
10 rounds
25 rounds 50 rounds
10 rounds
25 rounds
50 rounds
MLP, 10 clients
FedAvv/FedPProx
69.9+0.4
74.7+0.6
76.9+0.9
0.80+0.04
0.71+00.33
0.66++0.03
heterog. 10% FMNIST
FedDyn
73.0.10.6
74.610.4
74.6:0.55
0.75
0.70110.0
0.77:00.05
FedLap
71.31:09
74.310.4
77.6110.
0.75.10.04
0.71:10.06
0.6510.05
FedLap-Cov
74.6210.7
78.321.0
80.540.6
0.70:10.04
0.63:10.04
0.60:1..44
IVON ADMM@M, my
77.0=0.88
8144004
82.1#0.1
0.65=0002
0.53+0.01
0.51=0.00
IVON ADMM
77.0++0.8
81.5+0.55
82.3+0.2
0.65+00.22
0.52+0.01
0.50++0.00
F.3 CIFAR-10
We train on the full CIFAR-10 dataset. We consider two models, one is convolutional network used in Swaroop al. (2025) (from Zenke et (2017)) which uses dropout and the other model is the convolutional network used in Acar et al. (2021) which does not use dropout and has less convolutional but more fully connected layers. All methods use batch-size 64.
Results for baselines for the first CNN are taken from Swaroop et al. (2025) our setup is identical to theirs. For the second CNN, we rerun hyperparameter sweeps similar settings to Swaroop et al. (2025, App E6). In addition to those parameters, 10 4). We also try SGD with learning rate 0.1 and sweep over learning rate decay (0.992 or 1), so that we follow the hyperparameter sweep settings from Acar al. (2021).
also sweep
Adam learning rate (10 -3 or
F.4 CIFAR-100
We train the full CIFAR-100 dataset, split across 10 clients using the same Dirichlet parameters. We use a ResNet-20 model (He et al., 2016) which has around 250k parameters. All methods use batch-size 64. No data augmentation is used.
We use the same hyperparameter sweeps as in Swaroop et al. (2025) for CIFAR-10. We note that FedLap-Cor Cov is now prohibitively slow, since computing the diagonal Laplace approximation estimates the covariance every communication round takes a long time, and so we do not provide results for it.
G ADDITIONAL RESULTS
The additional results shown in Table confirm the findings from the main paper. IVON-ADMM gets the highest accuracy and lowest test-loss across communication rounds, improving significantly also over the recent previous state-of-the-art Bayesian federated learning FedLap-Cov which uses more expensive Laplace approximations.
We also perform various ablation studies in Bayesian-ADMM and IVON-ADMM.
G.1 ABLATION: PVI WITH IVON
PVI with damping (Ashman et al., 2022) implemented with IVON (Shen et al., 2024). The resulting method, which we call IVON-PVI is a special case of the server update.
Here, we compare IVON-ADMR
IVON-ADMM, where the client step size is set to p = and one uses a We use finely tuned dual damping, there are other step-size hyperparameters. We run ResNet-20 on CIFAR-IOS 100, as in the main paper. The results are summarized in the following Table 5,
We see that using the Bayesian -DMM update gives an advantage over PVI in the early rounds, while reaching comparable accuracies at the end. A benefit of PVI is that there no step-sizes to tune other than the dual damping. Finally, we note that PVI has not been implemented so far with oate-off-the-at methods like IVON, and IVON-PVI is special case of our IVON-ADMM.
24

---

Published as a conference paper at ICLR 2026
Table 5: Test accuracy and NLL after 25, 50 and 100 rounds, with mean and standard deviations over 3 runs. IVON-ADMM has faster convergence in early rounds than IVON-PVI, but they reach overall comparable accuracies after 100 rounds.
Test accurucy (1 larger is better)
Test NLL (1 smaller better)
Scenario
Method
25 rounds
50 rounds
rounds
25 rounds
50 rounds
rounds
ResNet-20,
WON-PVI@1H
31.5:10.5
41,240.6
47.420.5
2.740..0
2.2:0..0
2.0:0.00
10 clients,
IVON-PVI
31.5:10.
41.21:06
47,510,5
2.7:0000
2.240.0
2.0:0.00
heterog. CHFAR-100
WON-ADMM@m,
40,0=1.0
46.2:004
46.520.6
2.3:00.0
2.1:0.0
2.340.1
IVON-ADMM
40.0:110
46,2::04
46.620.6
2.340.00
2.1=0.00
2.240.1
G.2 ABLATION: SENSITIVITY TO INVERSE CLIENT STEP SIZE p AND TEMPERATURE T
We further study the sensitivity of IVON-ADMM to p and т. We again use the ResNet-20 on CIFAR-100 and highlight the parameters used in the main paper.
The results for varying temperature are summarized in Table 6.
Table 6: Temperature ablation for IVON-ADMM. Test accuracy and after 25, 50 and 100 rounds, with mean and standard deviations over 3 runs.
Test accuracy (t larger better)
Test NLL (L smaller is better)
Scenario
25 rounds
50 rounds 100 rounds
25 rounds 50 rounds 100 rounds
0.02
38. 6.0..88
23.9199.4
25.5+7.66
2.510.1
5.3+1..4
4.1+0.1
0.05
42.0
42.0+:0.6
41.1+-0.6
2.2+0.00
2.7:0.11
2.9+00.1
0.1
3..911.2
46.110.66
46.610.8
2.310.0
2.1:0.0
2.2:0.11
ResNet-20,
0.2
35.7+1.0
45.1=110
47.910.7
2.510.1
2.1:0.00
2.0:0.0
clients,
0.5
27.1+1.1
36.6+1.6
40.2+1.6
2.900.1
2.5+0.1
2.3+-0.1
heterog. CIFAR-100
20.0:11.2
27.5114
30.411.2
3.310.1
2.910.1
2.811..1
Smaller temperatures lead to faster convergence, but the overall best final accuracy is reached for
т = 0.2. Larger temperatures tend to degrade the performance. All experiments in the paper use 0.1 (highlighted as gray row in the table), but we can see, an additional tuning of 7 can further improve the results.
The results for varying client inverse step-size p summarized in the following Table 7.
Table 7: Client inverse step-size p ablation for IVON-ADMM. Test accuracy and NLL after 25,50 and 100 rounds, with mean and standard deviations over runs.
Test accuracy (t larger is better)
Test NLL (L smaller is better)
Scenario
25 rounds
50 rounds
100 rounds
25 rounds
rounds
100 rounds
0.1
1.0++..0
1.0+0..0
1.0++0.0
diverged
diverged
diverged
0.2
26.7410.6
14.3+10.8
14.0:1:0.6
diverged
diverged
diverged
0.5
39.8+1.2
46.2++0.8
46.6+0.9
2.3+0.0
2.1+0.0
2.2++0.1
ResNet-20,
31.410.6
40.610.77
45.2110.
2.710.0
2.310.0
2.1:0.00
10 clients,
2.0
25.4+00.7
32.6+0.77
38.9+0.77
3.0+0.00
2,7+0..0
2,4+00.0
heterog. CIFAR-100
5.0
18.210.9
23.71.0.7
29.0.10.8
3,4.0.0
3.110.0
2.940.00
Smaller p. which gives the client more freedom deviate from the server's consensus solution, leads
to instabilities. Larger constrains the client to be closer to the server's 's solution, and while leading stable convergence degrades the overall performance. The value of p = 0.5 which we used the main paper is at sweet-spot, balancing this trade-off.
25


Reviews:
[{"review_id": "AzsMwrRDyW", "reviewer": "Reviewer_B61X", "summary": "The authors propose an analogous extension to federated ADMM in the context of variational inference. This provides an extension of ADMM-like federated procedures based on the duality of various exponential distributions.", "questions": "Why is it valid to just switch back and forth between variational inference and MLE? I was not convinced that equation 5 was equivalent to equation 1. Along this line, the notation in the paper blurred distinction between parameters and distributions (see equation 4 for example and line 976).\nIn equation 12-14, $\\bar{\\lambda}$ is not defined. Perhaps include an update equation for it. Is that because it is a deterministic mapping from $\\bar{\\mu}$? Similarly for $\\bar{q}$, is the update from equation 4?\nUnder what circumstances does BayesADMM converge?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a novel ADMM-like extension to federated learning, with good experimental results.", "weaknesses": "It seemed that the argument was more by analogy than exact equality. The claim that ADMM is recovered exactly is misleading because it requires an approximation and therefore is not necessarily recovered exactly.\n\nA couple minor points:\n- In equation 3, a subscript k is missing\n- On line 195 A* should be defined in the main text rather than just in the appendix.\n- In figure 4, it would be helpful to mention that each line is numbered with the iteration number.", "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": "luNx6OA9N7", "reviewer": "Reviewer_NFZi", "summary": "Authors propose a new Bayesian approach to derive and extend the federated Alternating Direction Method of Multipliers (ADMM). We show that the solutions of variational-Bayesian objectives are associated with a duality structure that not only resembles ADMM but also extends it， which opens a new direction to use Bayes to extend ADMM and other primal-dual methods.", "questions": "In the federated learning ADMM framework, there are theoretical guarantees for communication complexity and iterative complexity. Can the author briefly discuss the communication complexity and iteration complexity of Bayesian ADMM.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Authors introduced a Bayesian duality, from which an extension of ADMM that optimizes over distributions naturally follows. For Gaussians with fixed variance, they recover regular ADMM and general Gaussians give Newton-like methods and IVON-ADMM. These show good performance when compared to recent baselines. Other approximating distributions may lead to new interesting splitting algorithms, and more generally, which opens up new research paths to extend and improve primal-dual algorithms using Bayesian ideas.", "weaknesses": "In the federated learning ADMM framework, there are theoretical guarantees for communication complexity and iterative complexity. Can the author briefly discuss the communication complexity and iteration complexity of Bayesian ADMM.", "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": "KjcOoJZHqH", "reviewer": "Reviewer_DtMH", "summary": "The paper proposes a general Bayesian perspective on federated ADMM. Under this framework, the classical ADMM emerges as a special case corresponding to isotropic Gaussian posteriors, while more expressive exponential-family posteriors yield new variants. One such variant called _IVON-ADMM_ is derived using diagonal Gaussian covariance and is claimed to perform better in heterogeneity and uncertainty cases of federated learning.", "questions": "1. What specific mechanism makes BayesADMM or IVON-ADMM handle heterogeneity and uncertainty better than existing method? Can you clarify the role of the posterior covariance in this improvement and demonstrate it experimentally?\n2. In Fig. 4, how are the gray uncertainty contours generated? Are they derived from posterior covariance?\n3. Why does the exposition (Eqns. 1, 2, 10, 11) rely on the plain Lagrangian, while the ADMM's implementation (Eqns. 3, 12-14) uses the augmented Lagrangian version?\n4. The notation $\\hat{\\lambda}_k$ seems to play dual roles -- as natural parameters of site functions in Sec. 3.2 and as dual variables in Sec. 3.3. Are both interpretations valid? If so, explain their precise connection.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "__Conceptual novelty:__ The paper establishes a novel Bayesian duality perspective that unifies ADMM and VB under a single framework. This is an interesting connection that could inspire extensions of primal-dual optimization methods.\n\n__Clear motivation and exposition:__ The introduction and backgrounds are well written and clearly position the work relative to prior ADMM and PVI approach (Swaroop et al., 2025).\n\n__Framework generality:__ The proposed Bayesian duality formulation provides a principled way to derive new ADMM-like algorithms by changing the exponential-family posterior.\n\n__Readable presentation:__ For the most part, the paper is well structured and logically progress from classical ADMM to its Bayesian interpretation and finally to _IVON-ADMM_.", "weaknesses": "__Soundness of the Formulation:__ While the high-level idea is promising, the derivation in section 3.3 raises concerns about mathematical consistency:\n- The \"Bayesian ADMM\" updates (Eqns. 12-14) are expected to follow from alternating optimization of the Lagrangian in Eqn. 11. However, the replacement of the dual update term $\\mu_k - \\bar{\\mu}$ with $\\lambda_k - \\bar{\\lambda}$ lacks justification within the Lagrangian formulation. The reasoning provided in Appendix E.2, appealing to Bayesian intuition, seems heuristic rather derivational.\n- Equation 11 itself may need reconsideration: shifting the linear term $<\\hat{\\lambda}_k, \\mu_k - \\hat{\\mu}>$ between the sub problems while keeping others fixed breaks symmetry between local and global updates, potentially undermining the claimed equivalence to ADMM.\n- Overall, the theoretical grounding of the \"Bayesian duality\" remains somewhat fragile: the proposed updates look reasonable by analogy but are not rigorously shown to correspond to valid saddle-point dynamics of the stated objective.\n\n__Relation to Existing Work:__\n- The method appears to be a straightforward extension of _PVI_ with modified update equations. The novelty over _PVI_ is mainly the introduction of the step size $\\rho$ and reinterpretation of dual variables. The paper should make a stronger argument for why this constitutes a _fundametal_ new framework rather  than a variant of _PVI_ with heuristic scaling.\nA direct empirical or theoretical  comparison with _PVI_ (as Eqn. 4) is missing. Including  such results would make the claimed advantages more credible.\n\n\n__Experimental Evaluation__:  The experiments, while broad, are not yet conclusive about the claimed benefits in heterogeneity and uncertainty.\n- Figure 3-4 provide illustrative  but small-scale toy examples; they show qualitative improvement but not a clear quantitative advantage.\n- The key claim that _IVON_ handles heterogeneity better by leveraging posterior covariance is not substantiated with ablation or analysis showing the role of uncertainty.\n- Comparisons with both BayesADMM (without _IVON_) and _PVI_ are missing. Including them would help isolate what _IVON_ adds.\n- The computational overhead relative to FedDyn should be quantified to see the computation gain compared to the performance gain.\n\n__Clarity and Notation:__ Several presentation issues reduce readability and reproducibility.\n- Step 3 of Fig. 2 is valid only for $\\alpha=\\frac{1}{1+\\rho K}$.\n- Many symbols and methods are used before being introduced:\n    + BLR first appears on lines 294 and 297 without citation.\n    + $q_{1:K}, \\hat{t}_{1:K},\\bar{q}$ are used before definition in line 215.\n- The discussion of natural vs. ordinary gradients is confusing. The paper should use distinct and properly defined notations for both.\n- Section 3.2  could be organized better: the correspondence between Eqns. 2 and 10 is conceptually interesting but presented unclearly, with inconsistent references to $\\hat{t}_k, \\lambda_k, \\text{ and } \\mu_k$.", "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": "tdhKUnWbGZ", "reviewer": "Reviewer_WVuS", "summary": "The paper proposes a novel method that combines Bayesian variational inference with the Alternating Direction Method of Multipliers (ADMM) for federated learning scenarios. Based on the concept of \"Bayesian Duality,\" the authors extend traditional parameter optimization to distribution optimization, developing the BayesADMM algorithm. This method not only encompasses standard federated ADMM as a special case (when using fixed-variance Gaussian distributions) but also naturally extends to more complex distributions, producing update rules similar to Newton's method. The authors further derive the IVON-ADMM variant, a computationally efficient implementation suitable for deep learning models. Experimental results show that this method outperforms existing baselines on multiple datasets (MNIST, FashionMNIST, and CIFAR-10), particularly excelling in heterogeneous data distribution scenarios.", "questions": "1.In practical federated learning scenarios, clients often have different computational capabilities and communication bandwidths. How does BayesADMM adapt to this system heterogeneity? Specifically, what is the algorithm's robustness when certain clients cannot complete the full variational inference update?\n2.From Figure 3(a), it can be seen that PVI diverges without damping while BayesADMM remains stable. Could you provide a detailed analysis of BayesADMM's convergence guarantees, particularly theoretical guarantees for non-convex optimization problems?\n3.IVON-ADMM uses diagonal covariance approximation in its implementation, which may lead to information loss in deep learning models. How much performance improvement do you think using low-rank approximation or other covariance structures would bring? How much would the computational overhead increase?\n4.In Algorithm 1, you mention \"implementation details in Appendix G,\" but the provided PDF excerpt does not contain this section. Could you briefly explain the key implementation differences between IVON-ADMM and standard ADMM, particularly the techniques used when handling high-dimensional parameter spaces?\n5.The paper states that BayesADMM can converge in just one communication round for certain loss functions. Is this property limited to specific types of loss functions? Could you provide more general conditions that specify when the algorithm can converge quickly?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper proposes a novel theoretical framework that integrates Bayesian variational inference with the ADMM algorithm, providing a unified perspective for federated learning. This approach of connecting optimization algorithms with probabilistic inference has theoretical depth.\n2. The authors clearly demonstrate how traditional federated ADMM serves as a special case of BayesADMM under specific conditions. This theoretical connection provides a new perspective for understanding existing methods and naturally leads to more powerful extensions.\n3. The implementation details of IVON-ADMM showcase a practical and efficient algorithm variant. Compared to existing federated learning optimizers (such as FedAvg and FedAdam), it incurs limited additional computational overhead while delivering significant performance improvements.\n4. The experimental design is comprehensive, covering both homogeneous and heterogeneous data distribution scenarios, and validates the method's effectiveness across multiple standard datasets. Particularly, Figure 3(b) demonstrates BayesADMM's property of converging in just one communication round for certain loss functions, which is a compelling empirical result.\n5. The authors situate their method within the broader context of optimization and Bayesian inference literature, clearly highlighting connections and distinctions with prior works such as Partitioned Variational Inference (PVI) and Bregman ADMM.", "weaknesses": "1.Although the paper claims BayesADMM is theoretically superior, it lacks detailed analysis of computational complexity. Specifically, how do the algorithm's computational and communication costs scale with model parameter size when using more complex distributions (such as full-covariance Gaussians)? This is crucial in practical federated learning scenarios.\n2.The experimental section only reports average performance without providing standard deviations or statistical significance tests, making it difficult to assess result reliability. Given the stochastic nature of federated learning (such as client selection and data partitioning), such statistical analysis is particularly necessary.\n3.Table 4 shows IVON-ADMM's performance across multiple rounds but does not analyze the trade-off between convergence speed and communication rounds. In practical applications, early stopping may be more useful, but the paper does not explore performance comparisons of different methods under limited communication budgets.\n4.While the authors mention the concept of \"Bayesian Duality,\" their explanation of its theoretical foundation is not sufficiently deep. The derivation of the BayesADMM algorithm in Section 3.3 is relatively brief, and certain key steps (such as the transition from Equation 26 to 27) lack detailed explanation, which may affect reader comprehension.\n5.The paper does not adequately discuss the impact of hyperparameter selection, particularly the sensitivity of regularization parameter ρ and temperature parameter τ. Although the appendix mentions hyperparameter search, it lacks systematic analysis of how these parameters affect algorithm performance.\n6.The comparison with recent state-of-the-art federated learning methods (such as Scaffold and FedProx) is not comprehensive. While the paper compares with FederatedADMM and BregmanADMM, these methods are relatively outdated, and more comparisons with recent works should be included.", "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)