Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach
Authors: Han Guo
Philip Greengard
Hongyi Wang
Andrew Gelman
Yoon Kim
Eric Xing
Venue: ICLR 2023 poster (2023)
Keywords: federated learning
variational inference
expectation propagation

Abstract: The canonical formulation of federated learning treats it as a distributed optimization problem where the model parameters are optimized against a global loss function that decomposes across client loss functions. A recent alternative formulation instead treats federated learning as a distributed inference problem, where the goal is to infer a global posterior from partitioned client data (Al-Shedivat et al., 2021). This paper extends the inference view and describes a variational inference formulation of federated learning where the goal is to find a global variational posterior that well-approximates the true posterior. This naturally motivates an expectation propagation approach to federated learning (FedEP), where approximations to the global posterior are iteratively refined through probabilistic message-passing between the central server and the clients. We conduct an extensive empirical study across various algorithmic considerations and describe practical strategies for scaling up expectation propagation to the modern federated setting. We apply FedEP on standard federated learning benchmarks and find that it outperforms strong baselines in terms of both convergence speed and accuracy.

Published as a conference paper at ICLR 2023
FEDERATED LEARNING AS VARIATIONAL INFERENCE: A SCALABLE EXPECTATION PROPAGATION APPROACH
Han Guo Philip Greengard Hongyi Wang' Andrew Gelmant Yoon Kim Eric P. Xing to
CCarnegie Mellon University, " "Massachusetts Institute of Technology, FColumbia University Mohamed bin Zayed University of Artificial Intelligence, Petuum Inc.
ABSTRACT
The canonical formulation of federated learning treats it distributed optimiza- tion problem where the model parameters are optimized against a global loss func- tion that decomposes across client loss functions. A recent alternative formulation instead treats federated learning as distributed inference problem, where the goal is to infer global posterior from partitioned client data (AI-Shedivat et al. (022). This paper extends the inference view and describes variational inference formu- lation of federated learning where the goal to find global variational posterior that well-approximates the true posterior. This naturally motivates an expecta- tion propagation approach to federated learning (FedEP), where approximations to the global posterior iteratively refined through probabilistic between the central server and the clients. We conduct an extensive empirical study across various algorithmic considerations and describe practical strategies for scaling up expectation propagation to the modern federated setting. We ap- ply FedEP on standard federated learning benchmarks and find that outperforms strong baselines in terms of both convergence speed and accuracy
passing
INTRODUCTION AND BACKGROUND
Many applications of machine learning require training a centralized model over decentralized, het- erogeneous, and potentially private datasets. For example, hospitals may be interested in collabora- tively training a model for predictive healthcare, but privacy rules might require each hospital's data to remain local. Federated Learning (FL. McMaaan er al. 2017 Kairouz et al. [021] Wang et al.)
2021] has emerged data to leave their local devices. FL introduces new challenges on top of classic distributed learning: expensive communication, statistical/hardwaar heterogeneity, and data privacy (Li et al. 2020a).
training paradigm that does not require clients' private
The canonical formulation of FL. treats it as distributed optimization problem where the model parameters are trained on K (potentially private) datasets D Uke(K] Dx.
= arggiing L(0), where L( E logx 0).
Standard distributed optimization algorithms (e.g., data-parallel SGD) are too communication- intensive to be practical under the FL setup. Federated Averaging (FedAvg. McMahan et al.) 2017) reduces communication costs by allowing clients to perform multiple local SGD steps/epochs be- fore the parameter updates are sent back to the central server and aggregated. However, due to client data heterogeneity, more local computations could lead to stale biased client updates, and hence sub-optimal behavior (Charles & Konečný 2020 Woodworth et al. 2020 Wang et al.) 2020a).
An alternative approach is to consider Bayesian formulation of the FL problem (Al- AllSSheiivat al. 2021). ). Here, we are interested in estimating the posterior of parameters p(0 DD given prior p(0) (such as an improper uniform or a Gaussian prior) and a collection of client likelihoods p(Dz (0) that are independent given the model parameters,
p(0 D) < p(0) [[cK P(Dx O).
In this case the posterior naturally factorizes across partitioned client data, wherein the global pos- terior equates to multiplicative aggregate of local factors (and the prior). However, exact posterior inference is in general intractable for even modestly-sized models and datasets and requires approx-
CCode: https: //g h///ggttul thub. com/ www/HanGuo997770 htwpa//s//eepppeaatatrrrrrrreeee
This work was completed while
Han Guo was visiting student MIT.

---

Published as a conference paper at ICLR 2023
imate inference techniques. In this paper we turn to variational inference, in effect transforming the federated optimization problem into a distributed inference problem. Concretely, we view the solution of federated learning the mode of variational (posterior) distribution q Q with some divergence function D-III-) (e.g., KI-diivrgeence,
0 arg maxo q(),, where q()) = arg min min qc (p ( l
Under this approach, clients use local computation to perform posterior inference (instead of pa- rameter/gradient estimation) in parallel. In exchange, possibly fewer lockstep synchronization and communication steps required between clients and servers.
One way to operationalize Eq. [ is through federated posterior averaging (FedPA, Ahedivii et al,, 2021), where each client independently runs approximate inference procedure and then sends the local posterior parameters to the server to be multiplicatively aggregated. However, there is no guarantee that independent approximations to local posteriors will lead to good global approxi- mate posterior. Motivated by the rich line of work on variational inference on streaming/partitional data (Broderick et al. 2013 Vehtari al. 2020), this work instead considers an expectation prop- agation (EP, Minka 2001) approach to FL. In EP, each partition of the data maintains its own local contribution to the global posterior that is iteratively refined through probabilistic
pessage-passing When applied to FL, this results in an intuitive training scheme where at each round, each client (1) receives the current approximation to the global posterior from centralized server, (2) carries out local inference to update its local approximation, and (3) sends the refined approximation to the server to be aggregated. Conceptually, this federated learning with expectation propagation (FedEP) approach extends FedPA by taking into account the current global approximation in step (2).
However, scaling up classic expectation propagation to the modern federated setting is challeng- ing due to the high dimensionality of model parameters and the large number of clients. Indeed, while there is some existing work on expectation propagation-base
federated learning (Corinzia et al. 2019 Kassab & Simeone 2022 Ashman et al. 2022), they typically focus small mod- els (fewer than TO0K parameters) ) and few clients (at most 100 clients). In this paper we conduct an extensive empirical study across various algorithmic considerations scale up expectation propaga- tion to contemporary benchmarks (e.g., models with many millions of parameters and datasets with hundreds of thousands of clients). When applied on top of modern FL benchmarks, our approach outperforms strong FedAvg and FedPA baselines.
FEDERATED LEARNING WITH EXPECTATION PROPAGATION
The probabilistic view from Eq. motivates an alternative formulation of federated learning based on variational inference. First observe that the global posterior p ( DD) given a collection of datasets D Uke [K] Dx factorizes as,
K P () 1-1
K
p (0 D) &
II (0), K=
where for convenience we define Po (0) == p() to be the prior and further use (0) p(Dx O) to refer to the local likelihood associated with k-th data partition. To simplify notation we hereon refer to the global posterior as Pglohal and drop the conditioning on D. Now consider an approximating global posterior 4global (0) that admits the same factorization the above, i.c.
Aglobul (0) o Io (x(( (0). Plugging in these terms into Eq. gives the following objective,
arg max
1 9 where (0)) arg min D (fIPL() k-0
K
K
[II(()) k=O
K
(2)
0
4
1-
Here Q is the variational family, which is assumed to be the same for all clients. This global objec- tive is in general intractable; evaluating I (0) requires accessing all clients' data and violates the standard FL assumption. This section presents a probabilistic message- saassaggpassingg algorithm based on expectation propagation (EP, Minka 2001).
2.1 EXPECTATION PROPAGATION
EP is an iterative algorithm in which intractable target density Pglobal (0) ()) is approximated by tractable density Aglohal (0) using collection of localized inference procedures. In EP, each local inference problem is a function of just pk and the current global estimate, making appropriate for the FL setting.
2

---

Published as a conference paper at ICLR 2023
Algorithm Federated Learning Inference Algorithm 3 Gaussian EP: Server Inference
1: for round f = ,, T do
1: Receive: (Aqu(0; Ang. AA2) hk
2: Sample subset of clients K.
2: geew globbll 9global Tl (Aqx)*
I/ Sec. 2.2.3
3: Broadcast Iglobal (0) the selected clients.
Nglobal < Tgglobal + ServerOptim
Amp)
each client & parallel
Aglohal f- Aglobal + ò ServerOptim
4AA)
5: ( ) + (aeentlnfer(ggaaalll
6: end for
3: Send: 4global (0; Agloball globbal
7: Collect (0) from the selected clients.
Algorithm 4 Gaussian EP: Client Inference
8: Aglobal (0) ServerInfer (Aqx (0))x)
1: Receive: Aglobal (0; Mglobal Aglobal If cavity distribution
end
2: 9-k 9global 9k
10: Return Hglobal
7-k Nglobul Tkk
A f- Aglabal Ag
Algorithm 2 Approximate Inference: MCMC
3: R =0 9re 0 Pk 9-x // tilted inference (Sec. 2.2.2
1: Imput: M
M me Aie + ApproxInference(qy, & Pk 9-x)
2: Si + 0
4: Aqk os qvve 4glohal
I/ client deltas (Sec. A.1
3: for i 1,...., N...,,N do
4: 0 SGDDDpcch log9 0g 1
Ank mh
Nglohale
AA + A Aglobal
5:
S₁ 5 S U
5: 9% qe (AqR)
# local update (Scc. 2.2.3
end
n ne + 5 ClientOptim(AA
7: 77 Ar - EstimateMoments(S))
Ag + A + &ClientOptim/ (Ax)
8: Output: 9 A.
Send: (0; Ar: AAx)
Concretely, EP iteratively solves the following problem (either in sequence or parallel),
(()) x()) where q (0) x 9global (
new (0) = arg min q
Pk (0) (0)
9k (0())
(()
(1)
Here Iglobal (0) and (0) are the global/local distributions from the current iteration. (See Sec. for further details). In the EP literature, 9 (0) is referred to as the caviry distribution and q%x
A.2
and (0) are referred to as the target/approximate rilted distributions. EP then uses 9%
new 0 to
derive gglobal (0). While the theoretical properties of EP are still not well understood (Minka 2001
Dehaene & Barthelmé 2015 2018), it has empirically been shown to produce good posterior ap- proximations in many cases (Li et al. 2015 Vehtari et al. (020). When applied to FL, the central server initiates the update by sending the parameters of the current global approximation Iglobal (0) as messages to the subset of clients IC. Upon receiving these messages, each client updates the re- spective local approximation 9tt (9) and sends back the changes in parameters as messages, which is then aggregated by the server. Algorithms illustrate the probabilistic message passing with the Gaussian variational family in more detail.
Remark. Consider the case where we set 9 (0) & 1 (i.e., an improper uniform distribu- tion that ignores the current estimate of the global parameters). Then Eq. reduces fed- erated learning with posterior averaging (FedPA) from Al-Shedivat et al. (2021), geew arg m D (Px(0) || (()).. Hence, FedEP improves upon FedPA by taking into account the global parameters and the previous local estimate while deriving the local posterior
2.2
SCALABLE EXPECTATION PROPAGATION
While federated learning with expectation propagation is conceptually straightforward, scaling up FedEP to modern models and datasets is challenging. For one, the high dimensionality of the param- eter space of contemporary models can make local inference difficult even with simple mean-field Gaussian variational families. This is compounded by the fact that classic expectation propaga- tion stateful and therefore requires that each client always maintains its local contribution the global posterior. These factors make classic EP potentially an unideal approach in settings where
When the parameters of 9global (0) and qu (0) 's are initialized as improper uniform distributions, the first round (but only the first round) of FedEP and FedPA is identical.
3

---

Published as a conference paper at ICLR 2023
the clients may be eesource-constraine and/or the number of clients is large enough that each client updated only few times during the course of training. This section discusses various algorithmic consideration when scaling up FedEP to contemporary federated learning benchmarks.
2.2.1
VARIATIONAL FAMILY
Following prior work on variational inference in ligh-dimensional parameter space (Graves 2011
Blundell et al. 2015 Zhang et al.] 2019; Osawa et al. 2019), we use the mean-field Gaussian variational family for Q. which corresponds to multivariate Gaussian distributions with diago- nal covariance. Although non-diagonal extensions are possible (e.g., through shrinkage estima- tors (Ledoit& Wolf 2004)). [ we empirically found the diagonal to work well while being simple and communication-efien effiiiennt.. For notational simplicity, we use the following two parameterizations of Gaussian distribution interchangeably.
( - N p4, = N (0;n,A,, where A 1 n
Conveniently, both products and quotients of Gaussian distributions distribtionss-opcratiion commonly used in EP RP-result in another Gaussian distribution, which simplifies the calculation of the cavity distribu- tion 9 (0) and the global distribution Iglobal
2.2.2 CLIENT INFERENCE
At each round of training, each client must estimate 7. (0), its own approximation to the tilted distribution (0) in Eq. We study various approaches for this estimation procedure.
Stochastic Gradient Markov Chain Monte Carlo (SG-MCMC.. SG-MCMC (Welling & Teh
2011 Ma et al., 2015) uses stochastic gradients to approximately sample from local posteriors. We
follow Al-Shedivat et al. (2021) and simple variant of SGD-bbaee SG-MCMC, where
collect a single sample per epoch to obtain a set of samples Sş = (0(1) 0 (..
(N)
The SGD
objective in this case is the unnormalized tilted distribution,
ExED, log p(z 0 +0+-00
0
A
0
n n_
()))
1og9-x(0) x(0)
which is simply the client negative log likelihood (- (0)) plus regularizer that penalizes parameters that have low probability under the cavity distribution (- 1o This connection makes it clear that the additional client computation compared FedAvg (which just minimizes the client negative tog-likelihood is negligible. Given a set of samples S₃ from SG-MCMC, we estimate the parameters of the tilted distribution ( with moment matching, i.e.,
(xx = pt. where f E,, < MomentFstimator((S))
While the mean obtained from S₁ via averaging empirically worked well, the covariance estimation was sometimes unstable. We next discuss three alternative techniques for estimating the covariance.
SG-MCMC with Scaled Identity Covariance. Our simplest approach approximates the covari- ance scaled identity matrix with tunable hyper-parameter deov, i.e., Σ + a I. This cuts down the communication cost in half since we longer have to send messages for the covariance parameters. While extremely simple, we found scaled identity covariance to work well in practice.
Laplace Approximation. Laplace's method approximates the covariance as the inverse Hessian of the negative log-likelline likelihood the (possibly approximate) MAP estimate. Since the exact inverse Hessian is intractable, we follow common practice and approximate it with the diagonal Fisher,
( (oolog ply 0,x) !))
27
Eis + H + ++++E E1)
where Hi ~ diag ExND, plylx,0
diagonal Fisher approximation
(4)
'Specifically, we have the following identities,
N (0; & &N 72, A
N' (0)m, (0; A2) x N (0:m + 72, A + A2),
A2).
N (0;11,,,))
Unlike Al-Shedivat et al. 2021), we do not apply Polyak averaging Mandt et al. 2019 as did find it improve results in our case.
2017
Maddox et al.
4

---

Published as a conference paper at ICLR 2023
This approach samples the input x from the data Dx and the output y from the current model ply x, f), as recommended by Kunstner et al. (2019). The Fisher approximation requires additional epochs of backpropagation on top of usual SGD (usually in our case), which requires additional client compute.
Natural Gradient Variational Inference. Our final approach uses natural-gradient variational inference (NGVI, Zhang et al. 2018: Khan et al. 2018 Osawa et al., 2019). which incorporates the geometry of the distribution to enable faster convergence. Most existing work NGVI assume a zero-mean isotropic Gaussian prior. We extend NGVI to work with arbitrary Gaussian priors- necessary for regularizing towards the cavity distribution in FedEP. Specifically, NGVI iteratively computes the following for f = 1... TNGVI and learning rate ΦNGVI.
I Ver < (DDxlst +I=I) where <- BNGVIS; (1 (--VV) E Ea~
Fisher
[Dx
Here Fisher(-) is the diagonal Fisher approximation in Eq. 4 but evaluated sample of parameters
from qv,, (0), the approximate posterior using the current estimate of L We give the exact
NGVI update (which algorithmically similar to the Adam optimizer (Kingma & Ba 2015)) in
Algorithm in the appendix.
2.2.3 ADAPTIVE OPTIMIZATION DAMPING
Given the approximate tilted distribution q. and the corresponding parameters p E k we can in principle follow the update equation in Eq. to estimate gglobal (9). However, adaptive optimizers have been shown to be crucial for scaling federated learning practical settings (Reddi al.) 2020). and the vanilla EP update does not immediately lend itself to adaptive updates. This section describes an adaptive extension to EP based on damping, in which we to re-interpret damped EP update as gradient update on the natural parameters, which allows for the use of adaptive optimizers.
Damping performs client updates only partially with step size and is commonly used in parallel EP
settings (Minka & Lafferty 2002] Vehtari et al. 2020). Letting Aqu (0) ) o q. (0) (0)/ / (0) denote the client "update" distribution, we can simplity the update and arrive at the following intuitive form /Vehtari et al. 2020) (see Sec. [.. for derivation).
Client: new (e) 91 (0) (Aq)()))
9 (0) 4global (0) II (a)x())))
Server:
k
Recalling that products of Gaussian distributions yields another Gaussian distribution that simply sums the natural parameters, the damped update for r7 given by,
Client:
Th < Mh 5Am:
Server:
Nglobal Aglobal + Σ Ark-
kEK
(The update the precision A is analogous.) By re-interpreting the update distribution Aqu (O; AAg) as "gradient", we can apply off-the-shelf adaptive optimizers.
Client:
7 + m + 5 oAtim(Anc), Server:
Nglobal + Nglobul + 5 optim( Ank).
kEAC
All our FedEP experiments (and the FedAvg and FedPA baselines) employ adaptive optimization.
2.2.4 STOCHASTIC EXPECTATION PROPAGATION FOR STATELESS CLIENTS
Clients are typically assumed to be stateful in the classic formulations of expectation propagation. However, there are scenarios in which stateful clients are infeasible (e.g., memory constraints) or even undesirable (e.g., large number of clients who only participate in a few update rounds, leading to stale messages). We thus additionally experiment with a stateless version of FedEP via stochastic expectation propagation (SEP, Li et al. 2015). SEP employs direct iterative refinement of a global approximation comprising the prior p( ()) and K copies of single approximating factor (0),
Aglobal (0) p(0) (0)
That is, clients assumed to capture the average effect. In practice, FedSEP is implemented in Algorithm via replacing the cavity update (step 2) with 9 (0) ()) x Iglohal (0) (0)/g((0) and removing the local update (step 5).
5

---

Published as a conference paper at ICLR 2023
Dataset
Model
Model size
Clients (train/test)
Examples per client (train/test)
CIFAR-100
ResNet-18
11.2M
500 / 100
100 / =00
StackOverflow
Lincar
5.0M
342,477 / 204,088
397 4301
EMNIST-62
CNN
1.2M
3,400 3,400
198 23
Table 1: Model and dataset statistics. The + in "Examples per client" client denotes standard deviation.
EXPERIMENTS
We empirically study FedEP across various benchmarks. We start with toy setting in Sec. where we examine cases where federated posterior average (FedPA, Al-Shedivat et al. 2021), which does not take into account global and other clients' approximations during client inference, performs sub-optimally. We then turn to realistic federated learning benchmarks in Sec. 3.2 where both the
size of the model and the number of clients are much larger than had been previously considered in
prior EP-based approaches to federated learning (Corinzia et al. 2019) Kassab & Simeone) 2022).
Here, we resort to the techniques discussed in Sec. approximate inference of the tilted distribu- tions, adaptive optimization, and possibly stateless clients. Finally, we conclude in Sec. 3.3 with an analysis of some of the observations from the benchmark experiments.
3.1 TOY EXPERIMENTS
We start with simple toy setting illustrate the differences between FedPA and FedEP. Here the task is to infer the global mean from two clients, each of which is parameterized as two- dimensional Gaussian, Pk (0) N(;; pk, I₃) ) for E {1,2]. Assuming an improper uniform prior, the global distribution is then also a Gaussian with its posterior mode coinciding with the global mean. We perform exact inference via analytically solving DKL(9%) ())) but restrict the variational family Gaussians with diagonal covariance (i.e., mean-field family). this case both the FedAvg and FedPA solution can be derived in "one-shot". Fig. illlustrates simple case where posterior averaging performs sub-optimally. On the other hand, expectation propagation iteratively refines the approximations toward the globally optimal estimation.
We study this phenomena more systematically by sampling random client distributions, where the client parameters are sampled from the normal-invers Wishart (NIW) distribution.
Client
Global (Target)
1 1xxx Mx ~ (N Po.
Σg W- (I V,)).
Federated Averaging
Here we set the hyper-prior mean P4D = O, degrees of freedom v 7, scale À = 0.2, and sample a random symmetric posi- tive definite matrix for Y. Table[]shov shows the average Euclidean distances between the estimated and target global mean for FedAvg, FedPA, and FedEP averaged over 200 random sam- ples of client distributions. Experimental results demonstrate that iterative message passing in FedEP consistently improves upon the sub-optimal solution from posterior averaging.
Expection Propagation (Trajectory)
Posterior
Averaging
Client
Figure 1: FedAvg, FedPA, and FedEP on a toy two dimensional dataset with two clients.
3.2 BENCHMARKS EXPERIMENTS
We next conduct experiments
suite of realistic benchmark tasks introduced by Reddi et al. (2020). Table I summarizes the model and raw dataset statistics, which is the same as in Al-Shedivat al. 2221). We use the dataset preprocessing provided in TensorFlow Federated (TFF Authors and implement the models in Jax (Bradbury et al., 2018: Hennigan et al. 2020 Ro et al.,
2018
2021 We compare against both FedAvg with adaptive optimizers and FedPA As in FedPA,
run a few rounds of FedAvg as burn-in before switching to FedEP. refer the reader to the appendix for the exact experimental setup.
For evaluation we consider both convergence speed and final performance. On CIFAR-100 and EMNIST-62, we measure the (1) number of rounds reach certain accuracy thresholds (based on 10-round running averages), and (2) the best accuracy attained within specific rounds (based on 100-round running averages). For StackOverflow, we measure the best precision, recall, miero- and
Reddi al. (2020) refer federated averaging with adaptive server optimizers FedAdam etc. We refer this FedAvg for simplicity.

---

Published as a conference paper at ICLR 2023
Accurucy
0%, t) 1500R
Rounds (A, L)
Method
prec.
recall
mi-F1
ma-F1
Method
L0COR
45%
50%
FedPA
74.66
19.94
30.78
11.63
FedPA FedAvg
45.8
48.4
811
FodAvg
23.32
8.02
44.7/0.)
(6.2
(1]
FodSEP (1)
71.32
2z)
(7.04(225)
13.61
FedEP ())
48.7
50.7 (0.4)
(73
1167 (177)
FedSEP (M)
58.31 (1004)
(.70
14.29 (2.70)
2.70
FedEP (M)
(63
1240 (133)
FedSEP (L)
70.98 (818))
25.88 (0:00)
37.80 (129)
13.97
13.97
FodEP
46.5au.)
47.7 (1:)
(23
FedSEP (V)
69.51 (035)
28.02
39.78 (225)
15.32
FedEP (V)
47.8
49.6gaus)
487 (24)
1290
Table 3: StackOverflow Experiments. Statistics shown are the averages and standard deviations (subscript in brackets) aggregated over 5 seeds. We measure the best precision (prec.). recall, micro- and macro-FI (mi/ma-F1) attained by round 1500 (based on 100- round running averages).
FedSEP (1)
48.2 18.2
48.9 (0.4)
438
FedSEP (M)
48.2
48. (0.4
438
FodSEP (L)
4..2
47 (0.8
442
FedSEP (V)
47.8(u.)
(8.5
(40
Table 2: CIFAR-100 Experiments. Statistics shown the averages and standard deviations (subscript in brackets) aggregated over 5 seeds. We measure the number of rounds to reach certain accuracy thresh- olds (based on 10-round running averages) and the best accuracy attained within specific rounds (based on 100-round running averages). FedEP and FedSEP re- fer to the stateful EP and stateless stochastic EP. We use I (Scaled Identity Covariance). M (MCMC), L (Laplace), and V (NGVT) to refer to different inference erages and standard deviations of Euclidean distances techniques. 'One seed does not reach the threshold. between the estimated and target global mean aggre- #Only one seed reaches the threshold.
Method
Euclidean Distance
FedAvg FedPA
5.4 x 10 1.77 x 10-1 2.6 x 10 2.6 10-1
FedEP
1.1 x 10 +988 x 10-8
Table 4: Toy Experiments. Statistics shown are the av-
gated over 200 random samples of client distributions.
macro-F1 I attained by round 1500 (based on 100-round running averages) 6 Due to the size of this dataset, the performance at each round is evaluated on 10K subsample. The evaluation setup is al- most exactly the same in prior work (Reddi et al.) (1-Sheeitte et al. (021). Due to space we mainly discuss the CIFAR- 100 ("CIFAR") and StackOverflow Tag Prediction ("StackOverfloww) re- sults in this section and defer the EMNIST-62 ("EMNIST") results (which are qualitatively similar) to the appendix (Sec. A.3)
CIFAR. In Table 2 and Fig. 2 (left, mid), we compare FedAvg, FedPA, and FedEP with various approaches for approximating the clients' tilted distributions (Sec. 2.2.2 A notable observation is the switch from FedAvg to FedPA/FedEP at the 400th round, where observe significant increases in performance. Somewhat surprisingly, we find that scaled identity is a simple yet strong base- line. (We conduct further experiments in Sec. 3.3 to analyze this phenomena in greater detail). We next experiment with stochastic EP (FedSEP, Sec. 2.2.4). stateless version of FedEP that is more memory-efficient. We find that FedSEP can almost match the performance of full EP despite being much simpler (Fig. 2 right).
StackOverflow. Experiments on CIFAR study the challenges when scaling FedEP to richly pa- rameterized neural models with millions of parameters. Our StackOverflow experiments are on the other hand intended to investigate whether FedEP can scale to regimes with a large number of clients (hundreds of thousands). Under this setup the number of clients is large enough that the average client will likely only ever participate in single update round, which renders the stateful version of FedEP meaningless. We thus mainly experiment with the stateless version of FedEP"" and
Table and Fig. (full figure available in the appendix Fig. [5) show the results comparing the same set of approximate client inference techniques. These experiments demonstrate the scalability of EP to large number of clients even when we assume clients are stateless.
3.3 ANALYSIS AND DISCUSSION
The Effeectivnness of Scaled Identity. Why does the scaled identity approximation work so well? We investigate this question in the same toy setting in Sec. 3.1. Fig. (left) compares the scaled- identity EP with FedEP, FedPA, and FedAvg. Unsurprisingly, this restriction leads to worse perfor- mance initially. However, as clients pass messages between each other, scaled-identity EP eventually converges to nearly the same approximation as diagonal EP.
The toy experiments demonstrate the effectiveness of scaled identity in terms of the final solution. However, this does not fully explain the benchmark experiments where we observed scaled iden-
6TFF by default considers a threshold-based precision and top-5 recall. Our early experiments found that threshold-basee metrics correlate better with loss, and use them in StackOverflow experiments. This was also due to the practical difficulty of storing all the elients' distributions.
7

---

Published as a conference paper at ICLR 2023
34
0.50
0.2
reding FalP DO MCCC TadE IMOVE
Fasleg
0.2
Futkay
28-
Foc pie NIGOOMO
26
TodP IVOVE
0.4
Me Ina
644
0--2-
Loss
Accuracy
Accuracy
Figure 2: 100 Experiments. Left and Middle: loss and accuracy of the server as function of rounds for FedAvg, FedPA, and (stateful) FedEP with various inference techniques. Right: accuracy as a function of rounds for FedAvg. FedPA, and (stateless) FedSEP. The transitions from FedAvg to FedPA, FedEP, and FedSEP happen at round 400. Lines and shaded regions refer to the averages and 2 standard deviations over 5 runs, resp.
8.54
FudF's
2,200
Tecking
Feating
3:755
Fisthag
Secilla (iG-aiacC
Foot iSc MCCV
TesSeP MOMOC
TucSEr IVG2
Fedt INCV
3:90
FazSEP Laptied
125
2:22-
Wit
2130
Micro-F1
Macro-F1
8375
Loss
135
299 66n
3000
3000
3400
031
69n 1006
Like
200 d2c n6C 350 1859 1201 1d89
Figure 3: StackOverflow Experiments. Curves represent loss, micro-F1, and macro-F1 of the global parameter estimation as a function of rounds for FedAvg, FedPA, and (statcless) FedSEP with various inference tech- niques. The transitions from FedAvg to FedPA and FedSEP happen round 800. Lines and shaded regions refer to the averages and 2 standard deviations over runs, resp.
tity EP to match more involved variants in terms of convergence speed. We hypothesize that as models grow more complex, the gap between scaled identity and other techniques might decrease due to the difficulty of obtaining credible estimates of (even diagonal) covariance in high dimen- sional settings. To test this, we revisit the CIFAR-I00 task and compare the following two settings: "small" setting which uses a smaller linear model on the PCA'ed features and has 10.1K parame- ters, and a "large" setting that uses linear model on the raw features and has 172.9K parameters. For each setting, we conduct experiments with EP using scaled-identity and NGVI and plot the results in Fig. (right). We observe that under the "small" setting, a more advanced approximate inference technique converges faster than scaled-identity EP, consistent with the toy experiments. As we increase the model size however ("large" setting), the gap between these two approaches disappears. This indicates that the model gets more complex, the convergence benefits of more advanced approximate inference decline due to covariance estimation's 's becoming more difficult.
Uncertainty Quantification. One motivation for Bayesian approach is uncertainty quantification. We thus explore whether a Bayesian treatment of federated learning results in models that have bet- ter expected calibration error (ECF, Naeinn et al.]
Accuracy 1%. Point Est. Marg.
ECE-15 4)
Method
Point Est. Marg.
FedPA
48.1
13.6 19.5
FodAvg
46.6
FedEPP (1)
50.8
49.6)..)
4.9
7.9
2015 Guo et al.] 2017). which is defined as ECE = Σ, Nais b; accuracy, - confidence;] Here accuracy, is the top-1 prediction accuracy in i-th bin, confidence; is the average confidence of pre- dictions in i-th bin, and b, is the fraction of data points in i-th bin. Bins are constructed in a uniform way in the [0, 1] range We consider accuracy and calibration from the resulting approximate posterior in two ways: (1) point estimation, which uses the fi- nal model (i.e., MAP estimate from the approximate posterior) to obtain the output probabilities for each data point, and (2) marginalized estimation, which samples 10 models from the approximate pos- terior and averages the output probabilities to obtain the final prediction probability. In Table 5
FedEP (M)
50.5
50.2
5.9
4.6ma4)
FedEP (L.
47.7 (05)
47.8 47.8
8.8
6.6744)
FedEP (V)
49.7
49.5guss
5.9
2.2
FodSEP (D)
49.0
48.5
10.0
3/413)
FedSEP (M)
48.9/04)
48.6y04)
10.1 10.1
3.5
FedSEP (L)
47.7
47.8 47.8
9.6
7:2
FedSEP (V)
48.5
48.7
9.3,a.)
3.7
Table 5: CIFAR-100 Calibration Experiments.
FedEP and FedSEP refer to the stateful EP and stateless stochastic EP. We use (Scaled Iden- tity Covariance), M (MCMC), I. (Laplace), and V (NGVT) to refer to different inference techniques.
We also experimented with an alternative binning method which puts an equal number of data points in each hin and observed qualitatively similar results.
8

---

Published as a conference paper at ICLR 2023
Accuracy (small)
Accuracy (large)
Euclidean Distance
FecEh Glegena
Fextii tinke Sieeey
10a
10'
143
Figure 4: Analysis Experiments. Left:: the average Euclidean distances between the estimated and target global mean function of rounds in the toy setting. Middle and Right: accuracy as function of rounds in the CIFAR-100 setting, with either (relatively) small model (middle) or large model (right).
we observe that FedEP/FedSEP improves both the accuracy (higher better) as well as expected calibration error (lower is better), with marginalization sometimes helping.
Hyperparameters. Table 8 shows the robustness of FedEP w.r.t. various hyperparameters.
Limitations. While FedEP outperforms strong baselines in terms of convergence speed and final accuracy, it has several limitations. The stateful variant requires clients to maintain its current con- tribution the global posterior, which increases the clients' memory requirements. The non-scaled- identity approaches also impose additional communication overhead due to the need to communi- cate the diagonal covariance vector. Further, while SG-MCMC/Scaled- Identity approaches have the same compute cost FedAvg the client side, Laplace/NGVi approaches require more compute to estimate the Fisher term. Finally, from a theoretical perspective, while the convergence properties of FedAvg under various assumptions have been extensively studied (Li et al. 2018 2020b), such guarantees for expectation propagation- approaches remains an open problem.
4 RELATED WORK
Federated Learning, FL is a paradigm for collaborative learning with decentralized private data (Konečný et al. 2016; McMahan et al. 2017 Li et 2020a Kairouz al., 2021 Wang al., 2021 Standard approach to FL tackles it as a distributed optimization problem where the global objective is defined by weighted combination of clients' local objectives (Mohri et al., 2019 Li et al. 2020a Reddi et al. 2020 Wang et al. 2020b). Theoretical analysis has demonstrated that federated optimization exhibits convergence guarantees but only under certain conditions, such as bounded number of local epochs (Li et al. 2020b). Other work has tried to improve the averaging- based aggregations Yurochkin et al. (2019); Wang et al. (2020a). Techniques such as secure ag- gregation (Bonawitz et al. [2017] 2019 He et al. 2020) and differential privacy (Sun et al.] 2019 McMahan et al., 2018) have been widely adopted to further improve privacy in FL (Fredrikson et al, 2013). Our proposed method is compatible with secure aggregation because it conducts server-side side reductions over Ank, AAx-
Expectation Propagation and Approximate Inference. This work considers EP general tech- nique for passing messages between clients and servers on partitioned data. Here, the cavity distri- bution "summarizes" the effect of inferences from all other partitions and can be used as a prior in the client's 's local inference. Historically. EP usually refers to specific choice of divergence func- tion ( (pllg) (Minka 2001). This is also known as Variational Message Passing (VMP. Winn et al. 2005) when ( is used instead, and Laplace propagation (L.P, (Soola et al. 2003) when Laplace approximation is used. There have been works that formulate federated learning as probabilistic inference problem. Most notably, AAl-hediiva et al. (2021) 1) formulate FL posterior inference problem. Achituve et al. (202T) apply Gaussian processes with deep kernel learning (Wil- et al. 2016) personalized FL. Finally, some prior works also consider applying EP to federated learning (Corinzia et al. 2019 Kassab & Simeone, 2022 Ashman et al.) 2022), but mostly on rela- tively small-scale tasks. In this work, instead discuss and empirically study various algorithmic considerations to scale up expectation propagation to contemporary benchmarks.
CONCLUSION
This work introduces a probabilistic message-passing algorithm for federated learning based on expectation propagation (FedEP). Messages (probability distributions) are passed to and from clients to iteratively refine global approximations. To scale up classic expectation propagation to the modern FL setting, discuss and empirically study various algorithmic considerations, such as choice of variational family, approximate inference techniques, adaptive optimization, and stateful/stateless clients. These enable practical EP algorithms for modern-scale federated learning models and data.

---

Published as a conference paper at ICLR 2023
Reproducibility Statement. For experiment details such the dataset, model, and hyperparame- ters, we provide detailed descriptions in Sec. 3|as as well as Sec. A.4 We also include in the Appendix additional derivations related to adaptive optimization and damping (Sec. A.I.
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their comments, and are grateful to Maruan Al-Shedivat for his feedback. EX supported by NSF IIS1563887, NSF CCF1629559, NSF IS1611583, NGA RM047622100022 NIGMS R01GM140467, NSF IIS195553,, NSF CNS2008248, NSF IIS2123952, and NSF BCS2040381. YK acknowledges the support of MIT-IBM Watson Al and Amazon.
REFERENCES
Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems, 34:8392- 8406, 2021.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. In ICLR, 2021.
Matthew Ashman, Thang D Bui, Cuong Nguyen, Efstratios Markou, Adrian Weller, Siddharth Swaroop, and Richard E Turner Partitioned variational inference: A framework for probabilistic federated learning. arXiv preprint arXiv:2202. 12275, 2022.
The TensorFlow Federated Authors. TensorFlow Federated, 12 2018. URLhttts://gittuubcco Lensorrflow/eedrrtte
Charles Blundell, Julien Corebiie,, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for privacy- preserving machine learning. In proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security, pp. 1191- 2017.
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp. Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, et al. Towards federated learning scale: System design. Proceedings of Machine Learning and Systems, 1: 374-388, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Men Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github.com cow.b/////mooogeeeae
Tamara Broderick, Nicholas Boyd, Andre Wibisono, Ashia C. Wilson, and Michael Jordan. Streaming variational bayes. In Christopher J. C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems, pp. 1727- 1735, 2013.
Zachary Charles and Jakub Konečný, On the outsized importance of learning rates in local update methods. arXiv preprint arXiv:2007.00878. 2020.
Luca Corinzia, Ami Beuret, and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint arXiv: / 01006996666... 2019.
Guillaume Dehaene and Simon Barthelmé. Expectation propagation in the large data limit. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 80(1):199-211 2018.
Guillaume Dehaene and Simon Barthelmé. Bounding errors of in Neural Information Processing Systems, 28, 2015.
Advances
10

---

Published as a conference paper at ICLR 2023
Matt Fredrikson, Somesh Jha, and Thomas Ristenpart. Model inversion attacks that exploit confi- dence information and basic countermeasures. In Proceedings of the 22nd ACM SIGSAC confer- computer and communications security, pp. 1322-13333 2015.
Alex Graves. Practical variational inference for neural networks. Advances neural information processing systems, 24, 2011.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, pp. 1321-1330. PMLR, 2017.
Chaoyang He, Songze Li, Jinhyun So, Xiao Zeng, Mi Zhang, Hongyi Wang, Xiaoyang Wang, Pra- neeth Vepakomma, Abhishek Singh, Hang Qiu, et al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint urXiv:2007.35518 2020.
Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX, 2020.
URL http://gi thub www.://c//d/eeemd/////ee/ heieepmiindddmhhaaikee
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, al. Ad- vances and open problems in federated learning. Foundations and Trends Machine Learning, 14(1-2):1-210, 2021.
Rahif Kassab and Osvaldo Simeone, Federated generalized bayesian learning via distributed stein variational gradient descent. IEEE Trunsactions Signal Processing, 70:21 10:2180-2192, 2022. doi: 10.1109/TSP.2222.316663
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by weight-perturbation adam. In International Conference on Machine Learning, pp. 2611 1-2620. PMLR, 2.18.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In ICLR, 2015.
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint arXiv:1610.02527, 2016.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical Fisher approx- imation for natural gradient descent. Advances in Neural Information Processing Systems, 32, 2019.
Olivier Ledoit and Michael Wolf. A well-conditioned estimator for large-dimensional covariance matrices. Journal of Multivariate Analysis, 88(2):365-411, 2004.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization heterogeneous networks. arXiv preprint 1rXiv:18/2.06112 2018.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-60, 2020a.
Xiang Li. Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of fedavg non-iid data. In International Conference on Learning Representations, 2020b.
Yingzhen Li, José Miguel Hermández-Lobato, and Richard E Turner. Stochastic expectation propa- gation. NeurIPS, 2015.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient MCMC. Ad- vances in Neural Information Processing Systems, 28, 2015.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple bascline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32, 2019.
Stephan Mandt, Matthew Hoffman, and David Blei. Stochastic gradient descent as approximate bayesian inference. arXiv preprint 3XX4v:7004.0440 2017.

---

Published as a conference paper at ICLR 2023
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effieee -efiiiinneffii learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282 PMLR, 2017.
H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. In International Conference on Learning Representations, 2018.
Thomas Minka and John Lafferty. Expectation Expectation-propagattooo for the generative aspect model. In Proceedings of the Eighteenth conference Uncertainty in artificial intelligence, pp- 352-359, 2002.
Thomas Minka. Expectation propagation for approximate bayesian inference. In UAI, 2001.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Interna- tional Conference on Machine Learning, pp. 4615-4625, PMLR, 2019.
Mahdi Pakdaman Nacini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated prob- abilities using bayesian binning. In Twenty-Ninth AAAI Conference on Artificial Intelligence, 2015.
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with Bayesian principles. Advances in neural information processing systems, 32, 2019.
Sashank J Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2020.
Jae Hun Ro, Ananda Theertha Suresh, and Ke Wu. FedJAX: Federated learning simulation with
JAX. arXiv preprint arXiv:2108.02111 2022
Alex Smola, S.v.n. Vishwanathan, and Eleazar Eskin. Laplace propagation. In Advances in Neural Information Processing Systems, volume 16, 2003.
Ziteng Sun, Peter Kairouz, Ananda Theertha Suresh, and H Brendan McMahan. Can you really backdoor federated learning? arXiv preprint arXiv:1911.0996. 2019.
Aki Vehtari, Andrew Gelman, Tuomas Sivula, Pasi Jylänki, Dustin Tran, Swupnil Sahai, Paul Blom- stedt, John P Cunningham, David Schiminovich, and Christian P Robert. Expectation propagation as a way of life: A framework for Bayesian inference on partitioned data. JMLR, 2020.
Hongyi Wang, Mikhail Yurochkin, Yuckai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference Learning Represen- tations, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623 2020b.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al- Shehddivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide feder- ated optimization. arXiv preprint arXiv:2107.06911 2021.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-I1), pp. 681-688. Citeseer, 2011.
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep kernel learning. In Artificial intelligence and statistics, pp. 370-378. PMLR, 2016.
John Winn, Christopher M Bishop, and Tommi Jaakkola. Variational message passing. Journal of Machine Learning Research, 6(4), 2005.
Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems, 33:6281-6299 2020.
12

---

Published as a conference paper at ICLR 2023
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. Interna- tional Conference on Machine Learning, pp. 7252-7261 PMLR, 2019.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational in- ference. IEEE Transactions on Pultern Analysis and Machine Intelligence, 41:2008-2026, 2019.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, pp. 5852-5861 PMLR, 2018.
13

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 DAMPED CLIENT AND SERVER UPDATES
To simplity the notations, observe that Eq. 3 could be re-written in the following way,
Pia(O (-)
9
new
where (0) arg min D (x)(0) x()) |l ())(0)
Ry
partially damped client update could be carried out by.
1-5
71x((0)
Client:
9k
new (0)
9x((0)
(-)(0)
8
5
(c(())
4global (0)/qx(( (/)
8. ( T
((<(0)
91: 91((0)
Aglobal (0)
(()
():(0)
Aglobal (0)
91: (0)
(0)
where define (0)
Aglobal (0)
Similarly, (damped) server updates could be written the following,
Server: Aglobal (0)
II 1k (0)
II (O) (AQx(O))*
II (qx (ATR(O))
II 9r (x)
9global (0) II (An)(0))
A.2 EXPECTATION PROPAGATION (EXTENDED)
Expectation propagation (EP) Minka 2001): Vehtari et al. 2020) constructs a posterior approxima- tion through iterating local computations that refine factors that approximate the posterior contribu- tion from each client. In this spirit, we would ideally like to solve the following localized version of Eq. where replace of the factors with its corresponding approximating factor,
new (0) arg min D (k(()p (0) |l (()) p x(0)). where p (0) Pelobal (0) (x
Q=o
Unfortunately, the right-hand side of the divergence is the intractable posterior we would like to approximate in the first place. Instead, EP solves the following problem (Eq.
new (0) = arg min D (0) 9 x0) ( 9 =(0)). qk
where q (0) Iglobal (
9x
A.3 ADDITIONAL. EXPERIMENTS AND DETAILS
StackOverflow. Please Fig. 5 for additional visualizations.
14

---

Published as a conference paper at ICLR 2023
PudFA
Natt
Scalectdeettt)
Ted53h Scaled icandto
E319
FodSEr EG SG-NCN
Fus522 (Sc (SB0
SEP (Si-0
FudEEP (NDV
1110
16-
Micro-F1
Macro-F1
14
922
CO15
E.56
Loss
CARS
1000
1290
1icn
Tualing
hadling Menetit)
31-
Fusher Scaler Mertiyy
hilter Listad
822
FodtEP suplon
FedE1
FedtED (VGV))
$15
35-
Precision
Recall
200
1000
1299
100
Figure S: Extended StackOverflow Visualizations.
0.140
FodEP
0.155
FodEP INS Tw21
TedEP (NGAS.
FedtPP (NDA.
RedEP Sealed itentity)
0.145
0.135
0.125
200
400
800
800
1000
1200
1400
Figure Accuracy function of rounds, and number of NGVI epochs (TNGVI) in the CIFAR-100 setting. with a (relatively) small model.
EMNIST. Please see Fig. 7 and Table 6 for experimental results.
Analysis. This section extends the experiments (the "small" setting) Sec, 3.3 It looks at the performance increase the complexity (a proxy of quality) of approximate inference techniques. We vary the number of iterations in NGVI from 1 (cheap) to 10 (expensive) epochs. We can observe in Fig. that as we increase NGVT's computations, the performance improves.
A.4 HYPERPARAMETERS
Please Table 7 for hyperparameter details. In Table also conduct experiments to understand their influence on the different algorithms.
15

---

Published as a conference paper at ICLR 2023
Fadti
hellieg
Fasheg
whyy
Fad (S0-))MMC)
Pustbon Nuobor (HOAO
FesS2r Euplict
255-
0.645
645
Loss
Accuracy
3.835
303
Figure 7: EMNIST-62 Experiments. Figures show the loss and accuracy of the global parameter estimation function of rounds for FedAvg, FedPA, and (stateless) FedSEP with various inference techniques. The transitions from FedAvg to FedPA and FedSEP happen at round 200.
accuracy (%, f)
rounds (#, L) 86.5%
Method
300R
500R
86%
FedPA
85,9
86.5
246
398
FedAvg
85.3
85.8
465
FedSEP (I)
86.1
86.6
228
399
FedSEP (M)
86.1
86.6
228
399
FedSEP
86.1
86.6
228
364
FedSEP (V)
86.1
86.6
228
382
Table 6: EMNIST-62 Experiments. We measure the number of rounds to reach certain accuracy thresholds (based 10-round running averages) and the best accuracy attained within specific rounds (based 100- round running averages). We (Scaled Identity Covariance), M (MCMC), L (Laplace), and V (NGVT) refer to different inference techniques.
Algorithm Approximate Inference: NGVI
Input: Mi I k, TNGVI, NNOVI PNGVI
2: Initialize sp, 1,,0
for 1,... TNGVI do
4: F+
5: for 1,..., NNGVI do
6: 0~~ (0: E1g--1)
7:
F FU D
Fisher Dz)-
8:
end for
9: F Average(F)
10: S; + BNGVIS:-1 + 11 (1-
-1 11: E (iDels+
12: end for
13: Output: E
16

---

Published as a conference paper at ICLR 2023
Hyperparameter
CIFAR-100
StackOverflow
EMNIST-62
Task Hyperparameters from Al-Shedivat al.
2021
Server Optimizer
SGD (m = 0.9)
Adagrad (T 10-9)
SGD (m 0.9)
Client Optimizer'
SGD (m 0.9)
SGD (m 0.9)
SGD (m 0.9)
Clients Per Round
20
10
100
Server Learning Rate
0.5
5.0
0.5
Client Learning Rate
0.01
50.0
0.01
Client Epochs
10
5
20
Burn In
400
800
200
Client Inference Hyperparameters
Scale Agov* $
5 x 10-2
10-8
5 10-3
MCMC Shrinkage
1 x 10
x 10
1 x 10
Laplace Epochs
5*
5
5
NGVI Epochs
5
10
5
NGVI Samples
NGVI ΩNUVI
0.99
0.99
0.99
Client Inference Hyperparameters Search Space
Scale Reav
(1,2,5,10) x 10-2
1 (10-,,10~$,0 10-9)
{1,5) x
MCMC Shrinkage
x {10 3,10-,,10 1,004 ,1
Laplace Epochs
(5,
NGVI Epochs
(5)
NGVI Samples
[5,10]
NGVI (NGVI
(0.9,0.99)
Table 7: Hyperparameters. *Client two separate optimizers, one used local optimization (SG-MCMC), and one used in local state updates (for stateful FedEP). When applied, the client state optimizer reuses the configuration the server optimizer. This per-data-point scale, and is also used in other approximate inference techniques. *The (stateful) FedEP uses Laplace epochs.
Accuracy (%, f)
Rounds (#, }
Method
Hyperparameter
1000R
1500R
45%
50%
x 10 -
49.1
50.3
457
1081
FedEP
Scale acov
5 x 10 10
48.9
50.5
464
1105
6 10
48.8
50.5
474
1206
10
47.8
48.8
FedEP (M)
MCMC Shrinkage
1 1 10
48.9
50.5
456
1179
5 x 10
49.2
49.2
436
46.7
47.9
513
FedEP (L)
Laplace Epechs
10
46.7
47.9
514
47.9
49.6
478
FedEP (V)
NGVI Epochs
10
46.6
48.5
523
4 x 10
48.2
48.7
431
FedSEP ())
Scale &cov
10 2
48.3
49.0
6 x 10
48.3
49.1
433
5 10
48.3
48,9
431
FedSEP (M)
MCMC Shrinkage
1 10
48.3
49.0
432
5 10
48.4
49.0
431
5
47.2
47.9
437
FedSEP (L)
Laplace Epochs
10
47.2
47.9
439
5
47.9
48.8
432
FedSEP (V)
NGVI Epochs
10
47.8
48.6
Table 8: CIFAR-100 Hyperparameter Analysis Experiments. FedEP and FedSEP refer to the stateful EP and stateless stochastic EP. We use (Scaled Identity Covariance), M (MCMC), L (Laplace), and V (NGVI) to refer different inference techniques.
17


Reviews:
[{"review_id": "zpkUxKGEw3", "reviewer": "Reviewer_jmPj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents an empirical study of different Expectation Propagation algorithmic choices applied to several federated learning settings. Important aspects to an empirical studies are missing.", "summary_of_the_review": "Interesting analysis showing the applicability of EP in FL settings. The R18 and SOF setting is convincing wrt. the aspect of scale. The empirical evaluation is lacking at the moment wrt. seeds and hyperparameter selection. ", "strength_and_weaknesses": "The strengths of this paper lie in the experimental settings and exposition. \nCifar100 with ResNet18 does serve as an example for model-scale (as far as bayesian methods are concerned), although SEP with (I) or maybe also (V) could be applied to even larger models. It would be interesting to see these methods compared against adaptive FedAVG on e.g. a wide resnet 50.\nStackoverflow experiments serve as a good example for a large number of clients with a simple LR model. \nThe exposition is good, explaining EP and its variants in sufficient details. \n\nThe weakness of this paper lie in the level of contribution to ICLR. None of the proposed algorithms is significantly novel as far as I can tell. This paper might be better suited as a workshop contribution. \n\nAnother weakness lies in that for an empirical study especially, I'd have expected more than a single seed for the proposed experiments. While most results seem significantly far apart based on my experience with these methods and data-sets, I would like to see standard-error across several seeds to judge significance. Learning curves show the strengths of EP in principle compared to FedAvg but don't allow to compare different methods and prior work (FedPA).\n\nFor an empirical effort, I would expect to see the search-space for the chosen Hyperparameters in Table 6 and a discussion on how (and if) the optimal hyperparameters chosen here generalize across the different algorithms. E.g. does FedEP(L) require the same hyper-parameters as FedSEP(V). Were these parameters tuned on the non-bayesian baseline? I can imagine that the stochasticity inherent to the bayesian optimization procedures warrants different hyper-parameters. \n\n\n\nSmall things:\nTypo section 3.1: \"Fig. 1 illustrates a simple case where posterior averaging [performs] sub-optimally.\"\nFigure 4: What is the difference between the two FedEP runs in middle and right sub figure?\nTables 1-4, please improve readability, especially by separating captions between Table 2 & 4. I might suggest reducing the font-size across the tables. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "loXWG6CoaM", "reviewer": "Reviewer_Z7LG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a curious algorithm to adapt federated learning in the variational inference framework. In particular, it uses a scalable version of EP to do so. The contribution seems solid although, given its practical application, I miss a real case scenario, for example in hospitals, of the algorithm.  ", "summary_of_the_review": "First thanks for the paper, I have enjoyed its reading. I would like to emphasize here the clarity of this paper and its organization. It is very easy to read although the topic is complex. My main issue that justifies not giving it a better recommendation is that this approach seems to be a practical real-case scenario but the experiments and somehow limited. For instance, your example of hospitals is great, but then you do not apply it to a real case scenario (maybe stackoverflow is the more complicated one). It would be a great paper if a more complex experiment could be done.\n\nConcerning technical details I find some decisions a little bit arbitrary. Why have you used EP and not some version of VI? (Justify this ;-)) Why do not you use EP in the client inference? You can also use it and it would be a full EP approach? Have you considered using PowerEP instead of EP? How can I determine or tune which is the optimal algorithm for the EP client inference? (It would be just great to use Bayesian optimization here to develop an auto-variational-inference approach here). And lastly you need to justify the dumping and whether the non-convergence of EP is a problem. It would also be great to talk about the complexity of your method and some drawbacks about this approach with respect to other variational inference approaches. This are basically all the questions that I have asked myself when reading this.   \n\nSome minor issues or enhancements to the paper are including a paragraph with the organization of the paper in the introduction, introduce a further work paragraph in the conclusions section as I believe that this opens new research lines, put the related work section (IMHO) before, better describe the results obtained by the toy experiment, include in section 1 a reference to other divergence functions instead of KL and in section 2 to variational families and justify why have you used an improper uniform prior (advantages and drawbacks of the decision).", "strength_and_weaknesses": "Strength: Solid methodology adapting EP, clarity.\nWeaknesses: Low empirical evidence. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "YKRcuOdvgNj", "reviewer": "Reviewer_Hh8V", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents FedEP, a federated learning approach based on expectation propagation in which a global inference task is constructed as local inference tasks. The authors also present several ways to scale the model for modern neural networks. The authors compared their method to baselines on CIFAR-100, StackOverflow, and EMNIST-62.", "summary_of_the_review": "I think that a Bayesian approach to FL is a good direction; however, in terms of novelty, and more importantly, the experiments, this paper can greatly be improved. I am willing to reevaluate my review based on the authors' response.", "strength_and_weaknesses": "Strength:\n* A Bayesian approach that in some sense generalizes previous approaches.\n* An important aspect of the work is the emphasis on scaling the model, which I deem as very important for Bayesian models.\n\nWeaknesses/Questions:\n* Novelty. While I appreciate the contribution in terms of scalability I am not sure that it is enough. I do not see how the approach is different than previous methods that also used EP for FL (and specifically, [1], [2]). Can the authors comment on that?\nFurthermore, the methods to scale the model are also common solutions. So overall it seems like the paper combines existing elements from previous studies. I will state though, that in my opinion, the bigger issue is with the experimental section which I will present next.\n* Experiments. I think that there are several issues with this part.\n  * First, information is missing. How did you split the data among clients? How were the hyper-parameters optimized? Was there a validation set? \n  * Second, in the paper you stated that you chose the best accuracy based on a running average of 100 rounds. What does that mean exactly? Why not take the metric results based on the last round or the best round based on a validation set? As far as I know, FL papers do not use this approach.\n  * Third, the baselines in this paper are not enough. I think the method should be compared to additional Bayesian methods (for example, [3]), and to leading methods in FL (and not only FedAvg and FedPA).\n  * In Fig. 2 and Fig. 3, where are the lines of FedEP until 400 and 800 steps respectively? currently, it looks a little bit odd. Also, how can you explain the jump at those points? Did you use a scheduler with a lr drop at those points?\n  * Finally, one of the motivations for being Bayesian is the ability to quantify uncertainty accurately. How is FedEP compared to other methods in that aspect?\n* Minor.\n  * Tables 1-4 could be organized better. Also, their numbering should be according to the reference in the text.\n  * The colors of the lines in the figures are too similar in my opinion.\n\n[1] Bui, T. D., Nguyen, C. V., Swaroop, S., & Turner, R. E. (2018). Partitioned variational inference: A unified framework encompassing federated and continual learning. arXiv preprint arXiv:1811.11206.\n\n[2]  Ashman, M., Bui, T. D., Nguyen, C. V., Markou, E., Weller, A., Swaroop, S., & Turner, R. E. (2022). Partitioned Variational Inference: A framework for probabilistic federated learning. arXiv preprint arXiv:2202.12275.\n\n[3] Achituve, I., Shamsian, A., Navon, A., Chechik, G., & Fetaya, E. (2021). Personalized Federated Learning with Gaussian Processes. Advances in Neural Information Processing Systems, 34, 8392-8406.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "K0zL6DxwXT", "reviewer": "Reviewer_za1u", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes to use EP together with approximate inference methods for federated learning. The approach shows promising results on toy examples and real data.  The method works with any approximate inference methods, and different choices are evaluated and compared. ", "summary_of_the_review": "The proposed FedEP algorithm is shown to perform well in practice on toy examples and real data.  The comparison of the different approximate inference methods is interesting, and highlights how the method can use any approximate inference scheme.  I believe this work will be interesting to the bayesian deep learning community as well as people working in federated learning. I am therefore leaning towards acceptance of this work, despite the outlined weaknesses.", "strength_and_weaknesses": "Strengths:\n- I found the toy experiment in Section 3.1 to be very illuminating, along with Table 4. FedEP greatly improves over FedAvg and FedPA.\n- The experiments on real data are also convincing, and the comparison of the difference approximate inference schemes was interesting to see. \n- The overall algorithm is agnostic to the approximate inference method, so advances in approximate inference may further improve FedEP.\n\nWeaknesses:\n- Some parts of the paper were not clear / easy to read (see next section).\n- The algorithm seems not too novel, since EP and the used approximate inference schemes are well-known in literature. \n- It was a bit disappointing to see that a very \"simple\" approximate inference method (scaled identity) works best. This somehow suggests that the good performance may not be due to the EP formalism, but some other effects.  The paper would be much stronger (and I would increase my rating in that case), if it could be shown that more accurate approximate inference would lead to more accurate results. For example the accuracy of the approximate inference methods could be scaled up on small problems by doing a gradually more expensive MCMC estimation of the posterior, or my increasing the number of samples in the NGVI algorithm. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster