Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms
Authors: Maruan Al-Shedivat
Jennifer Gillenwater
Eric Xing
Afshin Rostamizadeh
Venue: ICLR 2021 Poster (2021)
Keywords: federated learning
posterior inference
MCMC

Abstract: Federated learning is typically approached as an optimization problem, where the goal is to minimize a global loss function by distributing computation across client devices that possess local data and specify different parts of the global objective.  We present an alternative perspective and formulate federated learning as a posterior inference problem, where the goal is to infer a global posterior distribution by having client devices each infer the posterior of their local data.  While exact inference is often intractable, this perspective provides a principled way to search for global optima in federated settings.  Further, starting with the analysis of federated quadratic objectives, we develop a computation- and communication-efficient approximate posterior inference algorithm—federated posterior averaging (FedPA).  Our algorithm uses MCMC for approximate inference of local posteriors on the clients and efficiently communicates their statistics to the server, where the latter uses them to refine a global estimate of the posterior mode.  Finally, we show that FedPA generalizes federated averaging (FedAvg), can similarly benefit from adaptive optimizers, and yields state-of-the-art results on four realistic and challenging benchmarks, converging faster, to better optima.

Published as a conference paper at ICLR 2021
FEDERATED LEARNING VIA POSTERIOR AVERAGING: A NEW PERSPECTIVE AND PRACTICAL ALGORIHTMS
Maruan Al-Shedivat CMU
Jennifer Gillenwater Google
Eric Xing MBZUAI & CMU
Afshin Rostamizadeh Google
ABSTRACT
Federated learning is typically approached optimization problem, where the goal is to minimize global loss function by distributing computation across client devices that possess local data and specify different parts of the global objective. We present an alternative perspective and formulate federated learning as a posterior inference problem, where the goal is infer a global posterior distribution by having client devices each infer the posterior of their local data. While exact inference is often intractable, this perspective provides a principled way to search for global optima in federated settings. Further, starting with the analysis of federated quadratic objectives, we develop a computation- and communication- efficient approximate posterior inference federatm-feeerrted posterior averaging (FEDPA). Our algorithm uses MCMC for approximate inference local posteriors the clients and efficiently communicates their statistics to the server, where the latter uses them to refine a global estimate of the posterior mode. Finally, we show that FEDPA generalizes federated averaging (FEDAvG), can similarly benefit from adaptive optimizers, and yields state-of-ther of-the-art results on four realistic and challenging benchmarks, converging faster, to better optima.
INTRODUCTION
Federated learning (FL) is framework for learning statistical models from heterogeneous data scattered across multiple entities (or clients) under the coordination of a central server that has no direct access to the local data (Kairouz et al., 2019). To learn models without any data transfer, clients must process their own data locally and only infrequently communicate some model updates the server which aggregates these updates into global model (McMahan al., 2017). While this paradigm enables efficient distributed learning from data stored on millions of remote devices (Hard al., 2018), it comes with many challenges (Li et al., 2020), with the communication cost often being the critical bottleneck and the heterogeneity of client data affecting convergence.
Canonically, FL is formulated distributed optimization problem with few distinetive properties such as unbalanced and non-i.i.d data distribution across the clients and limited communication. The de facto standard algorithm for solving federated optimization federated averaging (FEDAVG, McMahan et al., 2017) 7), which proceeds in rounds of communication between the server and random subset of clients, synchronously updating the server model after each round (Bonawitz et al., 2019), 9). By allowing the clients perform multiple local SGD steps (or epochs) at each round, FEDAVG can reduce the required communication by orders of magnitude compared to mini-batch (MB) SGD.
However, due to heterogeneity of the client data, more local computation often leads to biased client updates and makes FEDAVG stagnate at inferior optima. As a result, while slow during initial training, MB-SGD ends up dominating FEDAVG at convergence (see example in Fig. 1). This has been observed in multiple empirical studies (e.g., Charles & Konečný, 2020), and recently was shown theoretically (Woodworth al., 2020a). Using stateful clients (Karimireddy et al., 2019; Pathak & Wainwright, 2020) can help to remedy the convergence issues in the cross-silo setting, where relatively few clients are queried repeatedly, but is not practical in the cross-device setting (i.e., when clients are mobile devices) for several reasons (Kairouz et al., 2019; Li et al., 2020; Lim et al., 2020). One key issue is that the number of clients in such a setting is extremely large and the average client will only ever participate in single FL round. Thus, the state of a stateful algorithm is never used.
Most of the work done Google. Correspondence: maruan. narnan.alshedivat.com

---

Published as a conference paper at ICLR 2021
Local objectives
FedAAg Convergence
FedPA Convergence
40
client
local optima
Distance to G.O.
10a
MB-SGD
obj eetive
FedAvg 10 sceps
FedAvg 100 steps
10"
FedPA -00 samples
Client #2 objective
FedPA 100 samples
-20 global optimum
10-1
30
01
Round Number
Round Number
An illustration of federated learning in toy 2D setting with two clients and quadratic objectives. Left: Contour plots of the client objectives, their local optima, as well as the corresponding global optimum. Middle: Learning curves for MB-SGD and FEDAVG with 10 and 100 steps per round. FEDAVG makes fast progress initially, but converges to point far away from the global optimum. Right: Learning curves for FEDPA with 10 and 100 posterior samples per round and shrinkage p 1. More posterior samples (i.c, more local computation) results in faster convergence and allows FEDPA to come closer to the global optimum. Shaded regions denote bootstrapped 95% based on runs with different initializations and random seeds. Best viewed in color.
Figure 1:
Is it possible design FL algorithms that exhibit both fast training and consistent convergence with stateless clients? In this work, we answer this question affirmatively, by approaching federated learning not as optimization but rather as posterior inference problem. We show that modes of the global posterior over the model parameters correspond to the desired optima of the federated optimization objective and can be inferred by aggregating information about local posteriors. Starting with an analysis of federated quadratics, we introduce a general class of federated posterior inference algorithms that run local posterior inference on the clients and global posterior inference on the server. In contrast with federated optimization, posterior inference can, with stateless clients, benefit from an increased amount of local computation without stagnating at inferior optima (illustrated in Fig. 1). However, a naïve approach to federated posterior inference is practically infcasible because its computation and communication costs are cubic and quadratic in the model parameters, respectively. Apart from the new perspective, our key technical contribution is the design of efficient algorithm with linear computation and communication costs.
Contributions. The main contributions of this paper can be summarized follows:
1. We introduce a new perspective on federated learning through the lens of posterior inference which broadens the design space for FL. algorithms beyond purely optimization techniques.
2. With this perspective, we design computation- and communication--effiieen approximate posterior inference
posterior averaging (FEDPA). FEDPA works with stateless clients and its computational complexity and memory footprint are similar FEDAVG.
Soorithmm-edereen
3. We show that FEDAVG with many local steps in fact a special case of FEDPA that estimates local posterior covariances with identities. These biased estimates are the source of inconsistent updates and explain why FEDAVG has suboptimal convergence even in simple quadratic settings.
4. Finally, we compare FEDPA with strong baselines on realistic FL benchmarks introduced by Reddi et al. (2020) and achieve state-of tatee-of-hh--nn art results with respect multiple metrics of interest.
RELATED WORK
Federated optimization. Starting with the seminal paper by McMahan al. (2017), lot of recent effort in federated learning has focused on understanding of FEDAVG (also known as local SGD) as an optimization algorithm. Multiple works have provided upper bounds on the convergence rate of FEDAVG in the homogencous i.i.d. setting (Yu et al., 2019; Karimireddy et al., 2019; Woodworth et al., 2020b) as well as explored various non-i.i.d. settings with different notions of heterogeneity (Zhao et al., 2018; Sahu et al., 2018: 8; Hsieh et al., 2019; Li et al., 2019; Wang et al., 2020; Woodworth et al., 2020a). Reddi et al. (2020) reformulated FEDAVG in a way that enabled adaptive optimization and derived corresponding convergence rates, noting that FEDAVG requires careful tuning of learning rate schedules in order to converge to the desired optimum, which was further analyzed by Charles Konečný (2020). To the best of our knowledge, our work is perhaps the first to connect, reinterpret, and analyze federated optimization from the probabilistic inference perspective.
Distributed MCMC. Part of our work builds on the idea of sub- posterior aggregation, which was originally proposed for scaling up Markov chain Monte Carlo techniques to large datasets (known as the concensus Monte Carlo, Neiswanger et al., 2013; Scott et al., 2016). One of the goals of this paper is to highlight the connection between distributed inference and federated optimization and develop inference techniques that can be used under FL-specific constraints.
2

---

Published as a conference paper at ICLR 2021
3 A POSTERIOR INFERENCE PERSPECTIVE ON FEDERATED LEARNING
Federated learning is typically formulated as the following optimization problem: 1 Eario)) fi(0) = The
F(O) [00)= E
min OER
5(0; zig),
(I)
where the global objective function F(O) is weighted average of the local objectives fi) over N clients; each client's objective is some loss f(0;z) 2) computed on the local data D; {=E1..... Fin; J In cross-device applications, the total number of clients N can be extremely large, and hence optimization of F(0) is done over multiple rounds with only a small subset of M clients participating in each round. The weights {%: are typically set proportional to the sizes of the local datasets (nt). which makes F(0) coincide with the training objective of the centralized setting.
Typically, =0; is negative log likelihood of z under some probabilistic model parametrized by 0, i.e. f(0;z) 2) ;= log IP (z 0). For example, least squares loss corresponds to likelihood under a Gaussian model, cross entropy loss corresponds likelihood under categorical model, etc. (Murphy, 2012). Thus, Eq. corresponds to maximum likelihood estimation (MLE) of the model parameters O.
An alternative (Bayesian) approach to maximum likelihood estimation is posterior inference or esti- mation of the posterior distribution of the parameters given all the data: IP (e D = UNU The posterior is proportional to the product of the likelihood and a prior, IP (e D) P(())))((
and, if the prior is uninformative (uniform over all 9), the modes of the global posterior coincide with MLE solutions or optima of F(0) in Eq. 1. While this simple observation establishes an equivalence between the inference of the posterior mode and optimization, the ad advantage of this perspective comes from the fact that the global posterior exactly decomposes into a product of local posteriors.'
Proposition 1 (Global Posterior Decomposition) Under the uniform prior, any global posterior distribution that exists decomposes into a product of local posteriors: P (0 D) x
N II P D.).
II
Proposition suggests that as long as we are able to compute local posterior distributions TP (0 D.) and communicate them to the server, we should be able to solve Eq. by multiplicatively aggregating them to find the mode of the global posterior IP (0 D) on the server. Note that posterior inference via multiplicative averaging has been successfully used to scale Monte Carlo methods large datasets, where the approach is embarrassingly parallel (Neiswanger et al., 2013; Scott et al., 2016). In the FL context, this means that once all clients have sent their local posteriors to the server, we can construct the global posterior without any additional communication. However, there remains the challenge of making the local and global inference and communication efficient enough for real federated settings. The example below illustrates how this can be difficult even for simple model and loss function.
Federated least squares. Consider federated least squares regression with a linear model, where 2 := (x, y) and the loss f(0;x, x, y) := [x 0 - y)² is quadratic. Then, the client objective becomes:
0 11 = log exp { 1 2((0 - m.: (0- - (1)) + const, 2
fi) = log exp
(2)
where X; E x is the design matrix, y; E IR the response vector, Σ, X, and P4 X x,) -xx,)>xx x,y,, y, Note that the expression in Eq. 2 is the log likelihood for multivariate Gaussian distribution with mean μı; and covariance Σ, Therefore, each local posterior (under the uniform prior) is Gaussian, and, as a product of Gaussians, the global posterior is also Gaussian with the following mean (which coincides with the posterior mode);
N
N
Σ
E. 11: Σ
(3)
Concretely, in the case of least squares regression, this suggests that it is sufficient for clients to infer the means (14)) and inverse covariances (Σ,") " of their local posteriors and communicate that information to server for the latter to be able to find the global optimum. However, a straightforward application of Eq. 3 would require (d2)) space and (dd)) computation, both on the clients and on the server, which is very expensive for the typical cross-device FL setting. Similarly, the communication cost would be (dd)) while standard FL algorithms have communication cost of (d).
'Not that from the optimization point of view, the global optimum generally cannot be represented any weighted combination of the local optima even in simple 2D settings (see Fig. 1, left).
3

---

Published as a conference paper at ICLR 2021
Algorithm 1 Generalized Federated Optimization
Algorithm 2 Client Update (FEDAVG)
input initial 0, CLIENTUPDATE, SERVERUPDATE
input initial On, loss fz f:) optimizer CLIENTOPT 1: for k L..... K do
1: for each round 1 = 1....,' do
2: Sample subset S of clients
2: 0x <- CLIENTOPT(OA-LL (0x(0=)))
all 4: for each client i E in parallel do
3: communicate
// server clients
3: end for
output A := 0p O2. client weight
5: A;,9: + CLIENTUPDATE(
Algorithm Client Update (FEDPA)
6: end for
7: communicate {Ai, ies
// server clients
input initial Ob. loss f(( sampler CLIENTMCMC 1: for k =,,... K do
A' 4 TST LiES A:A; // /aggregate updates
9: 0 + SERVERUPDATE((. A
2: 0₂ CLIENTCCMC(AA
10: end for
3: end for
output final 0
output A := I (0u jt), client weight qi
Approximate federated posterior inference. Apart from the computation and communication issues discussed in the simple example above, we also have to contend with the fact that, generally, posteriors are Gan-Gaussian and closed form expressions for global posterior modes may not exist." In such cases, propose to use the Laplace approximation for local and global posteriors, i.e., approximate them with the best-fitting Gaussians. While imperfect, this approximation will allow us to compute the (approximate) global posterior mode in computation- and communicationn-ffiiennn manner using the following three steps: (i) infer approximate local means {it,) and covariances {E,}. (ii) communicate these to the server, and (iii) compute the posterior mode given by Eq. 3. Note that directly computing and communicating these quantities would completely infeasible for the realistic setting where models are neural networks with millions of parameters. In the following section, we design a practical algorithm where all costs are linear in the number model parameters.
FEDERATED POSTERIOR AVERAGING: A PRACTICAL ALGORITHM
Federated averaging (FEDAvG, McMahan et al., 2017) solves the problem from Eq. I over T rounds by interacting with M random clients at each round the following way: (i) broadcasting the current model parameters 0 to the clients, (ii) running SGD for K steps on each client, and (iii) updating the global model parameters by collecting and averaging the final SGD iterates. Reddi et al. (2020) reformulated the same algorithm in the form of server- and client-level optimization (Algorithm 1), which allowed them to bring techniques from the adaptive optimization literature to FL.
FEDAVG is efficient in that it requires only (d) computation on both the clients and the server, and (d) communication between each client and the server. To arrive at a similarly efficient algorithm for posterior inference, we focus on the following questions: (a) how estimate local and global posterior moments efficiently? (b) how to communicate local statistics to the server efficiently?
(1) Efficient global posterior inference. There are two issues with computing an estimate of the global posterior mode fe directly using Eq. 3. First, it requires computing the inverse of a d x d matrix on the server, which is an (dd) operation. Second, it relies on acquiring local means and inverse covariances, which would require (dd)) communication from each client. We propose to solve both issues by converting the global posterior estimation into an equivalent optimization problem.
Proposition 2 (Global Posterior Inference) The global posterior mode μ given in Eq. 3 is the mini- mizer of quadratic Q(0) T A0 b" 0, where A := Em -, and b - 9, H.
Proposition allows us to obtain a good estimate of μu by running stochastic optimization of the quadratic objective ((0) on the server. Note that the gradient of ((0) has the following form: N
VQ(0) = Σ 2 (0 (1,)),
which suggests that we can obtain ja by using the same Algorithm 1 as FEDAVG but using different client updates: A; := Σ, (00 - (.)). Importantly, as long as clients are able to compute A this approach will result in Odd) communication and ((d) server computation cost per round.
GGaussian posteriors can be further generalized to the exponential family for which closed form expressions be obtained under appropriate priors (Wainwright & Jordan, 2008). We leave this extension future work.
4

---

Published as a conference paper at ICLR 2021
(2) Efficient local posterior inference. To compute A, each client needs to be able to esti- mate the local posterior means and covariances. We propose to use stochastic gradient Markov chain Monte Carlo (SG-MCMC, Welling & Teh, 2011; Ma et al., 2015) for approximate sam- pling from local posteriors on the clients, so that these samples can be used to estimate ju," 's and E 's. Specifically, we use variant of SG- MCMC with iterate averaging (IASG, Mandt et al., 2017), which involves: (a) running lo- cal SGD for some number of steps to mix in the Markov chain, then (b) continued running of SGD for more steps to periodically produce sam- ples via Polyak averaging (Polyak & Juditsky, 1992) of the intermediate iterates (Algorithm 4). The more computation we can run locally on the clients each round, the more posterior samples can be produced, resulting better estimates of the local moments.
Algorithm 4 LASG Sampling (CLIENTMCMC)
input initial 0, loss f2( optimizer CLIENTOPT(O), B: burn-in steps, K: steps per sample, é: samples.
// Burn-in
1: for step 1 = I..... B do
0 1- CLIENTOPT(O, ((())
3: end for
I/ Sampling
4: for sample = 1, ,,.., 8 do
5:
Sg +
#/ Initialize iterates
6:
for step = l..... do
7:
0 CLIENTOPT (
8:
So f- So u (0)
9:
end for
10:
0, + AVERAGE(Se)
|/ Average iterates
11: end for output samples (0)....,
(3) Efficient computation of the deltas. Even if we can obtain samples ( 0z) via MCMC and use them to estimate local moments, }4, and E,, computing A; naïvely would still require inverting d x d matrix, i.e., (dd) compute and (dd) memory. The good news is that we are able to show that clients can compute A much more efficiently, in ((d) time and memory, using a dynamic programming algorithm and appropriate mean and covariance estimators.
let py be the sample mean, Se be the sample covariance, and Σ, := ppl + (1 -p)Ss be shrinkage estimator (Ledoit & Wolf. 2004b) of the covariance with Pt == 1/(1 + (E - 1))) for some p E [0, +oo). Then, for any O we can compute Â = ΣS (0- (u)) in (f2d) time and using (ld) memory.
Theorem 3 Given l approximate posterior samples (0,....,0, 0,).
Proof [Sketch] We give constructive proof by designing efficient algorithm for computing Å Our approach is based on two key ideas:
1. We prove that the specified shrinkage estimator of the covariance has a recursive decomposition into rank-l updates, i.e., Σ. E + c. x x,, where o is constant and x; is some vector. This allows us to leverage the Sherman-Morrine formula for computing the inverse of Σ..
2. Further, we design dynamic programming algorithm for computing Å exactly without storing the covariance matrix or its inverse. Our algorithm is online and allows efficient updates of Á, as more posterior samples become available.
See Appendix C for the full proof and derivation of the algorithm.
Note that the computational cost of Å consists of two components: (i) the cost of producing é approximate local posterior samples using IASG and (ii) the cost of solving a linear system using dynamic programming. How much of an over- head does add compared to simply running local SGD? It turns out that in practical settings the over- head almost negligible, Table I shows the time it takes a client compute the updates based on 5 local epochs (100 steps per epoch) using different algorithms (FEDAvG vs. our approach with exact or dynamic programming (DP) matrix inversion) on synthetic linear regressions. As the dimensionality grows, computational complexity of DD-based estimation of A becomes nearly identical to FEDAVG, which indicates that the majority of the cost in practice would come from SGD steps rather than our dynamic programming procedure.
Table 1: Computational complexity of the client up- dates for methods that use 5 local epochs measured in milliseconds (% denotes relative increase).
Dim
FLBAAG
Å (DP)
Å (exact)
100
72
91
+26%
82
+12%
1K
76
92
+21%
104
+36%
10K
80
93
+16%
797
+896%
100K
149
155
+4%
While in this work we use a variant of SG-MCMC, other techniques such HMC (Neal et al., 201 1) or NUTs (Hoffman & Gelman, 2014) can he used, too. We leave analysis of alternatiee approaches future work.
5

---

Published as a conference paper at ICLR 2021
The final algorithm, discussion, and implications. Putting all the pieces together, we arrive at the federated posterior averaging (FEDPA) algorithm for approximately computing the mode of the global posterior over multiple communication rounds. Our algorithm is variant of generalized federated optimization (Al (Algorithm 1) with a new client update procedure (Algorithm 3). Importantly, this also implies that FEDAVG can be viewed as posterior inference algorithm that estimates Σ with an identity and, as result, obtains biased client deltas FEDAVG := I(0- pt).
In Fig. in the introduction, we demonstrate the differences in behavior between FEDAVG and FEDPA that stem from the differences in their client updates. Biased client updates make FEDAVG converge to a suboptimal point: moreover, increasing local computation only pushes the fixed point further away from the global optimum. On the other hand, FEDPA converges faster and to better optimum, trading off bias for slightly more variance (becomes visible only closer to convergence). We see that FEDPA also substantially benefits from more local computation (more local samples).
Since the main difference between FEDAVG and FEDPA is, in fact, the bias-variance trade off the server gradient estimates (Eq. 4), we can view both methods as biased SGD (Ajallocian & Stich, 2020) and reason about their convergence rates as well as distances between their fixed points and correct global optima functions of the gradient bias. In Appendix A, we provide further details, discuss convergence, empirically quantify the bias and variance of the client updates for both methods, and analyse the effects of the sampling- -aasengg-bseddd approximations on the behavior of FEDPA.
EXPERIMENTS
Using a suite of realistic benehmark tasks introduced by Reddi et al. (2020), we evaluate FEDPA against several competitive baselines: the best versions of FEDAVG with adaptive optimizers well as MIME (Karimireddy et al., 2020)-a recently-proposse FEDAVG variant that also works with stateless clients, but uses control-variates and server-level statistics mitigate convergence issues.
Table 2: Statistics on the data and tasks. The number of examples per client are given with one standard deviation across the corresponding set of clients (denoted with +).. See description of the tasks in the text.
Dataset
Task
classes
# clients (train/test)
# examples p/ client (train/test)
EMNIST-62
CR
62
3,400 / 3,400 500 / 100
198 + 77 / 23 + 9 100 1 0/ 100 1 0
CIFAR-100
IR
100
LR
500
StackOverflow
342,477 204,088
397 + 1279/ / 81 + 301
NWP
10,000
5.1 THE SETUP
Datasets and tasks. The four benchmark tasks are based on the following three datasets (Table 2): EMNIST (Cohen et al., 2017). CIFARI00 (Krizhevsky et al., 2009), and StackOverflow (StackOver- flow, 2016). EMNIST (handwritten characters) and CIFAR 100 (RGB images) are used for multi-class image classification tasks. StackOverflow (text) is used for next-word prediction (also a multi-class classification task, historically denoted NWP) and tag prediction (a multi-label classification task, historically denoted LR because a logistic regression model is used). EMNIST was partitioned by authors (Caldas et al., 2018), CIFAR 100 was partitioned randomly into 600 clients with realistic heterogeneous structure (Reddi et al., 2020), and StackOverflow was partitioned by its unique users. All datasets were preprocessed using the code provided by Reddi et al. (2020).
Methods and models. We use generalized framework for federated optimization (Algorithm 1), which admits arbitrary adaptive server optimizers and expects clients to compute model deltas. As a baseline, we use federated averaging with adaptive optimizers (or with momentum) on the server and refer to it as FEDAVG- 1E or FEDAVG-ME, which stands for 1 or multiple local epochs performed by clients each round, respectively.* 4 The number of local epochs in the multi-epoch versions is hyperparameter. We use the same framework for federated posterior averaging and refer to it as FEDPA-ME. As our clients use IASG to produce approximate posterior samples, collecting
'Reddi et al. (2020) referred to federated averaging with adaptive server optimizers FEDADAM, FEDYOGI, etc. Instead, we select the best optimizer for each task and refer to the corresponding method simply FEDAVG.

---

Published as a conference paper at ICLR 2021
(a) CIFAR-100 100: Evaluation loss (left) and accuracy (right) for FEDAVG-ME and PEDPA-ME.
FedAvg-ME
FedPA-ME
FedAvg-ME
FedPA-ME
SG SG-MCMC campling phose
Burnin phose
Eval Accuracy
0.40
Eval
0.30
Rum-in phose
SG-MCMC sampling phase
0.20
500
10C0
1500
600
1003
1503
Round Number
Round Number
(b) StackOverflow LR: Evaluation loss (Ieft) and macro-F1 (right) for FEDAVG- ME and FEDPA-ME.
Fed//vg ME
FedPA-ME
FedAvg ME
FedPA-ME
8.00
0.18
SGAMCMC sampling phane
ploson
8.00
Evall macro-F1
Evall Loss
7.00
8.00
Buttiin phais
SG-MCMC sampling phsse
0.10
0
500
1000
1500
0
800
1003
1500
Round Number
Round number
Figure 2: Evaluation metries for FEDAVG and FEDPA computed each training round on (a) CIFAR- 100 and (b) StackOverflow LR. During the initial rounds (the "burn-in phase"), FEDPA computes deltas the same way as FEDAVG; after that, FEDPA computes deltas using Algorithm 3 and approximate posterior samples.
a single sample per epoch is optimal (Mandt et al., 2017). Thus FEDPA-ME uses M samples to estimate client deltas and has the same local and global computational complexity as MEDAVG-ME but with two extra hyperparameters: the number of burn-in rounds and the shrinkage coefficient from Theorem 3. As in Reddi et al. (2020), we use the following model architectures for each task: CNN for EMNIST-62, ResNet-18 for CIFAR-100, LSTM for StackOverflow NWP, and multi-label logistic regression on bag-of-words vectors for StackOverflov LR (for details see Appendix D).
Hyperparameters. For hyperparameter tuning, we first ran small grid searches for FEDAVG-ME using the best server optimizer and corresponding learning rate grids from Reddi et al. (2020). Then, used the best FEDAVG-ME configuration and did small grid search to tune the additional hyperparameters of -ME, which turned out not to be very sensitive (i.e., many configurations provided results superior FEDAVG). More hyperparameter details can be found in Appendix D.
Metrics. Since both speed of learning as well final performance are important quantities for federated learning, we measure: (i) the number of rounds takes the algorithm to attain a desired level of evaluation metric and (ii) the best performance attained within a specified number of rounds. For EMNIST-62, we measure the number of rounds it takes different methods to achieve 84% and 86% evaluation accuracy and the best validation accuracy attained within 500 and 1500 rounds. For CIFAR-100, we use the same metrics but use 30% and 40% as evaluation accuracy cutoffs and 1000 and 1500 as round number cutoffs. Finally, for StackOverflow, we measure the the number of rounds it takes to the best performance and evaluation accuracy (for the NWP task) and precision, recall 5, macro- and micro-F1 (for the LR task) attained by round 1500. We note that the total number of rounds was selected based on computational considerations (to ensure reproducibility within a reasonable amount of computational cost) and the intermediate cutoffs were selected qualitatively to highlight some performance points of interest. In addition, we provide plots of the evaluation loss and other metries for all methods over the course of training which show a much fuller picture of the behavior of the algorithms (most of the plots are given in Appendix E).
Implementation and reproducibility. All our experiments on the benchmark tasks were conducted in simulation using TensorFlow Federated (TFF, Ingerman & Ostrowski, 2019). Synthetic experiments were conducted using JAX (Bradbury et al., 2018). The JAX implementation of the algorithms is available at httpss: /tt:///itttat wwwpah...ooooooaaaaaeaeaeaaaaaaa com/ shedd fedpa. The TFF implementation will be released through https: /twtt/gtttht///g com/ com/ wowoo//ooolleeeeoooo http:/googleeeeerrrrtteeeeee federated.
"Centralized optimization of the CNN model on EMNIST-62 attains the evaluation accuracy of 88%.
7

---

Published as a conference paper at ICLR 2021
Table 3: Comparison of FEDPA with baselines. All metrics were computed on the evaluation sets and averaged over the last 100 rounds before the round limit was reached. The "number of rounds to accuracy" was determined based the 10-round running average crossing the threshold for the first time. The arrows indicate whether higher (t) or lower (4) better. The performance in each column denoted in bold.
(a) EMNIST-62
(b) CIFAR-100
accuracy (%, 1)
rounds (#, J)
accuracy
(%, f) 1500R
rounds (#, 30%
Method \@
500R
1500R
84%
86%
Method I@
1000R
40%
AFO
80.4
86.8
546
1291
AFO
31.9
41.1
898
1401
MIME
83.1
*84.9
464
MIME
33.2
*33.9
680
FEDAVG-1E IE
83.9
86.5
451
1360
FEDAVG-1E IE
24.2
31.7
1379
FEDAVG-ME
85.8
85.9
86
FEDAVG-ME
40.2
42.1
348
896
FEDPA-ME
86.5
87.3
84
92
FEDPA-ME
44.3
46.3
348
543
(c) SterkOverflow
NWP
LR (all metrics in %, f) recall@5
Method Metric
accuracy (%, 1)
rounds (#, J)
precision
ma-F1
mi-F1
AFO
23.4
1049
68.0
FEDAVG-1E
22.8
1074
74.58
69.1
14.9
43.8
FEDAVG-ME
23.0
870
78.65
68.7
15.6
43.3
FEDPA-ME
23.4
805
72.8
68.6
17.3
44.0
the best results taken from (Reddi et al., 2020). the best results taken from (Karimireddy et al., 2020).
results were only availahle for the method trained 1000 rounds.
5.2 RESULTS ON BENCHMARK TASKS
The effects of posterior correction of client deltas. As we demonstrated in Section 4, FEDPA essentially generalizes FEDAVG and only differs in the computation done on the clients, where we compute client deltas using an estimator the local posterior inverse covariance matrix, Σ, which requires sampling from the posterior. To be able to SG-MCMC for local sampling, FEDPA the burn-in regime (which is identical to FEDAVG) for number of rounds to bring the server state closer to the clients' local optima, 6 after which we "turn on" the local posterior sampling. The effect of switching from FEDAVG to FEDPA for CIFAR- 100 (after 400 burn-in rounds) and StackOverflow LR (after 800 burn-in rounds) presented on Figs. 2a and 2b, respectively." During the murn-in phase, evaluation performance is identical for both methods, but once FEDPA starts computing client deltas using local posterior samples, the loss immediately drops and the convergence trajectory changes, indicating that FEDPA is able avoid stagnation and make progress towards better optimum. Similar effects are observed across all other tasks (see Appendix E).
first run
While the improvement of FEDPA over FEDAVG some of the tasks is visually apparent (Fig. 2), we provide a more detailed comparison of the methods in terms of the speed of learning and the attained performance on all four benchmark tasks, summarized in Table 3 and discussed below.
Results on EMNIST-62 and CIFAR-100. In Tables 3a and 3b, we present a comparison of FEDPA against: tuned FEDAVG with fixed client learning rate (denoted FEDAVG-1 E and MEDAVG-ME), the best variation of adaptive FEDAVE from Reddi et al. (2020) with exponentially decaying client learning rates (denoted AFO), and MIME of Karimireddy et al. (2020). With more local epochs. we see significant improvement in terms of speed of learning: both FEDPA-ME and FEDAVG-ME achieve 84% accuracy on EMNIST-62 under 100 rounds (similarly, both methods attain 30% on CIFAR-100 by round 350). However, more local computation eventually hurts FEDAVG leading to
'If SGD cannot reach vicinity clients' local optima within the specified number local steps or epochs, estimated local means and covariances based on the SGD iterates can be arbitrarily poor.
'The number of burn-in rounds is hyperparamter and was selected for each task maximize performance. See more details Appendix D.
'We note that rumning burn-in for fixed number of rounds before switching to sampling was a design choice; other, more adaptive strategies for determining when switch from burn-in sampling certainly possible (e.g., use local loss values determine when start sampling). We leave such alternatives future work.
8

---

Published as a conference paper at ICLR 2021
worse optima: on EMNISTT-62, FEDAVG-ME is not able to consistently achieve 86% accuracy within 1500 rounds; on CIFAR-100, it takes extra 350 rounds for FEDAVG-ME to get to 40% accuracy.
Finally, federated posterior averaging achieves the best performance both tasks in terms of evaluation accuracy within the specified limit on the number of training rounds. On EMNIST-62 in particular, the final performance of FEDPA-ME M after 1500 training rounds is 87.3%, which, while only a 0.5% absolute improvement, bridges 41.7% of the gap between the centralized model accuracy (88%) and the best federated accuracy from previous work (86.8% Reddi et al., 2020).
Results on StackOverflow NWP and LR. Results for StackOverflow Although not as pronounced as for image datasets, we observe some improvement of FEDPA over FEDAVG here well. For NWP, we have an accuracy gain of 0.4% over the best baseline. For the LR task, we compare methods in terms of average precision, recall at 5, and macro-/micro--F), The first two metrics have appeared in some prior FL work, while the latter two the primary evaluation metrics typically used in multi-label classification work (Gibaja & Ventura, 2015). Interestingly, while FEDPA underperforms in terms of precision and recall, it substantially outperforms in terms of micro- and macro- averaged FI, especially the macro-F1. This indicates that while FEDAVG learns a model that can better predict figh-ffrequenyy labels, FEDPA learns a model that better captures rare labels (Yang, 1999; Yang & Liu, 1999). Interestingly, note while FEDPA improves on F1 metrics and has almost the same recall at 5, it's precision after 1500 rounds worse than FEDAVG. A more detailed discussion along with training curves for each evaluation metric are provided in Appendix E.
presented in Table 3c.
CONCLUSION AND FUTURE DIRECTIONS
In this work, we presented a new perspective on federated learning based on the idea of global posterior inference via averaging of local posteriors. Applying this perspective, we designed new algorithm that generalizes federated averaging, is similarly practical and efficient, and yields state-of-the-ar results on multiple challenging benchmarks. While our algorithm required number of specific approximation and design choices, we believe that the underlying approach has potential to significantly broaden the design space for FL algorithms beyond purely optimization techniques.
Limitations and future work. As we mentioned throughout the paper, our method has a number of limitations due to the design choices, such as specific posterior sampling and covariance estimation techniques. While in the appendix analyzed the effects of some of these design choices, exploration of: (i) other sampling strategies, (ii) more efficient covariance estimators (Hsich et al., 2013), (iii) alternatives to MCMC (such variational inference), and (iv) more general connections with Bayesian deep learning are all interesting directions to pursue next. Finally, while there is a known, interesting connection between posterior sampling and differential privacy (Wang al., 2015), better understanding of privacy implications of posterior inference in federated settings is an open question.
ACKNOWLEDGMENTS
The authors would like to thank Zachary Charles for the invaluable feedback that influenced the design of the methods and experiments, and Brendan McMahan, Zachary Garrett, Sean Augenstein, Jakub Konečný, Daniel Ramage, Sanjiv Kumar, Sashank Reddi, Jean-François Kagy for many insightful discussions, and Willie Neiswanger for helpful comments on the early drafts.
REFERENCES
Ahmad Ajallocian and Sebastian U Stich. Analysis of sgd with biased gradient estimators. arXiv preprint ar arXiv:2008.000 2020.
Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, H Brendan McMahan, et al. Towards federated learning at scale: System design. arXiv preprint 19Xiv:1902.01046, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin. and Skye Wanderman-Milme JAX: composable transformations of Python+NumPy programs, 2018. URL http: / thu.. coog e jax.

---

Published as a conference paper at ICLR 2021
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 1812.011,, 2018.
Zachary Charles and Jakub Konečný, On the outsized importance of learning rates local update methods. arXiv preprint arXiv:2007.00877 2020.
Yilun Chen, Ami Wiesel, Yonina C Eldar, and Alfred Hero. Shrinkage algorithms for mmse 2010.
covariance estimation. IEEE Transactions on Signal Processing,
Gregory Cohen, Saced Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. In 20/7 International Joint Conference on Neural Networks (IJCNN), pp. 2921-2926. IEEE, 2017.
Eva Gibaja and Sebastián Ventura. A tutorial on multilabel learning. ACM Computing Surveys (CSUR), 47(3):1-38 2015.
Isabelle Guyon. Design of experiments the nips 2003 variable selection benchmark. NIPS 2003
workshop feature extraction and feature selection, volume 253, 2003.
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: 18/ 1.03604, 2018.
Matthew D Hoffmann and Andrew Gelman. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1593-166 2014.
Cho-Jui Hsieh, Mátyás A Sustik, Inderjit S Dhillon, Pradeep K Ravikumar, and Russell Poldrack. BIG & QUIC: Sparse inverse covariance estimation for million variables. In Advances in neural information processing systems, pp. 3165-3173, 2013.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B Gibbons. The non-iid data quagmire of decentralized machine learning. arXiv preprint arXiv:1910000189, 2019.
Alex crated, 2019.
Ingerman and Krzys
Ostrowski. URL
Introducing tensorflow /etess///mmuummmmm wtt.ccmm//ce/eeeee honn/tessorflow//
fed-
Peter Kairouz, Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv: 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for on-device federated learning. arXiv preprint arXiv:
2019.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mchryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008.036600 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Olivier Ledoit and Michael Wolf. Honey, i shrunk the sample covariance matrix. The Journal of Portfolio Management, 30(4): 10-119,, 2004a.
Olivier Ledoit and Michael Wolf. A well-conditioned estimator for large-dimensional covariance matrices. Journal of multivariate analysis, 88(2):365 411, 2004b.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated Iearning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):00-60, 2020.
Xiang Li. Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang, On the convergence of fedavg non-iid data. arXiv preprint arXiv: 1907.00000000000 0227.0022..9.. 2019.
10

---

Published as a conference paper at ICLR 2021
Wei Yang Bryan Lim, Nguyen Cong Luong, Dinh Thai Hoang, Yutao Jiao, Tin--Chngg Liang, Qiang Yang, Dusit Niyato, and Chunyan Miao. Federated learning in mobile edge networks: A comprehensive survey. IEEE Communications Surveys & Tutorials, 2020.
Jun S Liu. Metropolized independent sampling with comparisons to rejection sampling and impor- tance sampling. Statistics and computing, 6(2):113-119, 1996.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient memc. In Advances in Neural Information Processing Systems, pp. 2917-2925, 2015.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent approximate bayesian inference. The Journal of Machine Learning Research, 18(1):487344977, 2017.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication
learning of deep networks from decentralized data. Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017.
Kevin P Murphy. Machine learning: probabilistic perspective. MIT press, 2012.
Radford M Neal et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Willie Neiswanger, Chong Wang, and Eric Xing. Asymptotically exact, embarrassingly parallel meme. arXiv preprint arXiv: 1311,4780, 2013.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on optimization, 19(4):1574- 1609, 2009.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
Reese Pathak and Martin J Wainwright. FedSplit: an algorithmic framework for fast federated optimization. arXiv preprint arXiv:2005.05238, 2020.
Boris T Polyak. Some methods of speeding up the convergence iteration methods. USSR
Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, 30(4):838-88) -5188888885,,,, 1992.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295, 2020.
Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith. On the convergence of federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06177. 3, 2018.
Juliane Schäfer and Korbinian Strimmer. A shrinkage approach to large-scale covariance matrix estimation and implications for functional genomics. Statistical applications in genetics and molecular biology, 4(1), 2005.
Steven L Scott, Alexander W Blocker, Fernando V Bonassi, Hugh A Chipman, Edward I George, and Robert McCulloch. Bayes and big data: The consensus monte carlo algorithm. International Journal of Management Science and Engineering Management, 11(2):78-88 2016.
StackOverflow. Stack Overllow Data, 2016.
URL https:/ / www. www.kaggle.com.
Seackoverfloww/sttckveeerr
Martin Wainwright and Michael I Jordan. Graphical models, exponential familiis,,, and variational inference. Now Publishers Inc, 2008.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective
inconsistency problem in heterogeneous federated optimization. arXiv preprint arXiv:2007.074: 2020.
11

---

Published as a conference paper at ICLR 2021
Yu- Xu-Xiang Wang, Stephen Fienberg, and Alex Smola. Privacy for free: Posterior sampling and stochastic gradient monte carlo. In International Conference on Machine Learning, pp. 2492-2502, 2015.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681-688 2011.
Blake Woodworth, Kumar Kshitij Patel, and Nathan Srebro. Minibatch vs local sgd for heterogeneous distributed learning, arXiv preprint arXiv:2006.04755 2020a.
Blake Woodworth, Kumar Kshitij Patel, Sebastian U Stich, Zhen Dai, Brian Bullins, H Brendan McMahan, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? arXiv preprint arXiv:2002.07839, 2020b.
Yiming Yang. An evaluation statistical approaches to categorization. Information retrieval, (1-2):69-90, 1999.
Yiming Yang and Xin Liu. A re-examination of text categorization methods. In Proceedings of the 22nd annual international ACM SIGIR conference on Research and development in information retrieval, pp. 42-49, 1999.
Hao Yu. Sen Yang. and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.0058/ 2018.
12

---

Published as a conference paper at ICLR 2021
A PRELIMINARY ANALYSIS AND ABLATIONS
In Section 4, we derived federated posterior averaging (FEDPA) starting with the global posterior decomposition (Proposition I, which is exact) and applying the following three approximations:
1. The Laplace approximation of the local and global posterior distributions.
2. The shrinkage estimation of the local moments.
3, Approximate sampling from the local posteriors using MCMC.
We have also observed that FEDAVG is special case of FEDPA (from the algorithmic point of view), since it can be viewed as also using the Laplace approximation for the posteriors, but estimating local covariances Eg's with identities and local means using the final iterates of local SGD.
In this section, we analyze the effects of approximations and 3 the convergence of FEDPA. Specifically, we first discuss the convergence rates of FEDAVG and FEDPA as biased stochastic gradient optimization methods (Ajalloeian & Stich, 2020). We show how the bias and variance of the client deltas behave for FEDAVG and FEDPA as functions of the number samples. We also analyze the quality of samples produced by IASG (Mandt et al., 2017) and how they depend the amount of local computation and hyperparameters. Our analyses are conducted empirically on synthetic data.
A.1 DISCUSSION OF THE CONVERGENCE OF FEDPA VS. FEDAVG
First, observe that if each client able to perfectly estimate their A; = = 1e), the problem solved by Algorithm I simply becomes an optimization of quadratic objective using unbiased stochastic gradients, A := M Σn A. The noise in the gradients in this case comes from the fact that the server interacts with only a small subset of M out of N clients in each round. This is a classical stochastic optimization problem with well-known convergence rates under some assumptions on the norm of the stochastic gradients (e.g., Nemirovski et al., 2009). The rate convergence for SGD with a O(t decaying learning rate used on the server is (1/vV). can be further improved to (1/tt) using Polyak momentum (Polyak, 1964) iterate averaging (Polyak & Juditsky, 1992).
In reality, both FEDAVG and FEDPA produce biased estimates ÂFEDAVa and Å ÂFENPA, respectively. Thus, can analyze the problem as SGD with biased stochastic gradient estimates and let Å V (F(0,) + b(O)) + n(0.) where b(O,) and n (O,,)) < are bias and noise terms. Following Ajallocian Stich (2020), we can further assume that the bias and noise terms are boun-bounde as follows.
Assumption ( bbounded bias) There exist constants S m < 1 and C² 2 0 such that |((0) S ml (|] + VO E RRd.
(5)
Assumption 5 -bounddd noise) There exist constants 0 S M < 1 and 0² 0 such that Eg [|n E) 1 & (|) + +o² A0 Rd
Under these general assumptions, the following convergence result holds.
Theorem (Ajalloeian Stich (2020), Theorem 2) Let F(0) be L-smooth. Then SGD with a 1-m
{ E 2ML
LF 1/2 and gradients that satisfy Assumptions 4, 5 achieves
learning rate a == min the vicinity of stationary point,
[[|V [0) = =O 0 +8+ (8+ in T iterations, where
I-m)
1
s(1 - m) M +
0²
LF
T=o
1+ 1
1-m
Note that SGD with biased gradients is able to converge to vicinity of the optimum determined by the bias term <2/(1- m.. For FEDAVG, since the bias is not countered, this term determines the distance between the stationary point and the true global optimum. For FEDPA, since ÁFEDPA - A with more local samples, the bias should vanish as we increase the amount of local computation.
Determining the precise statistical dependence of the gradient bias on the local samples is beyond the scope of this work. However, to gain more intuition about the differences in behavior of FEDPA and FEDAVG, below we conduct an empirical analysis of the bias and variance of the estimated client deltas synthetic least squares problems, for which exact deltas can be computed analytically.
13

---

Published as a conference paper at ICLR 2021
(a) FEDAVG bias and variance as functions of the number of local steps.
Dimensionality: 10
Dimensionality: 100
Dimensionality: 1000
20
10a
300
40
Variance (Fro)
103
15
30
3
6
10
100
102
10
0
25
50
75
100
50
100
150
200
2000
4000
6000
Steps
Local Steps
Steps
(b) FEDPA bias and variance functions of the number of local steps. The burn-in steps were not included. For dimensionality 10, 100, and 1000, the shrinkage was fixed to 0.01, 0.005, and 0.001, respectively.
Dimensionality:
Dimensionality: 100
Dimensionality:1000
40
400
60000
Variance (Fro)
260
300
4000
10
200
240
25
2000
20
220
10a
10
10*
102
107
10*
10²
10a
10*
Local Steps
Local Steps
Local Steps
(c) FEDPA bias and variance functions of the shrinkage parameter. For dimensionality 10, 100, and 1000, the number local steps was fixed to 5,000, 10,000, and 50,000, respectively.
Dimensionality: 10
Dimensionality: 100
Dimensionality: 1000
10a
104
50
Varia Variance ance (Fro)
104
10a
400
3.5
40-
350
10a
3.0
109
30
10
300
10-3
10
10-1
10
10-1
10-2
Shrinkage
Shrinkage
Shrinkage
Figure 3: The hias and variance tradeoffs for FEDAVG and FEDPA as functions of the estimation parameters.
Quantifying empirically the bias and variance of Á for FEDPA and FEDAVG. We measure the empirical bias and variance of the client deltas computed by each of the methods on the syn- thetic least squares linear regression problems generated according to Guyon (2003) using the make_regrest i on function from scikit-learn. The problems were generated as follows: for each dimensionality (10, 100, and 1000 features), we generated 10 random least squares problems, each of which consisted of 500 synthetic data points. Next, for each of the problems we generated 10 random initial model parameters (O,...., 110} and for each of the parameters
computed the exact A; as well as Å FEDAVa, and A for different numbers of local steps; for Å FEDPA we also varied the shrinkage hyperparameter. Using these sample estimates, we further computed the L2 Lg-norm of the bias and the Frobenius norm of the covariance matrices as functions of the number of local steps.
The results are presented on Fig. 3. From Fig. 3a, see that the amount of local computation increases, the bias FEDAVG delta estimates grows and variance reduces. For FEDPA (Fig. 3b), the trends turn out be the opposite: as the number of local steps increases, the bias consistently reduces; the variance initially goes up. but with enough samples joins the downward trend. Note that the initial upward trend in the variance is due the fact that used the same fixed shrinkage p regardless of the number local steps. To avoid sharp increases the variance, must be selected for each number local steps separately; Fig. 3c demonstrates how the bias and variance depend on the shrinkage hyperparameter for some fixed number of local steps." 10
9htpps: scikicikittto.. www..oor../ggooooo.oooo.o.. httgg//tabb stabballllee..... e/moodu es/generated earr war.eaaataeett... datasets. make_regresse on. html
10nne could also posterior samples to estimate the hest possible p that balance the bias-variance tradeoff (e.g., Chen al., 2010) and avoids sharp increases in the variance.
14

---

Published as a conference paper at ICLR 2021
A.2 ANALYSIS OF THE QUALITY OF BASG-BASED SAMPLING AND COVARIANCE
The more and better samples we can obtain locally, the lower the bias and variance of the gradients of ((0) will be, resulting in faster convergence sampling, we proposed to use a variant of SG-MCMC called Iterate Averaged Stochastic Gradient (IASG) developed by Mandt al. (2017), given in Algorithm The algorithm generates samples by simply averaging every K intermediate iterates produced by a client optimizer (typically, SGD with some a fixed learning rate a) after skipping the first B iterates as burn-in phase."
fixed point closer to the global optimum. For local
How good are the samples produced by IASG and how do different parameters of the algorithm affect the quality of the samples? To answer this question, we run IASG on synthetic least squares problems, for which we can compute the actual posterior distribution and measure the quality of the samples by evaluating the effective sample size (ESS, Liu, 1996; Owen, 2013). Given approximate posterior samples (0,,...,0, ,r.,, the ESS statistic can computed as follows:
2
ESS ((0.))-1)=
Σ
where weights Wy must be proportional to the posterior probabilities, or equivalently to the loss. Effeets of the dimensionality, the number of data points, and IASG parameters on ESS. The results of our synthetic experiments are presented below in Fig. 4. The takeaways are as follows:
More burn-in steps (or epochs) generally improve the quality samples.
The larger the numher of steps per sample the better (less correlated) the samples are.
The learning rate is the most sensitive and important hyperparameter-if too large, IASG might diverge (happened the 1000 dimensional case); too small, the samples become correlated. Finally, the quality of the samples deteriorates with the increase in the number of dimensions.
(a) as function the number burn-in steps. (Steps per sample: 50.)
Dimensionality: 10
Dimensionality: 100
Dimensionality: 1000
100
100
100
ESS (%)
50
100
150
200
100
200
300
400
200
400
600
800
Burn-in Steps
(b) ESS as a function of number of steps per sample. (Burn-in steps: 100.)
Dimensionality: 10
Dimensionality: 100
Dimensionality: 1000
ESS (%6)
100.0
100
99.5
99.9
50
100
150
200
250
50
100
150
200
250
50
100
150
200
250
Steps per Sample
(c) ESS function of the learning rate. (Burn-in steps: 100, steps per sample: 50.)
Dimensionality: 10
Dimensionality: 100
Dimensionality: 1000
100
100
100
ESS (%)
10-2
10
10a
10-3
10 -2
10
10
10-2
10
Learning Rate
Figure 4: The ESS statistics for samples produced by IASG random synthetic least squares linear regression problems of dimensionality 10, 100, 1000. Total number of data points per problem: 500, batch size: 10. In (a) (b) the learning rate was set 0.1 for 10 and dimensions, and 0.01 for 1000 dimensions.
Note that in our experiments in Section 5, instead of using B local steps for burn-in at each round, we used several initial rounds burn-in-only rounds, running FEDPA in the FEDAVG regime.
15

---

Published as a conference paper at ICLR 2021
B PROOFS
Proposition 1 (Global Posterior Decomposition) Under the uniform prior, any global posterior
distribution that exists decomposes into product of local posteriors: P [0 D) x II (0 D.).
Proof Under the uniform prior, the following equivalence holds for (0 D) as a function 0:
D) 0)
II P (= 0)
III P 0= il
II
P (0 D,)
(8)
#ED
1=1
local likelihood
The proportionality constant between the left and right hand side in Eq. 8 is IT P P(D.)/P(D),
The global posterior mode μa given in Eq. 3 is the mini- mizer of a quadratic Q(0) := T A0 b- 6, where A == 9
Proposition 2 (Global Posterior Inference)
a. and b = 1; μ.
Proof The statement of the proposition (implicitly) assumes that all matrix inverses exist. Then, the quadratic 2(0) is positive definite (PD) since A PD as convex combination PD matrices Σ Thus, the quadratic has unique solution O where the gradient of the objective vanishes:
AØ =
r* == A-
(9)
1=1
which implies that μ is the unique minimizer of Q(0).
C COMPUTATION OF CLIENT DELTAS VIA DYNAMIC PROGRAMMING
In this section, provide constructive proof for the following theorem by designing efficient
algorithm for computing À Σ (0 jug) on the clients in time and memory linear in the number of dimensions of the parameter vector Rdd.
0₂}, let jag be the sample mean, Se be the sample covariance, and Σ, ppl (1 ps) be shrinkage estimator (Ledoit & Wolf, 2004b) of the covariance with Pt == 1/(1 (E - 1) for some p E (00). Then, for any O, we can compute Å -.- 0 (u.) in ((ad) time and using (Ed) memory.
Theorem Given approximate posterior samples (
The naïve computation of update vectors (iee, where we first estimate ja, and Eg from posterior samples and use them to compute deltas) requires (dd) storage and Odd) compute on the clients and is both computationally and memory intractable. We derive an algorithm that, given posterior samples, allows us compute Å using only ((dd) memory and ( compute.
The algorithm makes use of the following two components:
1. The shrinkage estimator of the covariance (Ledoit & Wolf, 2004b), which is known to be well-conditioned even in high-dimensionn smaller than the number of dimensions) and is widely used in econometrics (Ledoit & Wolf, 2004a) and computational biology (Schäfer & Strimmer, 2005).
settings (i.e., when the number samples is
2. Incremental computation of "
(0, - (iz) that exploits the fact that each new posterior sample only adds a rank-1 component to Σ, and applies the Cherman-Morrisee formula to derive dynamic program for updating Á
Notation.
For the sake of this discussion, denote 0 (i.e., the server state broadcasted to the clients at round f) as Xg. drop the client index i, denote posterior samples as x,, sample mean as x; & ES= Xj, and sample covariance as Se + Ej=((x xx)(x; x.)T.
16

---

Published as a conference paper at ICLR 2021
C.I THE SHRINKAGE ESTIMATOR OF THE COVARIANCE
Ledoit & Wolf (2004b) proposed to estimate a high-dimensional covariance matrix using a convex combination of identity matrix and sample covariance (known the LW shrinkage estimator):
Ee(pe) == I. Pe)
(10)
where Pr is scalar parameter that controls the bias-variance tradeoff of the estimator. As an aside, while Pg can be arbitrary and the optimal Pe requires knowing the true covariance Σ, there are near-optimal ways estimate Pe from the samples (Chen al., 2010), which we discuss at the end of this section.
In this section, we focus deriving expression for Pt function of t = 1,..., l that ensures that the difference between Σ, and Σ is rank-1 matrix (this is not the case for arbitrary p's).
Derivation of a shrinkage estimator that admits rank-1 updates. Consider the following matrix:
Σ
I+ B,S.,
(11)
where 3t is a scalar function of = 1, L. We would like to find 3, such that Σ, = - U,. where U, is rank-1 matrix, i.e., the following equality should hold:
B.S, = - 7: U.
(12)
To determine the functional form of B,, we need recurrent relationships for x2 and St. For the former, note that the following relationship holds for two consecutive estimates of the sample mean, - and X2:
_,- (t - 1) +
8,-1 + 7(xx &c-1)
(13)
This allows to expand S, follows:
(t - 1)5, = Σ (x,- - xx)(Xx/ - xx²)
1=1
Xx - X )
X1 8c-1
Xj X1-11
x, - x,-11
t
t
X; 81-1
(x) x1-1)) (x, xx-1) -2
(x) - x(-1))+ +
(14)
t
=(1-2)$1-1
=0
t t-
t--
(xi - xx-1)(xx - xx)
T +
(x, - xx-1)( - x(-1)
f2
= (1 2-)))-111 + 1 1 (x: xx-1)(xx - x(-1)
Thus, have the following recurrent relationship between 8, and S,
| - 2
1 [xx x, 1xx x1-1)
S.
t - $.-1+ + t
(15)
Now, we can plug (15) into (12) and obtain the following equation:
Be t 2
S.-1 i P (x. 8.-1) xx-1) -B- 72
(16)
which implies that U, == (x, xx-1)(x;; x(-1)__ n: = Belt, and the following telescoping expressions for Be:
t t-1
t
8 =
:-1= =
3r2= = ... -t-1)) 1)
(17)
t-
f
17

---

Published as a conference paper at ICLR 2021
where set $2 p E [0, +00) to be constant. Thus, if we define = I p(t -1)S, then the following recurrent relationships will hold:
Σ = I,
Σ2 I+ pS. = E + p + - (xx2xx))(x> - -x1)²
20
Σ3 I 2 = ΣΦ + [x3 (x3 - x2)(X3 - x2)
Σ, I+ (t- 1 (t-1)p 1)p (x, - xx) xx-1)
(18)
Finally, we can obtain shrinkage estimator of the covariance from E, by normalizing coefficients:
1
(t 1)p
Σ
I
S. =pp
(19)
1 1+ (t 1)p
1+ (t 1)p
Pt
1-Pt
Note that Σ₁ and Σg St t +
C.2 COMPUTING DELTAS USING SHERMAN-MORRISON AND DYNAMIC PROGRAMMING
Since Σg is proportional to Σ, and the latter satisfies recurrent rank-1 updates given in Eq. 18,
denoting u, : Xt. X I, we can express Ex- = Σ, /p: using the Sherman-Morrison formula:
78 E mnu -.
Σ = Ex1
(20)
1 1+ 1+98 78
Note that we would like to estimate Å = ES (xn xt), which can be done without computing or storing any matrices if we know ΣТ Ug and Σ (xo Xe).
Denoting A Σ (xx- x1), and knowing that xo (xo- xx-1) -ue/l (which follows
from Eq. 13), we can compute Å using the following recurrence:
Â == Xo X1. V1,2 = X2 X1.
// initial conditions
(21)
U, == Xx X1-1+ V1-1, = $ ur
// recurrence for u, and Vy-1,t
(22)
"t (1,] أ u, vo-1a)
Á = Â 1+
V recurrence for A
(23)
1 (n. v
Á, Å./pt
I/ final step for A,
(24)
Remember that our goal is to avoid storing d d matrices throughout the computation. In the above recursive equations, all expressions depend only on vector-vector products except the one for V,-1,t which needs matrix-vector product. To express the latter one in the form of vector-vector products,
we need another 2-index recurrence on ViJ = Σ, Uy:
V1
U2
V1,3 = Ug
V1,1 = Ut
// initial conditions
(25)
S 21-24 IT, 1
--1
7-1
Vt-18 =
1-2
u
// Sherman-Morrion
(26)
1 1-1
(v (V-22-1111 U1)
Vv-22
+ 7
(v
22-2,-11 Up-1) V2-2
(27)
1-1 7 1
=
Vk-1kk
I/ final expression for V.-
(28)
k= 1 7%
V
18

---

Published as a conference paper at ICLR 2021
Now, equipped with these two recurrences, given stream of samples X1 .g...... .,,.,.,,,,,...........,......... ..,..., we compute À for 2 2 based Xt, {ux (uu))=11 (V--2 and Å using the following two steps:
1. Compute i, and Vt-1, using the second recurrence.
2. Compute A from ue, Vt-1,6, and á using the first recurrence.
For each new sample in the sequence, repeat the two steps to obtain the updated Á, estimate, until we have processed all samples. Note that the first step requires O(1) vector-vector multiplies, ie. (td)) compute, and (dd) memory, and the second step a (1)) number of vector-vecte multiplies. As result, the computational complexity of estimating A is ( and the storage needed for the dynamic programming state represented by tuple (uux) k= ( Á.-1) is Old)).
The any-time property of the resulting algorithm. Interestingly, the above algorithm is online as well as any-time in the following sense: as we keep sampling more from the posterior, the estimate of A keeps improving, but if stopped at any time, the algorithm still produces the best possible estimate under the given time constraint. If the posterior sampler is stopped during the burn-in phase or after having produced only I posterior sample, the returned delta will be identical to FEDAVG. By spending more compute on the clients (and a bit of extra memory), with each additional posterior sample x,, we have A
E (x0 1t).
Optimal selection of p. Note that to be able to run the above described algorithm in online fashion, we have to select and commit to before seeing any samples. Altematively, if the online and unnecessary,
any-time properties of the algorithm then infer a near-optimal f. from these samples-e.g. using the Rao-Blackwellized version of the LW estimator (RBLW) the oracle approximating shrinkage (OAS), both proposed and analyzed by Chen et al. (2010)-am) and then use the inferred p. to compute the corresponding delta using our dynamic programming algorithm.
first obtain É posterior samples (xx 1==
D DETAILS ON THE EXPERIMENTAL SETUP
In this part, provide additional details experimental setup, including detailed description of the datasets and tasks, models, methods, and hyperparameters.
D.1 DATASETS, TASKS, AND MODELS
Statistics of the datasets used in our empirical study can be found in Table 2. All the datasets and tasks considered in our study are a subset of the tasks introduced by Reddi et al. (2020).
The dataset is comprised of 28 x 28 images of handwritten digits and lower and upper case English characters (62 different classes total). The federated version of the dataset was introduced by Caldas et al. (2018), and is partitioned by the author of each character. The heterogeneity the dataset is coming from the different writing style of each author. We use this dataset for the character recognition task, termed EMNIST CR in Reddi et al. (2020) and the same model architecture, which a 2-layer convolutional network with x 3 kernel, max pooling, and dropout, followed by 128-unit fully connected layer. The model adopted from the TensorFlow Federated library: https: /tap::////b/e /bttt . 1y/ /314//1440 3141LKv.
EMNIST-62.
The federated version of CIFAR-100 100 was introduced by Reddi et al. (2020). The training set of the dataset is partitioned among 500 clients, 100 data points per client. The partitioning was created using two-step latent Dirichlet allocation (LDA) over to "coarse" to "fine" labels which created a label distribution resembling a more realistic federated setting. For the model, also following Reddi et al. (2020), used modified ResNet-18 with group normalization layer instead of batch normalization, as suggested by Hsieh et al. (2019). The model was adopted from the TensorFlow Federated library: https: //bit t. 1y/ 3.3 6M.
CIFAR-100.
19

---

Published as a conference paper at ICLR 2021
Table 4: Selected optimizers for each task. For SGD, m denotes momentum. For Adam, & =.9, -0.99.
Hyperparameter
EMNIST-62
CIFAR-100
StackOverflow NWP
StackOverflow LR
SERVEROPT CLIENTÖPT
SGD (m 0.9)
SGD (m 0.9)
Adam (T 10-2) SGD 0.0)
Adagrad (r 1005) SGD 0.9)
SGI (m 0.9)
SGD 0.9)
# clients p/round
100
20
Table 5: Hyperparameter grids for each task.
Hyperparameter
EMNIST-62
CIFAR-100
StackOverflow
StackOverflow LR
NWP
Server learning rate
(0.01, 0.05, 0.1, 0.1, 1,5] {0.001, 0.005, 0.01, 0.05, 0.13
(0.01, 0.05, 0.1, 0.5, 1 [0.01, 0.05, 0.1)
{0.1, 1,5, 10) [1,5, 10,50, 100)
Client learning rate
Client epochs
[2 20)
FEDPA burn-in FEDPA shrinkage
(100, 200, 400, 600, 800]
{0.0001, 0.001, 0.01, 0.1,
StackOverflow. The dataset consists of text (questions and answers) asked and answered by the total of 342,477 unique users, collected from https: /stsskteeacoeooooooooooooo com. The federated version of the dataset partitions it into clients by the user. In addition, questions and answers in the dataset have associated metadata, which includes tags. We consider two tasks introduced by Reddi et al. (2020): the next word prediction task (NWP) and the tag prediction task via multi-label logistic regression. The vocabulary of the dataset is restricted to 10,000 most frequently used words for each task (i.e., the NWP task becomes a multi-class classification problem with 10,000 classes). The tags are similarly restricted 500 most frequent ones (i.e., the LR task becomes multi-label classification proble with 500 labels).
For tag prediction, we use simple linear regression model where each question or answer are represented by normalized bag-of-wordd vector. The model was adopted from the TensorFlow Federated library: https: /tttt 1 jEX)AeY.
For the NWP task, we restrict each client to the first 128 sentences in their dataset, perform padding and truncation ensure that sentences have 20 words, and then represent each sentence as sequence of indices corresponding the 10,000 frequently used words, as well as indices representing padding, out-of-vocabulary (OOV) words, beginning of sentence (BOS), and end of sentence (EOS), We note that accuracy of next word prediction is measured only on the content words and on the OOV, BOS, and EOS symbols. We use an RNN model with 96-dimensional word embeddings (trained from scratch), dio-eninnsional LSTM layer, followed by fully connected output softmax layer. The model was adopted from the TensorFlow Federated library: https:: :tttpp::/bii.. ly/ 225/ 3x.
D.2 METHODS
As mentioned in the main text, we used FEDAVG with adaptive server optimizers with or multiple local epochs per client as our baselines. For each task, we selected the best server optimizer based on the results reported by Reddi et al. (2020), given in Table 4. We emphasize, even though we refer to all our baseline methods as FEDAVG, the names of the methods as given by Reddi et al. (2020) should be FEDAVGM for EMNIST-62 and CIFAR-100, FEDADAM for StackOverflow NWP and FEDADAGRAN for StackOverflow LR. Another difference between our baselines and Reddi et al. (2020) is that ran SGD with momentum the clients for EMNIST-62, CIFAR-100 100, and StackOverflow LR, as that improved performance of the methods with multiple epochs per client.
Our FEDPA methods used the configurations FEDAYG baselines; moreover, FEDPA and FEDAVG were identical (algorithmically) during the burn-in phase and only different in the client-side computation during the sampling phase of FEDPA.
20

---

Published as a conference paper at ICLR 2021
Table 6: The best selected hyperparameters for each task.
Hyperparameter
EMNIST-62
CIFAR-100
StackOverflow NWP
StackOverflov LR
Server learning rate
0.5
0.5
1.0
5.0
Client learning rate
0.01
0.01
0.1
50.0
Client epochs
5
10
5
5
FEDPA burn-in FEDPA shrinkage
100
400
800
800
0.1
0.01
0.01
0.01
D.3 HYPERPARAMETERS AND GRIDS
All hyperparameter grids are given Table The best server and client learning rates were selected based on the FEDAVG performance and used for FEDPA. The best selected hyperparameters are given Table
E ADDITIONAL EXPERIMENTAL RESULTS
We provide additional experimental results. As mentioned in the main text, results presented in Table 3 were selected to highlight the differences between the methods with respect to two metrics of interest: (i) the number of rounds until the desired performance, and (ii) the performance achievable within fixed number of rounds. A much fuller picture is given by the learning curves of each method. Therefore, we plot evaluation losses, accuracies, and metrics interest over the course of training. On the plots, individual values at each round are indicated with x-markers and the 10-round running average with a line of the corresponding color.
EMNIST-62. Learning curves for FEDAVG and FEDPA on EMNIST-62 are given Fig. Fig. 5a shows the best FEDAVG- EE, FEDAVG-5E. and FEDPA-5E models and Fig. 5b shows the best FEDAVG-20E: 20F, and FEDPA-200 Apart from the fact that multi-epoch versions converge significantly faster than the 1-epoch FEDAVG-1 E, note that the effect of bias reduction when switching from the burn-in to sampling in FEDPA becomes much more pronounced in the 20-epoch version.
CIFAR-100 and StackOverflow. Learning curves for various models CIFAR-100 and Stack- Overflow tasks are presented in Figs. 6 and 7. The takeaways for CIFAR-100 and StackOverflow NWP essentially the same for EMNIST-62 -uchh faster convergence with the increased number of local epochs and visually noticeable improvement in losses and accuracies due to compling-based bias correction in client deltas after the burn-in phase over. Interestingly, we that on StackOver- flow LR task FEDAVG- EE clearly dominates multi-epoch methods in terms of the loss and recall 5, losing in precision and macro-F1. Even more puzzling the significant drop in the average precision of FEDPA-ME after the switching to sampling. while at the same time a jump in recall and Fl metrics. This indicates that the global model moves to different fixed point where over-predicts positive labels (i.e., less precise) but also less likely to miss labels (i.c., higher recall rare labels, and as a result jump in macro-Fl), The reason why this happens, however, is unclear.
21

---

Published as a conference paper at ICLR 2021
EMNIST-62: Evaluation accuracy for FEDAVG- IE, FEDAVG-SEL and FEDPA-S5 SE.
FedAvg-1E
FE
FedPA-ME
Feelvg-- 1E
FedAvg-ME
FedAug-ME
0.7
Eval Accurscy
0.6
0.84
Evall
0.5
0.22
0.80
250
750
1000
500
750
1000
Round Number
Round Number
(b) EMNIST-62: Evaluation loss and accuracy for FEDAvG-20E and PEDPA-20E.
FadAug-ME
FodAug-ME
0.80
0.88
0.70
Eval Accuracy
0.86
Eval
0.84
0.50
0.82
Round Number
Figure 5: Evaluation metrics for FEDAVG and FEDPA computed at each training round on EMNIST-62.
CIFAR-10% 100: Evaluation loss and accuracy for FEDAVG- IE, FEDAVG-10E, and FEDPA-
FedAvg-1E
FedAvg-ME
FedPA-ME
FedAug-1E
FedAvg-ME
FedPA-ME
0.50
0.40
Eval Accuracy
0.30
0.0
(b) StackOverflow-NWWW -WPP Evaluation perplexity and accuracy for FEDAVG- 1E, FEDAVG-5E, and FEDPA-55 P5-
FedAug-1E
FedAug-ME
FedAug-1E
FedAvg
0.25
Eval Perplexty
Eval Accuresy
0.25
0.20
0.18
0.15
Round Number
Round Number
Figure 6: Evaluation metries for FEDAYG and FEDPA computed at each training round on (a) CIFAR-100 and (b) StackOverflow NWP tasks.
22

---

Published as a conference paper at ICLR 2021
FadAvg-1E
FodAvg
FodAug-1E
FodAug-ME
8.00
0.80
Precision
Evall Loss
0.63
FedAvg-1E
FedAvg-MM
FedPA-ME
FedAvg--E 1E
FedAvg-ME
FedPA-ME
0.70
0.45
Eval Recaliig
0.58
0.43
0.68
0.40
0.64
Evall
0.38
FodAug-1E
FadAug-ME
FadPA-ME
0.58
Eval macroF1
0.95
000
1000
1000
Round Number
Figure Evaluation metries for FEDAVG and FEDPA computed at each training round StackOverflow LR. Evaluation loss, average precision and recall, micro- macro-averaged FEDAVG-15 1E, FEDAYG-SE.
FEDPA-5E.
23


Reviews:
[{"review_id": "XCbkfYRtUDq", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is nice and seems technically sound. However, some part must be clarified.\nFor instance, the discussion of the state-of-the-art must be improved.\nMy main concerns is the degree of novelty. It is not clear the difference with other works.\nSee my comments below.\n\n- Regarding Eq. (1): you state that \"For example, least squares loss corresponds to likelihood under a\nGaussian model, cross entropy loss corresponds to likelihood under a categorical model, etc. Thus, Eq. 1 corresponds to maximum likelihood estimation (MLE) of the model parameters \u0012.\" This sentence must be clarified since it is not straightforward to see it from Eq. (1) and it is an important equation and sentence for your work.  Do you mean that F(\\theta) is a log-posterior or a  log-likelihood?\nGenerally,  it is not a mixture of components. Clarify this point.\n\n- Regarding Eq. (3):  the state-of-the-art related to this equation must be improved. For instance, the following relevant contributions must be considered\n\nLavancier, F., Rochet, P.: A general procedure to combine estimators. preprint arXiv:1401.6371 (2014)\n\nD. Luengo et al, \"Efficient linear fusion of partial estimators\", Digital Signal Processing, Volume 78, Pages: 265-283, 2018.\n\nCattivelli, F.S., Sayed, A.H.: Diffusion LMS strategies for distributed estimation. IEEE Transactions on Signal Processing 58(3), 1035–1048 (2010)\n\nBordley, R.F.: The combination of forecasts: A Bayesian approach. Journal of the Operational Research Society 33(2), 171–174 (1982)\n\n- It is not clear the difference among your strategy and the consensus one and the others in the papers above. Please clarify.\n\n- Regarding parallel MCMC chains, other schemes could be considered such as:\n\nR. Craiu et al. Learn from thy neighbor: Parallel-chain and regional adaptive MCMC. Journal of the American\nStatistical Association, 104(448):1454–1466, 2009.\n\nF. Llorente et al, \"Parallel Metropolis-Hastings Coupler\", IEEE Signal Processing Letters, Volume 26, Number 6, Pages 953-957, 2019.\n\nJ.Corander et al. Parallel interacting MCMC for learning of topologies of graphical models. Data Mining and Knowledge Discovery, 17(3):431–456, 2008\n\n- Regarding the Laplace approximations: you still need the knowledge of the MAP  (maxima a-posteriori). You still need to compute/reach these maxima. Please clarify the computational cost of these previous steps for finding these maxima.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "MwJP-xQX-jG", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces a new perspective on federated learning through the lens of posterior inference. The paper designs a computation- and communication-efficient posterior inference algorithm—federated posterior averaging (FEDPA), which generalizes FedAvg. FEDPA is compared with the strong baselines in Reddi et al. (2020) on realistic FL benchmarks, which achieves state-of-the-art results with respect to multiple metrics of interest. \n\nOverall, the paper is well written and easy to follow. I tend to accept the paper. The detailed comments follow. \n\nPros: \n1. Viewing federated learning through the lens of posterior inference is new. The motivating example in Figure 1 gives a nice explanation why FedPA may outperform FedAvg. \n2. The proposed FedPA algorithm enjoys provable performance guarantee.  \n3. The simulations are extensive, which demonstrate the clear advantage of FedPA over FedAvg on multiple benchmark tasks. \n\nCons:\n1. The algorithm itself is not well explained in the first eight pages. Consider moving some key equations from appendix to main text to explain clientMCMC. \n2. The convergence result is weaker than those on optimization for federated learning. Specifically, no finite-time analysis is provided, and the dependence of the number of local update K on the performance is not clear in the current form. \n3. The new challenge relative to distributed MCMC literature is not well explained. Communication? Privacy? or else. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "8dBWt3UAeMD", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a new method of generating local (client) updates in Federated Learning (FL), where the clients return an adjusted version of their usual local updates to the server. The authors derive this new local update rigorously from the viewpoint of estimating the posterior distribution of the data (under Gaussianity assumptions). They also provide an efficient method for calculating this new update, and show that it outperforms Federated Averaging on several datasets.\n\nPros:\n\nThe new method (FEDPA) is well motivated and is 'as simple as possible, but not simpler' based on the derivation.\n\nFEDPA has standard Federated Averaging (FEDAVG) as a special case, and suggests a family of new methods based on approximations of the covariance matrix, which likely would exhibit bias-variance-tradeoff-esque behaviour.\n\nThe authors provide a practical way of calculating the required new quantities efficiently.\n\nExperimental results show FEDPA has superior performance compared to FEDAVG, especially in regimes where client compute is high. \n\nCons:\n\nThe big O analysis of the dynamic programming method for computing the local updates is very useful, but it would also be good to have empirical results on the additional cost on the client. It seems like the cost should not be too significant, but evidence of this would be very valuable, since the cost of FEDPA is strictly greater than the cost of FEDAVG. \n\nThe addition of more tuning parameters in FL is never ideal, especially with the knowledge that using too small of a burn in time can lead to arbitrarily bad behaviour. However since the positive effects of FEDPA appear very quickly after the burn in ends, this may be less of a concern in practice. \n\nWould be valuable to include the results of the FEDAVG-1E in the experiments in the main paper, especially since it outperforms both FEDAVG-5E and FEDPA-5E on the StackOverflow LR, which is a surprising and unexpected result. \n\nIn all experiments the FEDAVG locally updated using SGD (likely to maintain the connection and comparison with SGD used in IASG). However for FEDAVG, the optimization procedure CLIENTOPT could be something else, such as adam. It would be valuable to know how this FEDPA compares to FEDAVG when the local optimizer is a more powerful method than SGD, since FEDAVG has the freedom to change the local optimizer. (Of course it is only fair to also empirically compare FEDPA where the CLIENTOPT can be the same method. However since the use of SGD in IASG provides it with certain properties, it is less clear if this substitution can be safely made). \n\nThe authors provide ample justification for their new method and sufficient evidence that it outperforms the existing standard method FEDAVG.  ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers raised a number of concerns which are addressed by the authors. The paper provides an interesting/novel perspective for federated learning (as a posterior inference problem rather than an optimization problem) which can potentially allow for faster and more accurate solutions. 
Decision: Accept (Poster)