Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms

1 INTRODUCTION . 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 efficient 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 ( FEDAVG , 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 , 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 & Konečnỳ , 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 . ∗Most of the work done at Google . Correspondence : maruan.alshedivat.com 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 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 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 efficient algorithm with linear computation and communication costs . Contributions . The main contributions of this paper can be summarized as 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 a computation- and communication-efficient approximate posterior inference algorithm—federated posterior averaging ( FEDPA ) . FEDPA works with stateless clients and its computational complexity and memory footprint are similar to FEDAVG . 3 . We show that FEDAVG with many local steps is 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-the-art results with respect to multiple metrics of interest . 2 RELATED WORK . 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 FEDAVG ( also known as local SGD ) as an optimization algorithm . Multiple works have provided upper bounds on the convergence rate of FEDAVG 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 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 & Konečnỳ ( 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 . 3 A POSTERIOR INFERENCE PERSPECTIVE ON FEDERATED LEARNING . Federated learning is typically formulated as the following optimization problem : min θ∈Rd { F ( θ ) : = N∑ i=1 qifi ( θ ) } , fi ( θ ) : = 1 ni ni∑ j=1 f ( θ ; zij ) , ( 1 ) where the global objective function F ( θ ) is a weighted average of the local objectives fi ( θ ) over N clients ; each client ’ s objective is some loss f ( θ ; z ) computed on the local data Di = { zi1 , . . . , zini } . 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 { qi } are typically set proportional to the sizes of the local datasets { ni } , 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 ) : = − logP ( 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 ≡ D1 ∪ · · · ∪DN ) . 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 . 1 . 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.1 Proposition 1 ( Global Posterior Decomposition ) Under the uniform prior , any global posterior distribution that exists decomposes into a product of local posteriors : P ( θ | D ) ∝ ∏N i=1 P ( θ | Di ) . Proposition 1 suggests that as long as we are able to compute local posterior distributions P ( θ | Di ) and communicate them to the server , we should be able to solve Eq . 1 by multiplicatively aggregating them to find 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 efficient enough for real federated settings . The example below illustrates how this can be difficult 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 > θ − y ) 2 is quadratic . Then , the client objective becomes : fi ( θ ) = log exp { 1 2 ‖Xiθ − yi‖2 } = log exp { 1 2 ( θ − µi ) > Σ −1 i ( θ − µi ) } + const , ( 2 ) where Xi ∈ Rni×d is the design matrix , yi ∈ Rni is the response vector , Σ−1i : = X > i Xi and µi : = ( X > i Xi ) −1 X > i yi . 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 ) : µ : = ( N∑ i=1 qiΣ −1 i ) −1 ( N∑ i=1 qiΣ −1 i µi ) . ( 3 ) Concretely , in the case of least squares regression , this suggests that it is sufficient 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 find the global optimum . However , a straightforward application of Eq . 3 would requireO ( d2 ) space andO ( d3 ) 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 ( d2 ) , while standard FL algorithms have communication cost of O ( d ) . 1Note that from the optimization point of view , the global optimum generally can not be represented as any weighted combination of the local optima even in simple 2D settings ( see Fig . 1 , left ) . Algorithm 1 Generalized Federated Optimization input initial θ , CLIENTUPDATE , SERVERUPDATE 1 : for each round t = 1 , . . . , T do 2 : Sample a subset S of clients 3 : communicate θ to all i ∈ S // server→ clients 4 : for each client i ∈ S in parallel do 5 : ∆ti , qi ← CLIENTUPDATE ( θ ) 6 : end for 7 : communicate { ∆ti , qi } i∈S // server← clients 8 : ∆t ← 1|S| ∑ i∈S qi∆ t i // aggregate updates 9 : θ ← SERVERUPDATE ( θ , ∆t ) 10 : end for output final θ Algorithm 2 Client Update ( FEDAVG ) input initial θ0 , loss fi ( θ ) , optimizer CLIENTOPT 1 : for k = 1 , . . . , K do 2 : θk ← CLIENTOPT ( θk−1 , ∇̂fi ( θk−1 ) ) 3 : end for output ∆ : = θ0 − θk , client weight qi Algorithm 3 Client Update ( FEDPA ) input initial θ0 , loss fi ( θ ) , sampler CLIENTMCMC 1 : for k = 1 , . . . , K do 2 : θk ∼ CLIENTMCMC ( θk−1 , fi ) 3 : end for output ∆ : = Σ̂ −1 ( θ0 − µ̂ ) , 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 non-Gaussian and closed form expressions for global posterior modes may not exist.2 In such cases , we 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 a computation- and communication-efficient 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 .

Summary: 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.