F EDERATED L EARNING VIA P OSTERIOR A VERAGING : A N EW P ERSPECTIVE AND P RACTICAL A LGORITHMS
A BSTRACT
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 communicationefﬁcient approximate posterior inference algorithm— federated posterior averaging (F ED PA).
Our algorithm uses MCMC for approximate inference of local posteriors on the clients and efﬁciently communicates their statistics to the server, where the latter uses them to reﬁne a global estimate of the posterior mode.
Finally, we show that F ED PA generalizes federated averaging (F ED A VG ), can similarly beneﬁt from adaptive optimizers, and yields state-of-the-art results on four realistic and challenging benchmarks, converging faster, to better optima.
I NTRODUCTION
Federated learning (FL) is a 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 to the server which aggregates these updates into a global model (McMahan et al., 2017).
While this paradigm enables efﬁcient distributed learning from data stored on millions of remote devices (Hard et 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 as a distributed optimization problem with a few distinctive 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 is federated averaging (F ED A VG ,
McMahan et al., 2017), which proceeds in rounds of communication between the server and a random subset of clients, synchronously updating the server model after each round (Bonawitz et al., 2019).
By allowing the clients perform multiple local SGD steps (or epochs) at each round, F ED A VG 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 F ED A VG stagnate at inferior optima.
As a result, while slow during initial training,
MB-SGD ends up dominating F ED A VG at convergence (see example in Fig. 1).
This has been observed in multiple empirical studies ( e.g. , Charles & Koneˇcn`y, 2020), and recently was shown theoretically (Woodworth et 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 a single FL round.
Thus, the state of a stateful algorithm is never used.
[Figure]
Is it possible to design FL algorithms that exhibit both fast training and consistent convergence with stateless clients?
In this work, we answer this question afﬁrmatively, 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, beneﬁt 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 infeasible 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 an efﬁcient algorithm with linear computation and communication costs.
Contributions.
The main contributions of this paper can be summarized as follows:
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.
With this perspective, we design a computation- and communication-efﬁcient approximate posterior inference algorithm— federated posterior averaging (F ED PA).
F ED PA works with stateless clients and its computational complexity and memory footprint are similar to F ED A VG .
We show that F ED A VG with many local steps is in fact a special case of F ED PA that estimates local posterior covariances with identities.
These biased estimates are the source of inconsistent updates and explain why F ED A VG has suboptimal convergence even in simple quadratic settings.
Finally, we compare F ED PA with strong baselines on realistic FL benchmarks introduced by Reddi et al.
and achieve state-of-the-art results with respect to multiple metrics of interest.
R ELATED W ORK
Federated optimization.
Starting with the seminal paper by McMahan et al.
(2017), a lot of recent effort in federated learning has focused on understanding of F ED A VG (also known as local SGD) as an optimization algorithm.
Multiple works have provided upper bounds on the convergence rate of
F ED A VG in the homogeneous 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; Hsieh et al., 2019; Li et al., 2019; Wang et al., 2020; Woodworth et al., 2020a).
Reddi et al.
(2020) reformulated F ED A VG in a way that enabled adaptive optimization and derived corresponding convergence rates, noting that F ED A VG requires careful tuning of learning rate schedules in order to converge to the desired optimum, which was further analyzed by Charles &
Koneˇcn`y (2020).
To the best of our knowledge, our work is perhaps the ﬁrst 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-speciﬁc constraints.
A P OSTERIOR I NFERENCE P ERSPECTIVE ON F EDERATED L EARNING
Federated learning is typically formulated as the following optimization problem:
[Equation]
where the global objective function F ( θ ) is a weighted average of the local objectives
f i ( θ ) over N clients; each client’s objective is some loss f
( θ ; z ) computed on the local data D i
=
{ z i 1 , . . .
, z in i } .
In real-world cross-device applications, the total number of clients N can be extremely large, and hence optimization of F ( θ ) is done over multiple rounds with only a small subset of M clients participating in each round.
The weights { q i } are typically set proportional to the sizes of the local datasets { n i } , which makes F ( θ ) coincide with the training objective of the centralized setting.
Typically, f ( θ ; z ) is negative log likelihood of z under some probabilistic model parametrized by θ , i.e. , f ( θ ; z ) :
= − log P ( z | θ ) .
For example, least squares loss corresponds to likelihood under a
Gaussian model, cross entropy loss corresponds to likelihood under a categorical model, etc.
(Murphy, 2012).
Thus, Eq.
1 corresponds to maximum likelihood estimation (MLE) of the model parameters θ .
An alternative (Bayesian) approach to maximum likelihood estimation is posterior inference or estimation of the posterior distribution of the parameters given all the data: P ( θ | D ≡ D 1 ∪ · · · ∪ D N ) .
The posterior is proportional to the product of the likelihood and a prior, P ( θ | D ) ∝ P ( D | θ ) P ( θ ) , and, if the prior is uninformative (uniform over all θ ), the modes of the global posterior coincide with
MLE solutions or optima of F ( θ ) in Eq.
While this simple observation establishes an equivalence between the inference of the posterior mode and optimization, the 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 ( θ | D ) ∝ (cid:81) Ni =1 P ( θ | D i ) .
Proposition 1 suggests that as long as we are able to compute local posterior distributions P ( θ | D i ) and communicate them to the server, we should be able to solve Eq.
1 by multiplicatively aggregating them to ﬁnd the mode of the global posterior P ( θ | D ) on the server.
Note that posterior inference via multiplicative averaging has been successfully used to scale Monte Carlo methods to 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 efﬁcient enough for real federated settings.
The example below illustrates how this can be difﬁcult even for a simple model and loss function.
Federated least squares.
Consider federated least squares regression with a linear model, where z :=
( x , y ) and the loss f
( θ ; x , y ) :
= 12 ( x (cid:62) θ − y ) 2 is quadratic.
Then, the client objective becomes:
[Equation]
where X i ∈ R n i × d is the design matrix,
y i ∈ R n i is the response vector, Σ − 1 i
[Equation]
i X i and µ i := (cid:0) X (cid:62)
i X i (cid:1) − 1 X (cid:62) i
y i .
Note that the expression in Eq.
2 is the log likelihood for a multivariate
Gaussian distribution with mean µ i and covariance Σ i .
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):
[Equation]
Concretely, in the case of least squares regression, this suggests that it is sufﬁcient for clients to infer the means { µ i } and inverse covariances { Σ − 1 i } of their local posteriors and communicate that information to server for the latter to be able to ﬁnd the global optimum.
However, a straightforward application of Eq.
3 would require O ( d 2 ) space and O ( d 3 ) 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 O ( d 2 ) , while standard FL algorithms have communication cost of O ( d ) .
input
3: communicate do 5:
7: communicate t
9: θ
10: end input initial θ 0 , loss f i θ ) optimizer C LIENT O PT 1: for k =
K do 2: θ k ← C LIENT
[Table] [Table]
k −
[Table] [Table]
f
[Table]
) 3: end for output i input
LIENT MCMC
[Table] [Table]
do
[Table]
output
[Table]
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 non-Gaussian and closed form expressions for global posterior modes may not exist.
In such cases, we propose to use the Laplace approximation for local and global posteriors , i.e. , approximate them with the best-ﬁtting Gaussians.
While imperfect, this approximation will allow us to compute the (approximate) global posterior mode in a computation- and communication-efﬁcient manner using the following three steps: (i) infer approximate local means { ˆ µ i } and covariances { ˆ Σ i } , (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 be 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 of model parameters.
F EDERATED P OSTERIOR A VERAGING : A P RACTICAL A LGORITHM
Federated averaging (F ED A VG , McMahan et al., 2017) solves the problem from Eq.
1 over T rounds by interacting with M random clients at each round in the following way: (i) broadcasting the current model parameters θ to the clients, (ii) running SGD for K steps on each client, and (iii) updating the global model parameters by collecting and averaging the ﬁnal 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.
F ED A VG is efﬁcient in that it requires only O ( d ) computation on both the clients and the server, and
O ( d ) communication between each client and the server.
To arrive at a similarly efﬁcient algorithm for posterior inference, we focus on the following questions: (a) how to estimate local and global posterior moments efﬁciently?
(b) how to communicate local statistics to the server efﬁciently?
(1) Efﬁcient global posterior inference.
There are two issues with computing an estimate of the global posterior mode µ directly using Eq. 3.
First, it requires computing the inverse of a d × d matrix on the server, which is an O ( d 3 ) operation.
Second, it relies on acquiring local means and inverse covariances, which would require O ( d 2 ) 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 minimizer of a quadratic Q ( θ ) :
= 12 θ (cid:62) A θ − b (cid:62) θ , where A := (cid:80) Ni =1 q i Σ − 1 i and b := (cid:80) Ni =1 q i Σ − 1 i µ i .
Proposition 2 allows us to obtain a good estimate of µ by running stochastic optimization of the quadratic objective Q ( θ ) on the server.
Note that the gradient of Q ( θ ) has the following form:
[Equation]
which suggests that we can obtain µ by using the same Algorithm 1 as F ED A VG but using different client updates: ∆ i :
= Σ − 1 i ( θ − µ i ) .
Importantly, as long as clients are able to compute ∆ i ’s, this approach will result in O ( d ) communication and O ( d ) server computation cost per round.
input initial θ , loss f i ( θ ) , optimizer C LIENT O PT ( α ) , B : burn-in steps, K : steps per sample, (cid:96) : # samples.
// Burn-in 1: for step t = 1 , . . .
, B do 2: θ
← C LIENT
O PT ( θ , ˆ ∇ f i ( θ )) 3: end for //
Sampling 4: for sample s = 1 , . .
, (cid:96) do 5:
S θ ← ∅ // Initialize iterates 6: for step t = 1 , . .
, K do 7: θ ← C LIENT
O PT ( θ , ˆ ∇ f i ( θ )) 8:
S θ ← S θ ∪ { θ } 9: end for 10: θ s ← A VERAGE ( S θ ) //
Average iterates 11: end for output samples { θ 1 , . .
, θ (cid:96) } (2) Efﬁcient local posterior inference.
To compute ∆ i , each client needs to be able to estimate 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 sampling from local posteriors on the clients, so that these samples can be used to estimate ˆ µ i ’s and ˆ Σ i ’s. Speciﬁcally, we use a variant of SGMCMC 3 with iterate averaging (IASG, Mandt et al., 2017), which involves: (a) running local SGD for some number of steps to mix in the Markov chain, then (b) continued running of
SGD for more steps to periodically produce samples 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 in better estimates of the local moments.
(3) Efﬁcient computation of the deltas.
Even if we can obtain samples { ˆ θ 1 , . .
, ˆ θ (cid:96) } via MCMC and use them to estimate local moments, ˆ µ i and ˆ Σ i , computing ∆ i naïvely would still require inverting a d × d matrix, i.e. , O ( d 3 ) compute and O ( d 2 ) memory.
The good news is that we are able to show that clients can compute ∆ i ’s much more efﬁciently, in O ( d ) time and memory, using a dynamic programming algorithm and appropriate mean and covariance estimators.
Theorem 3 Given (cid:96) approximate posterior samples { ˆ θ 1 , . .
, ˆ θ (cid:96) } , let ˆ µ (cid:96) be the sample mean, ˆ S (cid:96) be the sample covariance, and ˆ Σ (cid:96) :
= ρ (cid:96) I + (1 − ρ (cid:96) )
ˆ S (cid:96) be a shrinkage estimator (Ledoit & Wolf, 2004b) of the covariance with ρ (cid:96) :
= 1 / (1 + ( (cid:96) − 1) ρ ) for some ρ ∈ [0 , + ∞ ) .
Then, for any θ , we can compute ˆ ∆ (cid:96) =
ˆ Σ − 1 (cid:96) ( θ − ˆ µ (cid:96) ) in O ( (cid:96) 2 d ) time and using O ( (cid:96)d ) memory.
Proof [Sketch] We give a constructive proof by designing an efﬁcient algorithm for computing ˆ ∆ (cid:96) .
Our approach is based on two key ideas:
We prove that the speciﬁed shrinkage estimator of the covariance has a recursive decomposition into rank-1 updates, i.e. , ˆ Σ t
= ˆ Σ t − 1 + c t · x (cid:62) t x t , where c t is a constant and x t is some vector.
This allows us to leverage the Sherman-Morrison formula for computing the inverse of ˆ Σ (cid:96) .
Further, we design a dynamic programming algorithm
for computing ˆ ∆ (cid:96) exactly without storing the covariance matrix or its inverse.
Our algorithm is online and allows efﬁcient updates of ˆ ∆ (cid:96) as more posterior samples become available.
See Appendix C for the full proof and derivation of the algorithm.
[Table]
Note that the computational cost of ˆ ∆ (cid:96) consists of two components: (i) the cost of producing (cid:96) approximate local posterior samples using IASG and (ii) the cost of solving a linear system using dynamic programming.
How much of an overhead does it add compared to simply running local
SGD?
It turns out that in practical settings the overhead is almost negligible.
Table 1 shows the time it takes a client to compute the updates based onlocal epochs (100 steps per epoch) using different algorithms (F ED A VG vs. our approach with exact or dynamic programming (DP) matrix inversion) on synthetic linear regressions.
As the dimensionality grows, computational complexity of DP-based estimation of ˆ ∆ (cid:96) becomes nearly identical to F ED A VG , which indicates that the majority of the cost in practice would come from SGD steps rather than our dynamic programming procedure.
The ﬁnal algorithm, discussion, and implications.
Putting all the pieces together, we arrive at the federated posterior averaging (F ED PA) algorithm for approximately computing the mode of the global posterior over multiple communication rounds.
Our algorithm is a variant of generalized federated optimization (Algorithm 1) with a new client update procedure (Algorithm 3).
Importantly, this also implies that F ED A VG can be viewed as posterior inference algorithm that estimates ˆ Σ with an identity and, as a result, obtains biased client deltas ˆ ∆ F ED A VG :
= I ( θ − ˆ µ ) .
In Fig.
1 in the introduction, we demonstrate the differences in behavior between F ED A VG and
F ED PA that stem from the differences in their client updates.
Biased client updates make F ED A VG converge to a suboptimal point; moreover, increasing local computation only pushes the ﬁxed point further away from the global optimum.
On the other hand, F ED PA converges faster and to a better optimum, trading off bias for slightly more variance (becomes visible only closer to convergence).
We see that F ED PA also substantially beneﬁts from more local computation (more local samples).
Since the main difference between F ED A VG and F ED PA is, in fact, the bias-variance trade off in the server gradient estimates (Eq. 4), we can view both methods as biased SGD (Ajalloeian & Stich, 2020) and reason about their convergence rates as well as distances between their ﬁxed points and correct global optima as 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-based approximations on the behavior of F ED PA.
E XPERIMENTS
Using a suite of realistic benchmark tasks introduced by Reddi et al.
(2020), we evaluate F ED PA against several competitive baselines: the best versions of F ED A VG with adaptive optimizers as well as MIME (Karimireddy et al., 2020)—a recently-proposed F ED A VG variant that also works with stateless clients, but uses control-variates and server-level statistics to mitigate convergence issues.
[Table]
5. T HE S ETUP
Datasets and tasks.
The four benchmark tasks are based on the following three datasets (Table 2):
EMNIST (Cohen et al., 2017), CIFAR100 (Krizhevsky et al., 2009), and StackOverﬂow (StackOverﬂow, 2016).
EMNIST (handwritten characters) and CIFAR100 (RGB images) are used for multi-class image classiﬁcation tasks.
StackOverﬂow (text) is used for next-word prediction (also a multi-class classiﬁcation task, historically denoted NWP) and tag prediction (a multi-label classiﬁcation task, historically denoted LR because a logistic regression model is used).
EMNIST was partitioned by authors (Caldas et al., 2018), CIFAR100 was partitioned randomly into 600 clients with a realistic heterogeneous structure (Reddi et al., 2020), and StackOverﬂow was partitioned by its unique users.
All datasets were preprocessed using the code provided by Reddi et al.
Methods and models.
We use a 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 F ED A VG -1E or F ED A VG -ME, which stands for 1 or multiple local epochs performed by clients at each round, respectively.
The number of local epochs in the multi-epoch versions is a hyperparameter.
We use the same framework for federated posterior averaging and refer to it as F ED PA-ME.
As our clients use IASG to produce approximate posterior samples, collecting
[Table] [Table]
[Figure]
a single sample per epoch is optimal (Mandt et al., 2017).
Thus F ED PA-ME uses M samples to estimate client deltas and has the same local and global computational complexity as F ED A VG -ME but with two extra hyperparameters: the number of burn-in rounds and the shrinkage coefﬁcient ρ 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 StackOverﬂow NWP, and multi-label logistic regression on bag-of-words vectors for StackOverﬂow LR (for details see Appendix D).
Hyperparameters.
For hyperparameter tuning, we ﬁrst ran small grid searches for F ED A VG -ME using the best server optimizer and corresponding learning rate grids from Reddi et al.
Then, we used the best F ED A VG -ME conﬁguration and did a small grid search to tune the additional hyperparameters of F ED PA-ME, which turned out not to be very sensitive ( i.e. , many conﬁgurations provided results superior to F ED A VG ).
More hyperparameter details can be found in Appendix D.
Metrics.
Since both speed of learning as well as ﬁnal performance are important quantities for federated learning, we measure: (i) the number of rounds it takes the algorithm to attain a desired level of an evaluation metric and (ii) the best performance attained within a speciﬁed number of rounds.
For EMNIST-62, we measure the number of rounds it takes different methods to achieve 84% and 86% evaluation accuracy 5 , 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 andand 1500 as round number cutoffs.
Finally, for StackOverﬂow, we measure the the number of rounds it takes to the best performance and evaluation accuracy (for the NWP task) and precision, recall at 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 metrics 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 https://github.com/alshedivat/fedpa .
The TFF implementation will be released through https://github.com/google-research/federated .
[Table]
5. R ESULTS ON B ENCHMARK T ASKS
The effects of posterior correction of client deltas.
As we demonstrated in Section 4, F ED PA essentially generalizes F ED A VG and only differs in the computation done on the clients, where we compute client deltas using an estimator of the local posterior inverse covariance matrix, Σ − 1 i , which requires sampling from the posterior.
To be able to use SG-MCMC for local sampling, we ﬁrst run
F ED PA in the burn-in regime (which is identical to F ED A VG ) for a 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 F ED A VG to F ED PA for CIFAR-100 (after 400 burn-in rounds) and
StackOverﬂow LR (after 800 burn-in rounds) is presented on Figs.
2a and 2b, respectively.
During the burn-in phase, evaluation performance is identical for both methods, but once F ED PA starts computing client deltas using local posterior samples, the loss immediately drops and the convergence trajectory changes, indicating that F ED PA is able to avoid stagnation and make progress towards a better optimum.
Similar effects are observed across all other tasks (see Appendix E).
While the improvement of F ED PA over F ED A VG on 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 F ED PA against: tuned F ED A VG with a ﬁxed client learning rate (denoted F ED A VG -1E and F ED A VG -ME), the best variation of adaptive F ED A VG from Reddi et al.
(2020) with exponentially decaying client learning rates (denoted AFO), and MIME of Karimireddy et al.
With more local epochs, we see signiﬁcant improvement in terms of speed of learning: both F ED PA-ME and F ED A VG -ME achieve 84% accuracy on EMNIST-62 in under 100 rounds (similarly, both methods attain 30% on
CIFAR-100 by round 350).
However, more local computation eventually hurts F ED A VG leading to worse optima: on EMNIST-62, F ED A VG -ME is not able to consistently achieve 86% accuracy within 1500 rounds; on CIFAR-100, it takes extra 350 rounds for F ED A VG -ME to get to 40% accuracy.
Finally, federated posterior averaging achieves the best performance on both tasks in terms of evaluation accuracy within the speciﬁed limit on the number of training rounds.
On EMNIST-62 in particular, the ﬁnal performance of F ED PA-ME 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 StackOverﬂow NWP and LR.
Results for StackOverﬂow are presented in Table 3c.
Although not as pronounced as for image datasets, we observe some improvement of F ED PA over
F ED A VG here as 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-F1.
The ﬁrst two metrics have appeared in some prior FL work, while the latter two are the primary evaluation metrics typically used in multi-label classiﬁcation work (Gibaja & Ventura, 2015).
Interestingly, while F ED PA underperforms in terms of precision and recall, it substantially outperforms in terms of micro- and macro-averaged F1, especially the macro-F1.
This indicates that while F ED A VG learns a model that can better predict high-frequency labels, F ED PA learns a model that better captures rare labels (Yang, 1999; Yang & Liu, 1999).
Interestingly, note while F ED PA improves on F1 metrics and has almost the same recall at 5, it’s precision after 1500 rounds is worse than F ED A VG .
A more detailed discussion along with training curves for each evaluation metric are provided in Appendix E.
C ONCLUSION AND F UTURE D IRECTIONS
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 a new algorithm that generalizes federated averaging, is similarly practical and efﬁcient, and yields state-of-the-art results on multiple challenging benchmarks.
While our algorithm required a number of speciﬁc approximation and design choices, we believe that the underlying approach has potential to signiﬁcantly 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 speciﬁc posterior sampling and covariance estimation techniques.
While in the appendix we analyzed the effects of some of these design choices, exploration of: (i) other sampling strategies, (ii) more efﬁcient covariance estimators (Hsieh et al., 2013), (iii) alternatives to MCMC (such as 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 et al., 2015), better understanding of privacy implications of posterior inference in federated settings is an open question. A CKNOWLEDGMENTS
The authors would like to thank Zachary Charles for the invaluable feedback that inﬂuenced the design of the methods and experiments, and Brendan McMahan, Zachary Garrett, Sean Augenstein, Jakub
Koneˇcný, Daniel Ramage, Sanjiv Kumar, Sashank Reddi, Jean-François Kagy for many insightful discussions, and Willie Neiswanger for helpful comments on the early drafts.
A P RELIMINARY A NALYSIS AND A BLATIONS
In Section 4, we derived federated posterior averaging (F ED PA) starting with the global posterior decomposition (Proposition 1, which is exact) and applying the following three approximations:
The Laplace approximation of the local and global posterior distributions.
The shrinkage estimation of the local moments.
Approximate sampling from the local posteriors using MCMC.
We have also observed that F ED A VG is a special case of F ED PA (
from
the algorithmic point of view), since it can be viewed as also using the Laplace approximation for the posteriors, but estimating local covariances ˆ Σ i ’s with identities and local means using the ﬁnal iterates of local SGD.
In this section, we analyze the effects of approximations 2 and 3 on the convergence of F ED PA.
Speciﬁcally, we ﬁrst discuss the convergence rates of F ED A VG and F ED PA as biased stochastic gradient optimization methods (Ajalloeian & Stich, 2020).
We show how the bias and variance of the client deltas behave for F ED A VG and F ED PA as functions of the number samples.
We also analyze the quality of samples produced by IASG (Mandt et al., 2017) and how they depend on the amount of local computation and hyperparameters.
Our analyses are conducted empirically on synthetic data.
A. D ISCUSSION OF THE C ONVERGENCE OF F ED PA VS . F ED A VG
First, observe that if each client is able to perfectly estimate their ∆ i = Σ − 1 i ( θ − µ ) , the problem solved by Algorithm 1 simply becomes an optimization of a quadratic objective using unbiased stochastic gradients, ∆ := 1 M (cid:80) Mi =1 ∆ i .
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 of convergence for
SGD with a O ( t − 1 ) decaying learning rate used on the server is O (1 / √ t ) .
It can be further improved to O (1 /t ) using Polyak momentum (Polyak, 1964) or iterate averaging (Polyak & Juditsky, 1992).
In reality, both F ED A VG and F ED PA produce biased estimates ˆ ∆ F ED A VG and ˆ ∆ F ED PA , respectively.
Thus, we can analyze the problem as SGD with biased stochastic gradient estimates and let ˆ ∆ t := ∇ F ( θ t )
+ b ( θ t )
+ n ( θ t ) where b ( θ t ) and n ( θ t , ξ ) are bias and noise terms.
Following Ajalloeian & Stich (2020), we can further assume that the bias and noise terms are norm-bounded as follows.
Assumption 4 ( ( m, ζ 2 ) -bounded bias)
There exist constants 0 ≤ m < 1 and ζ 2 ≥ 0 such that
[Equation]
Assumption 5 ( ( M, σ 2 ) -bounded noise)
There exist constants 0 ≤ M < 1 and σ 2 ≥ 0 such that
[Equation]
Under these general assumptions, the following convergence result holds.
Theorem 6 (Ajalloeian & Stich (2020), Theorem 2) Let F ( θ ) be L -smooth.
Then SGD with a learning rate α := min (cid:110) 1 L , 1 − m 2 ML , (cid:0) LFσ 2 T (cid:1) 1 / 2 (cid:111) and gradients that satisfy Assumptions 4, 5 achieves the vicinity of a stationary point, E (cid:2) (cid:107)∇ F ( θ ) (cid:107) 2 (cid:3) =
O (cid:16) ε +
[Equation] [Equation]
− m (cid:17) , in T iterations, where
[Equation]
Note that SGD with biased gradients is able to converge to a vicinity of the optimum determined by the bias term ζ 2 / (1 − m ) .
For F ED A VG , since the bias is not countered, this term determines the distance between the stationary point and the true global optimum.
For F ED PA, since ˆ ∆ F ED PA → ∆ 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 F ED PA and
F ED A VG , below we conduct an empirical analysis of the bias and variance of the estimated client deltas on synthetic least squares problems, for which exact deltas can be computed analytically.
[Figure]
[Figure] [Figure]
[Figure]
Quantifying empirically the bias and variance of ˆ ∆ for F ED PA and F ED A VG .
We measure the empirical bias and variance of the client deltas computed by each of the methods on the synthetic least squares linear regression problems generated according to Guyon (2003) using the make_regression 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 { θ 1 , . . .
, θ 10 } and for each of the parameters we computed the exact ∆ i as well as ˆ ∆ F ED A VG ,i and ˆ ∆ F ED PA ,i for different numbers of local steps; for ˆ ∆ F ED PA we also varied the shrinkage hyperparameter.
Using these sample estimates, we further computed the L 2 -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.
From Fig.
3a, we see that as the amount of local computation increases, the bias in F ED A VG delta estimates grows and the variance reduces.
For F ED PA (Fig. 3b), the trends turn out to 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 to the fact that we used the same ﬁxed shrinkage ρ regardless of the number of local steps.
To avoid sharp increases in the variance, ρ must be selected for each number of local steps separately; Fig.
3c demonstrates how the bias and variance depend on the shrinkage hyperparameter for some ﬁxed number of local steps.
A. A NALYSIS OF THE Q UALITY OF IASG- BASED S AMPLING AND C OVARIANCE
The more and better samples we can obtain locally, the lower the bias and variance of the gradients of
Q ( θ ) will be, resulting in faster convergence to a ﬁxed point closer to the global optimum.
For local sampling, we proposed to use a variant of SG-MCMC called Iterate Averaged Stochastic Gradient (IASG) developed by Mandt et al.
(2017), given in Algorithm 4.
The algorithm generates samples by simply averaging every K intermediate iterates produced by a client optimizer (typically, SGD with some a ﬁxed learning rate α ) after skipping the ﬁrst B iterates as a burn-in phase.
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 (cid:96) approximate posterior samples { θ 1 , . .
, θ (cid:96) } , the ESS statistic can be computed as follows:
[Equation]
where weights w j must be proportional to the posterior probabilities, or equivalently to the loss.
Effects of the dimensionality, the number of data points, and IASG parameters on ESS.
The results of our synthetic experiments are presented below in Fig.
The takeaways are as follows:
• More burn-in steps (or epochs) generally improve the quality of samples.
• The larger the number 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 in the 1000 dimensional case); if too small, the samples become correlated.
• Finally, the quality of the samples deteriorates with the increase in the number of dimensions.
[Figure]
B P ROOFS
Proposition 1 (Global Posterior Decomposition)
Under the uniform prior, any global posterior distribution that exists decomposes into a product of local posteriors: P ( θ | D ) ∝ (cid:81) Ni =1 P ( θ | D i ) .
Proof Under the uniform prior, the following equivalence holds for P ( θ | D ) as a function of θ :
[Equation]
The proportionality constant between the left and right hand side in Eq.
8 is (cid:81) Ni =1 P
( D i ) / P ( D ) .
Proposition 2 (Global Posterior Inference)
The global posterior mode µ given in Eq.
3 is the minimizer of a quadratic Q ( θ ) :
= 12 θ (cid:62) A θ − b (cid:62) θ , where A := (cid:80) Ni =1 q i Σ − 1 i and b :
[Equation]
Proof The statement of the proposition (implicitly) assumes that all matrix inverses exist.
Then, the quadratic Q ( θ ) is positive deﬁnite (PD) since A is PD as a convex combination of PD matrices
Σ − 1 i .
Thus, the quadratic has a unique solution θ (cid:63) where the gradient of the objective vanishes:
[Equation]
which implies that µ is the unique minimizer of Q ( θ ) .
C C OMPUTATION OF C LIENT D ELTAS VIA D YNAMIC P ROGRAMMING
In this section, we provide a constructive proof for the following theorem by designing an efﬁcient algorithm for computing ˆ ∆ (cid:96) :
= ˆ Σ − 1 (cid:96) ( θ − ˆ µ (cid:96) ) on the clients in time and memory linear in the number of dimensions d of the parameter vector θ ∈ R d .
Theorem 3 Given (cid:96) approximate posterior samples { ˆ θ 1 , . .
, ˆ θ (cid:96) } , let ˆ µ (cid:96) be the sample mean, ˆ S (cid:96) be the sample covariance, and ˆ Σ (cid:96) :
= ρ (cid:96) I + (1 − ρ (cid:96) )
ˆ S (cid:96) be a shrinkage estimator (Ledoit & Wolf, 2004b) of the covariance with ρ (cid:96) :
= 1 / (1 + ( (cid:96) − 1) ρ ) for some ρ ∈ [0 , + ∞ ) .
Then, for any θ , we can compute ˆ ∆ (cid:96) =
ˆ Σ − 1 (cid:96) ( θ − ˆ µ (cid:96) ) in O ( (cid:96) 2 d ) time and using O ( (cid:96)d ) memory.
The naïve computation of update vectors ( i.e. , where we ﬁrst estimate ˆ µ (cid:96) and ˆ Σ (cid:96) from posterior samples and use them to compute deltas) requires O ( d 2 ) storage and O ( d 3 ) compute on the clients and is both computationally and memory intractable.
We derive an algorithm that, given (cid:96) posterior samples, allows us to compute ˆ ∆ (cid:96) using only O ( (cid:96)d ) memory and O ( (cid:96)
2 d ) compute.
The algorithm makes use of the following two components:
The shrinkage estimator of the covariance (Ledoit & Wolf, 2004b), which is known to be well-conditioned even in high-dimensional settings ( i
e. , when the number of samples is smaller than the number of dimensions) and is widely used in econometrics (Ledoit & Wolf, 2004a) and computational biology (Schäfer & Strimmer, 2005).
Incremental computation of ˆ Σ − 1 (cid:96) ( θ (cid:96) − ˆ µ (cid:96) ) that exploits the fact that each new posterior sample only adds a rank-1 component to ˆ Σ (cid:96) and applies the Sherman-Morrison formula to derive a dynamic program for updating ˆ ∆ (cid:96) .
Notation.
For the sake of this discussion, we denote θ ( i.e. , the server state broadcasted to the clients at round t ) as x 0 , drop the client index i , denote posterior samples as x j , sample mean as ¯ x (cid:96) :
= 1 (cid:96) (cid:80) (
cid:96)j =1 x j , and sample covariance as ˆ S (cid:96) :
= 1 (cid:96) − 1 (cid:80) (cid:96)j =1 ( x j − ¯ x (cid:96) )
( x j − ¯ x (cid:96) ) (cid:62) .
C. T HE S HRINKAGE E STIMATOR OF THE C OVARIANCE
Ledoit & Wolf (2004b) proposed to estimate a high-dimensional covariance matrix using a convex combination of identity matrix and sample covariance (known as the LW or shrinkage estimator):
[Equation]
where ρ (cid:96) is a scalar parameter that controls the bias-variance tradeoff of the estimator.
As an aside, while ρ (cid:96) can be arbitrary and the optimal ρ (cid:96) requires knowing the true covariance Σ , there are near-optimal ways to estimate ˆ ρ (cid:96) from the samples (Chen et al., 2010), which we discuss at the end of this section.
In this section, we focus on deriving an expression for ρ t as a function of t = 1 , . . .
, (cid:96) that ensures that the difference between ˆ Σ t and ˆ Σ t − 1 is a rank-1 matrix (this is not the case for arbitrary ρ ’s).
Derivation of a shrinkage estimator that admits rank-1 updates.
Consider the following matrix:
[Equation]
(11) where β t is a scalar function of t = 1 , 2 , . . .
, (cid:96) .
We would like to ﬁnd β t such that ˜ Σ t = ˜ Σ t − 1 + γ t U t , where U t is a rank-1 matrix, i.e. , the following equality should hold:
[Equation]
To determine the functional form of β t , we need recurrent relationships for ¯ x t and ˆ S t .
For the former, note that the following relationship holds for two consecutive estimates of the sample mean, ¯ x t −and ¯ x t :
[Equation]
This allows us to expand ˆ S t as follows:
[Equation]
Thus
[Equation]
we have the following recurrent relationship between ˆ S t and ˆ S t − 1 :
[Equation]
Now
[Equation]
we can plug (15) into (12) and obtain the following equation:
[Equation]
which implies that U t :
= ( x t − ¯ x t − 1 )( x t − ¯ x t − 1 ) (cid:62) , γ t := β t /t
[Equation]
and the following telescoping expressions for β t :
[Equation]
where we set β 2 ≡ ρ ∈ [0 , + ∞ ) to be a constant.
Thus, if we deﬁne ˜ Σ t :=
I + ρ ( t − 1)ˆ S t , then the following recurrent relationships will hold:
[Equation]
Finally, we can obtain a shrinkage estimator of the covariance from ˜ Σ n by normalizing coefﬁcients:
[Equation]
Note that ˆ Σ 1 ≡ I and ˆ Σ t → S t as t → ∞ .
C. C OMPUTING D ELTAS USING S HERMAN -M ORRISON AND D YNAMIC P ROGRAMMING
Since ˆ Σ (cid:96) is proportional to ˜ Σ (cid:96) and the latter satisﬁes recurrent rank-1 updates given in Eq.
18, denoting u (cid:96) :
= x (cid:96) − ¯ x (cid:96) − 1 , we can express ˆ Σ − 1 (cid:96) =
˜ Σ − 1 (cid:96) /ρ
(cid:96) using the Sherman-Morrison formula:
[Equation]
Note that we would like to estimate ˆ ∆ (cid:96) :
= ˆ Σ − 1 (cid:96) ( x 0 − ¯ x (cid:96) ) , which can be done without computing or storing any matrices if we know ˜ Σ − 1 (cid:96) − 1 u (cid:96) and ˜ Σ − 1 (cid:96) − 1 ( x 0 − ¯ x (cid:96) ) .
Denoting ˜ ∆ t := ˜ Σ − 1 t ( x 0 − ¯ x t ) , and knowing that x 0 − ¯ x (cid:96) =
( x 0 − ¯ x (cid:96) − 1 ) − u (cid:96) /(cid:96)
(which follows from Eq. 13), we can compute ˆ ∆ (cid:96) using the following recurrence:
[Equation]
// initial conditions (21)
[Equation]
// recurrence for u t and v t
[Equation]
[Equation] [Equation]
recurrence for ˜ ∆
[Equation]
(23) ˆ ∆
[Equation] [Equation]
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 t − 1 ,t which needs a matrix-vector product.
To express the latter one in the form of vector-vector products, we need another 2-index recurrence on v i,j := ˜ Σ − 1 i u j :
[Equation] [Equation]
initial conditions
[Equation] [Equation]
ﬁnal expression for v t
[Equation]
[Equation] [Equation]
Now, equipped with these two recurrences, given a stream of samples x 1 , x 2 , . . .
, x t , . .
, we compute ˆ ∆ t for t ≥ 2 based on x t , { u k } t − 1 k =1 , { v k − 2 ,k − 1 } t − 1 k
=1 and ˆ ∆ t − 1 using the following two steps:
Compute u t and v t − 1 ,t using the second recurrence.
Compute ˆ ∆ t from u t , v t − 1 ,t , and ˆ ∆ t − 1 using the ﬁrst recurrence.
For each new sample in the sequence, we repeat the two steps to obtain the updated ˆ ∆ t estimate, until we have processed all (cid:96) samples.
Note that the ﬁrst step requires O ( t ) vector-vector multiplies, i.e. ,
O ( td ) compute, and O ( d ) memory, and the second step a O (1) number of vector-vector multiplies.
As a result, the computational complexity of estimating ˆ ∆ (cid:96) is O ( (cid:96) 2 d ) and the storage needed for the dynamic programming state represented by a tuple (cid:16) { u k } t − 1 k =1 , { v k − 2 ,k − 1 } t − 1 k
=
1 , ˆ ∆ t − 1 (cid:17) is O ( (cid:96)d ) .
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 ˆ ∆ 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 1 posterior sample, the returned delta will be identical to F ED A VG .
By spending more compute on the clients (and a bit of extra memory), with each additional posterior sample x t , we have ˆ ∆ t −→ t →∞ Σ − 1 ( x 0 − µ ) .
Optimal selection of ρ .
Note that to be able to run the above described algorithm in an online fashion, we have to select and commit to a ρ before seeing any samples.
Alternatively, if the online and any-time properties of the algorithm are unnecessary, we can ﬁrst obtain (cid:96) posterior samples { x k } (cid:96)k =1 , then infer a near-optimal ˆ ρ (cid:63) from these samples— e.g. , using the Rao-Blackwellized version of the
LW estimator (RBLW) or the oracle approximating shrinkage (OAS), both proposed and analyzed by Chen et al.
(2010)—and then use the inferred ˆ ρ (cid:63) to compute the corresponding delta using our dynamic programming algorithm.
D D ETAILS ON THE E XPERIMENTAL S ETUP
In this part, we provide additional details on our experimental setup, including a more detailed description of the datasets and tasks, models, methods, and hyperparameters.
D. D ATASETS , T ASKS , AND M ODELS
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.
EMNIST-62.
The dataset is comprised of 28 × 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 of 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 is a 2-layer convolutional network with 3 × 3 kernel, max pooling, and dropout, followed by a 128-unit fully connected layer.
The model was adopted from the TensorFlow
Federated library: https://bit.ly/3l41LKv .
CIFAR-100.
The federated version of CIFAR-100 was introduced by Reddi et al.
The training set of the dataset is partitioned among 500 clients, 100 data points per client.
The partitioning was created using a two-step latent Dirichlet allocation (LDA) over to “coarse” to “ﬁne” labels which created a label distribution resembling a more realistic federated setting.
For the model, also following
Reddi et al.
(2020), we used a modiﬁed ResNet-18 with group normalization layer instead of batch normalization, as suggested by Hsieh et al.
The model was adopted from the TensorFlow
Federated library: https://bit.ly/33jMv6g .
[Table] [Table]
StackOverﬂow.
The dataset consists of text (questions and answers) asked and answered by the total of 342,477 unique users, collected from https://stackoverflow.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 classiﬁcation problem with 10,000 classes).
The tags are similarly restricted to 500 most frequent ones ( i.e. , the LR task becomes a multi-label classiﬁcation proble with 500 labels).
For tag prediction, we use a simple linear regression model where each question or answer are represented by a normalized bag-of-words vector.
The model was adopted from the TensorFlow
Federated library: https://bit.ly/2EXjAeY .
For the NWP task, we restrict each client to the ﬁrst 128 sentences in their dataset, perform padding and truncation to ensure that sentences have 20 words, and then represent each sentence as a sequence of indices corresponding to 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 not on the OOV,
BOS, and EOS symbols.
We use an RNN model with 96-dimensional word embeddings (trained from scratch), 670-dimensional LSTM layer, followed by a fully connected output softmax layer.
The model was adopted from the TensorFlow Federated library: https://bit.ly/2SoSi3X .
D. M ETHODS
As mentioned in the main text, we used F ED A VG with adaptive server optimizers with 1 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 F ED A VG , the names of the methods as given by Reddi et al.
(2020) should be F ED A VG M for EMNIST-62 and CIFAR-100, F ED A DAM for StackOverﬂow NWP and
F ED A DAGRAD for StackOverﬂow LR.
Another difference between our baselines and Reddi et al.
(2020) is that we ran SGD with momentum on the clients for EMNIST-62, CIFAR-100, and
StackOverﬂow LR, as that improved performance of the methods with multiple epochs per client.
Our F ED PA methods used the same conﬁgurations as F ED A VG baselines; moreover, F ED PA and
F ED A VG were identical (algorithmically) during the burn-in phase and only different in the client-side computation during the sampling phase of F ED PA.
[Table]
D. H YPERPARAMETERS AND G RIDS
All hyperparameter grids are given in Table 5.
The best server and client learning rates were selected based on the F ED A VG performance and used for F ED PA.
The best selected hyperparameters are given in Table 6.
E A DDITIONAL E XPERIMENTAL R ESULTS
We provide additional experimental results.
As mentioned in the main text, the 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 a ﬁxed number of rounds.
A much fuller picture is given by the learning curves of each method.
Therefore, we plot evaluation losses, accuracies, and metrics of interest over the course of training.
On the plots, individual values at each round are indicated with × -markers and the 10-round running average with a line of the corresponding color.
EMNIST-62.
Learning curves for F ED A VG and F ED PA on EMNIST-62 are given in Fig.
Fig.
5a shows the best F ED A VG -1E, F ED A VG -5E, and F ED PA-5E models and Fig.
5b shows the best
F ED A VG -20E, and F ED PA-20E. Apart from the fact that multi-epoch versions converge signiﬁcantly faster than the 1-epoch F ED A VG -1E, note that the effect of bias reduction when switching from the burn-in to sampling in F ED PA becomes much more pronounced in the 20-epoch version.
CIFAR-100 and StackOverﬂow.
Learning curves for various models on CIFAR-100 and StackOverﬂow tasks are presented in Figs.
6 and 7.
The takeaways for CIFAR-100 and StackOverﬂow NWP are essentially the same as for EMNIST-62—much faster convergence with the increased number of local epochs and visually noticeable improvement in losses and accuracies due to sampling-based bias correction in client deltas after the burn-in phase is over.
Interestingly, we see that on StackOverﬂow LR task
F ED A VG -1E clearly dominates multi-epoch methods in terms of the loss and recall at 5, losing in precision and macro-F1.
Even more puzzling is the signiﬁcant drop in the average precision of F ED PA-ME after the switching to sampling, while at the same time a jump in recall and
F1 metrics.
This indicates that the global model moves to a different ﬁxed point where it over-predicts positive labels ( i.e. , less precise) but also less likely to miss rare labels ( i.e. , higher recall on rare labels, and as a result a jump in macro-F1).
The reason why this happens, however, is unclear.
[Figure] [Figure]
[Figure]
ED A VG -1E, F ED A VG -5E, and F ED PA-5E.
[Figure] [Figure]