Federated Generalized Bayesian Learning via  Distributed Stein Variational Gradient Descent

1 INTRODUCTION . Federated learning refers to the collaborative training of a machine learning model across agents with distinct data sets , and it applies at different scales , from industrial data silos to mobile devices ( Kairouz et al. , 2019 ) . While some common challenges exist , such as the general statistical heterogeneity – “ non-iidnes ” – of the distributed data sets , each setting also brings its own distinct problems . In this paper , we are specifically interested in a small-scale federated learning setting consisting of mobile or embedded devices , each having a limited data set and running a small-sized model due to their constrained memory . As an example , consider the deployment of health monitors based on data from smart-watch ECG data . In this context , we argue that it is essential to tackle the following challenges , which are largely not addressed by existing solutions : • Trustworthiness : In applications such as personal health assistants , the learning agents ’ recommendations need to be reliable and trustworthy , e.g. , to decide when to contact a doctor in case of a possible emergency ; • Number of communication rounds : When models are small , the payload per communication round may not be the main contributor to the overall latency of the training process . In contrast , accommodating many communication rounds requiring arbitrating channel access among multiple devices may yield slow wall-clock time convergence ( Lin et al. , 2020 ) . Most existing federated learning algorithms , such as Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) , are based on frequentist principles , relying on the identification of a single model parameter vector . Frequentist learning is known to be unable to capture epistemic uncertainty , yielding overconfident decisions ( Guo et al. , 2017 ) . Furthermore , the focus of most existing works is on reducing the load per-communication round via compression , rather than decreasing the number of rounds by providing more informative updates at each round ( Kairouz et al. , 2019 ) . This paper introduces a trustworthy solution that is able to reduce the number of communication rounds via a non-parametric variational inference-based implementation of federated Bayesian learning . Federated Bayesian learning has the general aim of computing the global posterior distribution in the model parameter space . Existing decentralized , or federated , Bayesian learning protocols are either based on Variational Inference ( VI ) ( Angelino et al. , 2016 ; Neiswanger et al. , 2015 ; Broderick et al. , 2013 ; Corinzia & Buhmann , 2019b ) or Monte Carlo ( MC ) sampling ( Ahn et al. , 2014 ; Mesquita et al. , 2020 ; Wei & Conlon , 2019 ) . State-of-the-art methods in either category include Partitioned Variational Inference ( PVI ) , which has been recently introduced as a unifying distributed VI framework that relies on the optimization over parametric posteriors ; and Distributed Stochastic Gradient Langevin Dynamics ( DSGLD ) , which is an MC sampling technique that maintains a number of Markov chains updated via local Stochastic Gradient Descent ( SGD ) with the addition of Gaussian noise ( Ahn et al. , 2014 ; Welling & Teh , 2011 ) . The performance of VI-based protocols is generally limited by the bias entailed by the variational approximation , while MC sampling is slow and suffers from the difficulty of assessing convergence ( Angelino et al. , 2016 ) . Stein Variational Gradient Descent ( SVGD ) has been introduced in ( Liu & Wang , 2016 ) as a nonparametric Bayesian framework that approximates a target posterior distribution via non-random and interacting particles . SVGD inherits the flexibility of non-parametric Bayesian inference methods , while improving the convergence speed of MC sampling ( Liu & Wang , 2016 ) . By controlling the number of particles , SVGD can provide flexible performance in terms of bias , convergence speed , and per-iteration complexity . This paper introduces a novel non-parametric distributed learning algorithm , termed Distributed Stein Variational Gradient Descent ( DSVGD ) , that transfers the mentioned benefits of SVGD to federated learning . As illustrated in Fig . 1 , DSVGD targets a generalized Bayesian learning formulation , with arbitrary loss functions ( Knoblauch et al. , 2019 ) ; and maintains a number of non-random and interacting particles at a central server to represent the current iterate of the global posterior . At each iteration , the particles are downloaded and updated by one of the agents by minimizing a local free energy functional before being uploaded to the server . DSVGD is shown to enable ( i ) a trade-off between per-iteration communication load and number of communication rounds by varying the number of particles ; while ( ii ) being able to make trustworthy decisions through Bayesian inference . 2 SYSTEM SET-UP . We consider the federated learning set-up in Fig . 1 , where each agent k = 1 , . . . , K has a distinct local dataset with associated training loss Lk ( θ ) for model parameter θ . The agents communicate through a central node with the goal of computing the global posterior distribution q ( θ ) over the shared model parameter θ ∈ Rd for some prior distribution p0 ( θ ) ( Angelino et al. , 2016 ) . Specifically , following the generalized Bayesian learning framework ( Knoblauch et al. , 2019 ) , the agents aim at obtaining the distribution q ( θ ) that minimizes the global free energy min q ( θ ) { F ( q ( θ ) ) = K∑ k=1 Eθ∼q ( θ ) [ Lk ( θ ) ] + αD ( q ( θ ) ||p0 ( θ ) ) } , ( 1 ) where α > 0 is a temperature parameter . The ( generalized , or Gibbs ) global posterior qopt ( θ ) solving problem ( 1 ) must strike a balance between minimizing the sum loss function ( first term in F ( q ) ) and the model complexity defined by the divergence from a reference prior ( second term in F ( q ) ) . It is given as qopt ( θ ) = 1 Z · q̃opt ( θ ) , with q̃opt ( θ ) = p0 ( θ ) exp ( − 1 α K∑ k=1 Lk ( θ ) ) , ( 2 ) where we denoted as Z the normalization constant . It is useful to note that the global free energy can also be written as the scaled KL F ( q ( θ ) ) = αD ( q ( θ ) ||q̃opt ( θ ) ) . The main challenge in computing the optimal posterior qopt ( θ ) in a distributed manner is that each agent k is only aware of its local loss Lk ( θ ) . By exchanging information through the server , the K agents wish to obtain an estimate of the global posterior ( 2 ) without disclosing their local datasets neither to the server nor to the other agents . In this paper , we introduce a novel non-parametric distributed generalized Bayesian learning framework that addresses this challenge by integrating Distributed VI ( DVI ) and SVGD ( Liu & Wang , 2016 ) . 3 DISTRIBUTED VARIATIONAL INFERENCE . In this section , we describe a general Expectation Propagation ( EP ) -based framework ( Vehtari et al. , 2020 ) , which we term as DVI , that aims at computing the global posterior in a federated fashion ( Bui et al. , 2018 ; Corinzia & Buhmann , 2019b ) . DVI starts from the observation that the posterior ( 2 ) factorizes as the product q ( θ ) = p0 ( θ ) K∏ k=1 tk ( θ ) , ( 3 ) where the term tk ( · ) is given by the scaled local likelihood exp ( α−1Lk ( θ ) ) /Z . Since the normalization constant Z depends on all data sets , the true scaled local likelihood tk ( · ) can not be directly computed at agent k. The idea of DVI is to iteratively update approximate likelihood factors tk ( θ ) for k = 1 , ... , K by means of local optimization steps at the agents and communication through the server , with the aim of minimizing the global free energy ( 1 ) over distribution ( 3 ) . We give here the standard implementation of DVI in which a single agent is schedule at each time , although parallel implementations are possible and discussed below . Accordingly , at each communication round i = 1 , 2 , ... , the server maintains the current iterate q ( i−1 ) ( θ ) of the global posterior , and schedules an agent k ∈ { 1 , 2 , . . . , K } , which proceeds as follows : 1 . Agent k downloads the current global variational posterior distribution q ( i−1 ) ( θ ) from the server ( see Fig . 1 ( a ) , step 1 ) ; 2 . Agent k updates the global posterior by minimizing the local free energy F ( i ) k ( q ( θ ) ) ( see Fig . 1 ( a ) , step 2 ) q ( i ) ( θ ) = argmin q ( θ ) { F ( i ) k ( q ( θ ) ) = Eθ∼q ( θ ) [ Lk ( θ ) ] + αD ( q ( θ ) ||p̂ ( i ) k ( θ ) ) } , ( 4 ) where we have defined the ( unnormalized ) cavity distribution p̂ ( i ) k ( θ ) as p̂ ( i ) k ( θ ) = q ( i−1 ) ( θ ) t ( i−1 ) k ( θ ) . ( 5 ) The cavity distribution p̂ ( i ) k ( θ ) , which removes the contribution of the current approximate likelihood of agent k from the current global posterior iterate , serves as a prior for the update in ( 4 ) . In a manner similar to ( 2 ) , the local free energy is minimized by the tilted distribution p ( i ) k ( θ ) ∝ p̃ ( i ) k ( θ ) with p̃ ( i ) k ( θ ) = p̂ ( i ) k ( θ ) exp ( − 1 α Lk ( θ ) ) ; ( 6 ) 3 . Agent k sends the updated posterior q ( i ) ( · ) = p ( i ) k ( · ) to the server ( see Fig . 1 ( a ) , step 3 ) , and updates its approximate likelihood accordingly as t ( i ) k ( θ ) = q ( i ) ( θ ) q ( i−1 ) ( θ ) t ( i−1 ) k ( θ ) ; ( 7 ) Finally , non-scheduled agents k′ 6= k set t ( i ) k′ ( θ ) = t ( i−1 ) k′ ( θ ) , and the server sets the next iterate as q ( i ) ( θ ) . We have the following key property of DVI . Theorem 1 . The global posterior qopt ( θ ) in ( 2 ) is the unique fixed point of the DVI algorithm . The fixed-point property in Theorem 1 can be verified directly by setting q ( i−1 ) ( θ ) = qopt ( θ ) and t ( i−1 ) k ( θ ) = exp ( α −1Lk ( θ ) ) /Z and by observing that this leads to the fixed point condition q ( i ) ( θ ) = q ( i−1 ) ( θ ) = qopt ( θ ) . The proof is provided in Sec . A.6 . Importantly , this property is not tied to the sequential implementation detailed above , and it applies also if multiple devices are scheduled in parallel , as long as one sets the next iterate as q ( i ) ( θ ) = p0 ( θ ) ∏ k∈K ( i ) t ( i ) k ( θ ) ∏ k′ 6∈K ( i ) t ( i ) k′ ( θ ) , where K ( i ) denotes the set of scheduled agents at communication round i and we have t ( i ) k′ ( θ ) = t ( i−1 ) k′ ( θ ) and t ( i ) k ( θ ) updated following ( 7 ) .

Summary: This paper proposes a Bayesian optimization algorithm in the context of federated learning. The whole framework is built on top of generalized Bayesian learning. To overcome the locality of clients' distributions, the authors propose their solution as an integration of Partitioned Variational Inference (PVI) and Stein Variational Gradient Descent (SVGD). Numerical experiments have been conducted on a synthetic dataset and some standard benchmark datasets, and evaluated on both regression and classification tasks.