Federated Generalized Bayesian Learning via  Distributed Stein Variational Gradient Descent
Authors: Rahif Kassab
Osvaldo Simeone
Venue:  (2021)
Keywords: Federated Learning
Distributed Variational Inference

Abstract: This paper introduces Distributed Stein Variational Gradient Descent (DSVGD), a non-parametric generalized Bayesian inference framework for federated learning. DSVGD maintains a number of non-random and interacting particles at a central server to represent the current iterate of the model global posterior. The particles are iteratively downloaded and updated by one of the agents with the end goal of minimizing the global free energy. By varying the number of particles, DSVGD enables a flexible trade-off between per-iteration communication load and number of communication rounds. DSVGD is shown to compare favorably to benchmark frequentist and Bayesian federated learning strategies, also scheduling a single device per iteration, in terms of accuracy and scalability with respect to the number of agents, while also providing well-calibrated, and hence trustworthy, predictions.

Under review as a conference paper at ICLR 2021
FEDERATED GENERALIZED BAYESIAN LEARNING VIA DISTRIBUTED STEIN VARIATIONAL GRADIENT DESCENT
Anonymous authors Paper under double-bling review
ABSTRACT
This paper introduces Distributed Stein Variational Gradient Descent (DSVGD), a non-parametric generalized Bayesian inference framework for federated learn- ing. DSVGD maintains number of non-random and interacting particles at central server to represent the current iterate of the model global posterior. The particles are iteratively downloaded and updated by one of the agents with the end goal of minimizing the global free energy. By varying the number of particles, DSVGD enables flexible trade-off between per-iteration communication load and number of communication rounds. DSVGD is shown to compare favorably to benchmark frequentist and Bayesian federated learning strategies in terms of ac- curacy and scalability with respect to the number of agents, while also providing well-caliie vell-calibrated, and hence trustworthy, predictions.
INTRODUCTION
Federated learning refers to the collaborative training of machine learning model across agents with distinct data sets, and it applies at different scales, from industrial data silos to mobile de- vices (Karruz et al. 2019). While some common challenges exist, such the general statistical heterogeneity - "non-iidnes - of the distributed data sets, each setting also brings its own distinct problems. In this paper, we specifically interested in 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' recom- mendations need to be reliable and trustworthy, e.g., to decide when to contact a doctor in case of 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, accom- modating 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 al. 2017) are based frequentist principles, relying the identification of single model pa- rameter 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 trustworthy solution that is able to reduce the number of communication rounds via a non-parametric variational inference-baue 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 (VT) (Angelino et al.] 2016 Neiswanger et al. 2015 Brod- crick 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-a methods in either category include Par- titioned Variational Inference (PVI), which has been recently introduced as unifying distributed VI framework that relies the optimization over parametric posteriors; and Distributed Stochas- tie Gradient Langevin Dynamics (DSGLD), which an MC sampling technique that maintains a number of Markov chains updated via local Stochastic Gradient Descent (SGD) with the addition of

---

Under review as a conference paper at ICLR 2021
Server
Server
f
M
fa Sal
(D)
(
Agents
Agents
argiaxaex Pa (o dlll
9 grg (((e))) " (8))
& (0(0! (0!
(a)
Federated learning across agents equipped with local datasets and assisted by a central server: (a) in DVI agents exchange the current model posterior (0) with the server, while (b) in DSVGD agents exchange particles (0, [0 providing non-parametric estimate of the posterior.
Figure I:
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 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 non- parametric 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 terms of bias, convergence speed, and per-iteration complexity. This paper introduces 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. T DSVGD targets a generalized Bayesian learning formulation, with arbitrary loss functions (Knoblauch 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 downloaded and updated by of the agents by minimizing local free energy functional hefore being uploaded to the server. DSVGD is shown to enable (i) a trade-off between per-iteration communication load and number communication rounds by varying the number of particles; while (ii) being able to make trustworthy decisions through Bayesian inference.
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 (0) for model parameter 0. The agents communicate through a central node with the goal of computing the global posterior distribution g(0) over the shared model parameter ically, following the generalized Bayesian learning framework (Knoblauch et al. 2019), the agents aim at obtaining the distribution (0) that minimizes the global free energy
IR for some prior distribution Po(0) (Angelino al. 2016)). Specif-
min
F(a(O)) = Σ Ea a (
+ +D((()))po)())
(1)
where 0 is a temperature parameter. The (generalized, Gibbs) global posterior gope (0) solving problem (T) must strike balance between minimizing the sum loss function (first term in F(g)) and the model complexity defined by the divergence from a reference prior (second term in F(g)). It given as
K (+))
Topt (0) (f), with qopt (0) = pf(() exp (-
where we denoted as Z the normalization constant. It is useful to note that the global free energy also be written the scaled KI. F(q(O)) aTD((()))) (0)).
2

---

Under review as a conference paper at ICLR 2021
The main challenge in computing the optimal posterior (ope in a distributed manner is that each agent & is only aware of its local loss (0). By exchanging information through the server, the K agents wish to obtain 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 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 BBui et al. 2018 Corinzia & Buhmann 2019b). DVI starts from the observation that the posterior factorizes as the product
K
( = Po( II (0),
where the term tx(-) is given by the scaled local likelihood Lxp(a=>>x(0)) Since the normal- ization constant Z depends on all data sets, the true scaled local likelihood (.) cannot be directly computed at agent k. The idea of DVI is iteratively update approximate likelihood factors (0) for k = 1, 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 ① over distribution
We give here the standard implementation of DVI in which a single agent is schedule at each time, although parallel implementations
possible and discussed below. Accordingly, at each commu- nication round 1,2...., the server maintains the current iterate (0) of the global posterior, and schedules an agent & E {1,2,... K), which proceeds as follows:
1. Agent & downloads the current global variational posterior distribution (0) from the server
(see Fig. ((a)
step
2. Agent h: updates the global posterior by minimizing the local free energy F(0) (q(0) (see Fig.
k
((a) step
2
(0) argmin F (g(0) = Ee [00 (
(4)
{0)
where we have defined the (unnormalized) cavity distribution Pk (0)
(0)(() as
(1)
Pk
(5)
(0)
The cavity distribution (0), which removes the contribution of the current approximate likelihood of agent & from the current global posterior iterate, serves a prior for the update in similar 2). the local free energy is minimized by the tilted distribution pk
(4). In manner
(0) x P (0) with
Pk (0) Pk (6) exp
Ia (+)
(6)
3. Agent i sends the updated posterior (0) updates its approximate likelihood accordingly as
Pk
(4) (-) the server (see Fig.
1(a) step
3) and
00
q(- 1) (i))(0);
(0 (0) = 1/ 1) (0), and the server sets the next iterate as Finally, non-scheduled agents 1/ * l set 10
((). We have the following key property of DVI.
Theorem 1. The global posterior (8) in (2) is the unique fixed point of the DVI algorithm.
The fixed-point property in Theorem I can be verified directly by setting (15-))(0) Zopt (0) and ((-=)(0) =
and by observing that this leads to the fixed point con- (0). The proof is provided in Sec. A.6 Importantly. this property is not tied to the sequential implementation detailed above, and it applics also multiple devices are scheduled in parallel, as long as one sets the next iterate as (0)
dition (0) = (0)
(o(() IIEEKO (0) (0) Iw exC ((i) (0), where K(O) denotes the set of scheduled agents at communi-
cation round i and we have (0) (0) (f-1) (0) and ( (0) updated following 7
3

---

for n 1, N. where k(..-) is the positive definite kernel associated with RKHS H. The first (8) drives the particles towards the regions of the target distribution (i()) with high probability, while the second term drives the particles away from each other, encouraging exploration in the model parameter space. It is known that, in the asymptotic limit of large number N of particles, the empirical distribution encoded by the particles {o!) N converges to the normalized Under review as a conference paper at ICLR 2021
4 PRELIMINARIES
In this section, we briefly review PVI, which serves as an important benchmark, and SVGD, on which we build the proposed Bayesian federated learning solution.
4.1 PARTITIONED VARIATIONAL INFERENCE
4 assumed by DVI is often not tractable. To address this problem, in its most typical form, PVI constrains the local free energy minimization
The exact minimization of the local free energy function
4 to the space of parametric distributions that factorize q(0(n) po I L-11 tx (0)m), where
prior po = ExpFamm (-I11) and approximate likelihood t (-177) ExpFFam (m are selected from the same exponential faaimnl distribution, with natural parameters To and 77. respectively. PVI
follows the same steps as DVI with the caveat that the local free energy (4) for agent & is minimized
over the natural parameter 7/- This can be done efficiently, albeit approximately, using for e.g.,
natural gradient descent [Amari 1998).
The bias imposed by the parametrization in PVI significantly affects the quality of the approximation of the obtained posterior (()) with respect to the true global posterior Gone (0) in the presence of model misspecification. In this case, the fixed point property in Theorem I no longer applies.
4.2 STEIN VARIATIONAL. GRADIENT DESCENT (SVGD)
SVGD tackles the minimization of the (scaled) free energy functional ((g(0) for unnor- malized target distribution p(0), over non paramaaramteccc generalized posterior g()) defined over the model parameters E R The posterior (() is represented by a set of particles {0,, J with 0n E R In practice, an approximation of g(0) can be obtained from the particles {0, [ through Kernel Density Estimator (KDE) (()) N Ex= K(0, 0..) for some kernel function K(--) BBshhop 2006). The particles are iteratively updated through a series of transformations that are optimized minimize the free energy. The transformations
restricted to lie within the unit ball
of a Reproducing Kernel Hilbert Space (RKHS) H. = H. .... x H. It is shown by Liu & Wang (2016) that this optimization yields the SVGD update
NV
819 01--11
> (())) -
e -.l- 11 Vo,, log:(o))")) ") Vo, k(tk "
(8)
N
term in the update
target distribution p(0) x p(0) (Liu 2017b).
DISTRIBUTED STEIN VARIATIONAL GRADIENT DESCENT
In this section, introduce DSVGD, novel distributed algorithm that tackles the generalized Bayesian inference problem (I) via DVI over a non- par-aaaramtic particle-based representation of the global posterior. As illustrated in Fig. 1(b) DSVGD is based the iterative optimization of local free energy functionals (4) via SVGD (see Sec. 4). and on the exchange of particles between the central server and agents. Given the flexibility of the non-parametrie form of the posterior, DSVGD doesn't It suffer from the bias caused by the parametrization assumed by PVI. As result, in the limit of sufficiently large number of particles, DSVGD benefits from the fixed point property of DVI stated in Theorem 1, recovering the true global posterior as a fixed point of its iterations. Furthermore, as we will discuss, DSVGD enables devices to exchange more informative messages regarding the current iterate of the posterior by increasing the number of particles. This can in turn reduce the number of communication rounds and the overall communication load to convergence, at the cost of larger per-round load. In this regard, we note that, in practice, small number of particles is sufficient to obtain statc-of-the-an performance (Liu & Wang 2016). as verified in Sec.
In order to facilitate the presentation, we first introduce a simpler version of DSVGD that has the practical drawback of requiring each agent to store a number of partieles that increases linearly with the number of iterations in which the agent is scheduled. Then, we present a more practical algorithm, for which the memory requirements do not scale with the number of iterations as each agent must only memorize set of N local particles across different iterations. Algorithmic table for
4

---

Under review as a conference paper at ICLR 2021
Algorithm 1: Distributed Stein Variational Gradient Descent (DSVGD)
Input: prior (0), local loss functions (0)) temperature >> 0, kemelss K(-,,) .) and k(::-)
Output: global approximate posterior q(0) = N Σ ===1 K(0,, 0n)
initialize q (0) = po)0); (0.h I==1 ~ po(0): ( = 0 y and (0) (0) = for 1,....A K
iid
for i= I,..,,
Server schedules an Agent &
Agent &t downloads current global particles {0% 1
In=1 from server
Agent & obtains updated global particles {6h) p using 13 foh J and ( J==
Agent l sends the updated global particles {0,² Y to the server
Agent & carries distillation to obtain {o [ J==1 encoding (k (0) using (7 and (o J
end
return (0) = N Σ Σn mm K(0,0k")
MicconttriinddDSVGGD (U-DSVGD) in addition to discussions on complexity and convergence, can be found respectively in Sec. A. and Sec. A..4in in the supplementary materials. A direct extension of DSVGD, termed ParalleI-DSVGD (P-DSVGD), where multiple agents are scheduled per round can be found in Sec. A.5 of the Appendix.
5.1 U-DSVGD
In this section, we present simplified DSVGD variant, which we refer U-DSVGD. We follow the standard implementation of DVI with a single agent scheduled at each communication round i 1.2,.... ., although, as discussed, parallel implementations are also possible. Let define as ( S (1,.... the subset of rounds at which agent l is scheduled prior, and including, iteration i. At the beginning of each round i, the server maintains the iterate of the current global particles
&
(0% [o Jh while each agent & keeps local buffer of particles (o(--1) 0f ) for all previous rounds j E I I((-1) at which agent & was scheduled. The growing memory requirements at the agents will be dealt with by the final version of DSVGD to be introduced in Sec. 5.2 Furthermore, as illustrated in Fig. [(b] at each iteration i, U-DSVGD schedules an agent & E (1,2..... K} and carries out the following steps.
1. Agent k downloads the current global particles {o!! ( I from the server (see Fig. 1(b) step and includes them in the local buffer.
2. Agent ht updates each downloaded particle as
0/0
- 1] + co(( for = 1,..,,, L,
where I. is the number of local iterations: [1] denotes the local iteration index; we have the initializa-
tion 0/0 = of: 1) and the function (6-) is to be optimized within the unit ball of a RKHS Hd. The function (()) is specifically optimized to maximize the steepest descent decrease of a particle-based
approximation of the local energy (4). To elaborate, we denote ( = Σ K(0, o!5-")
the KDE of the current global posterior iterate encoded by particles {o!! ( 1 In-1- Adopting the fac- torization for the global posterior (ef. Z7), we define the current local approximate likelihood
(0) g (()))(()) (0).
(0)
II
(10)
jEI(--1 q0-1)(0) 9 (0)
Note that 100 can be computed using all the particles in the buffer at agent & at iteration i. FFnally, the (unnormalized) tilted distribution Pk (cf. (6) is written
=(6)
Pk (0) g
exp a ( (-1)
(11)
Following SVGD, the update 9 is optimized to maximize the steepest descent decrease of the Kullback-Leibler (KL) divergence between the approximate global posterior (l) encoded via particles (0%) N and the tilted distribution (((0 (0) in 11 (see Fig. ((b) step 2 )), i.e.,
5

---

Under review as a conference paper at ICLR 2021
D- 100)
1 [L- 1007 Agant
2 (C- 100) Agant
82 100)
1 CC - (00)
Agant
Agant
Gaussian toy example with uniform prior and K = 2. Dashed lines represent local poste- riors, the shaded area represents the true global posterior, while the solid blue line is the approximate posterior obtained using a KDE over the particles.DSVC DSVGD schedules agent and odd and even number of communication rounds i, respectively.
Figure 2:
-e 1- (O)IEE (0)), s.t. (O)Ib (0)),
1 $1
(-) arg max d())H
(12)
de
Thus, recalling 8 the particles updated
a! n
011-11
0l
log Pe Vo, (of- ko,k( " 1-
+
N
3. Agent k sets 0 for n ,,...,.. N. Particles (0 Jn=z are added to the buffer and sent (13)
the server (see Fig.
I(b) step 3) that updates the current global particles as {0m] (0h² JNN
n==1-
In order to implement the described U-DSVGD algorithm, need compute the gradient in [3) at agent k. First, by III we have
Vo log Pkc (0) Vo log (1-1) 1) (6) - Vo log th ( Lo)
(() 1
(14)
a
Using 10 the second gradient term can be obtained a recursive manner using the local buffer as
Ve log (i-2) (0) if agent not scheduled at iteration (i (0-1) & (1-2) & (f) + Vo (i-1 Vo log (()-2(()) otherwise.
V log ( k 11 (0) =
(15)
Vo log
Finally, the gradients V log g00 (0) can be directly computed from the KDE expression of (0), with initializations (0)) (0) asymptotic convergence properties of SVGD in terms of local free energies, but existing results do not imply that the global free energy decreases across the iterations. This result provided in the next theorem, whose precise formulation can be found in Sec. Af the Appendix.
and (0) Po()). The inner loop of U-DSVGI inherits the
Theorem 2 (Guaranteed per-iteration decrease of the global free energy.). The decrease in the global free energy from local iteration I to I 1 during communication round for which agent is sched- uled can be lower bounded as
(0) - (())+1((0) 2 4 Pk
(1) ery) - -2a(K- - 1))(() max
(16)
where (1) sup max log(t (0)) exp( (0)) S(g. p) denotes the Kernalized Stein Dis- ( 6 mffk
crepancy between distributions q and p Liu al. 2016 and y constant depending on the
RKHS kernel and the target distribution.
The first term bound 16) quantifies the decrease in the local free energy at agent k. which depends on the "distance" between current iterate rlf and the local target given by the tilted distribution ( (0); while the second term quantifies the effect of the update on the local free energies of other Ph
agents. the presence of only one agent, the second terms reduce to zero, and one recovers the upper bound on the guaranteed per-iteration improvement for SVGD derived in Korba et al. 2020
5.2 DSVGD
In this section, we describe the final version of DSVGD, which, unlike U-DSVGD, requires each agent & to maintain only N local partieles (0(1) (. x=1 across the communication rounds i 1,2..... To this end, in each round i, at the end of the L local SVGD updates in (13), DSVGD carries out

---

Under review as a conference paper at ICLR 2021
Covertype dataset
Twonorm dataset
Covertypa dataset
Twononn dataset
(2 agante)
Accur sarryay (20 aggunta)
Accuracy
02 appe
Accurag
(20
DOVID N-
SHOT
DAVOD
SAND FSLLO
DOVOD Perr
Pubde
Puddeg
Number rounds
Number of rounds
Number of rounds
i
Number of rounds !
Figure 4: Accuracy for Bayesian logistic regression with (left) K agents and (right) K 20
agents as function of the number of communication rounds i (N = 6 particles, L = 200).
form of model distillation (Hinton et al. 2015 Chen & Chao 2020 via SVGD. Specifically, L' ad- ditional SVGD steps are used to approximate the term b (0) using the N local particles {0 1/
is noted that this approximation step is not necessarily harmful to the overall performance, since describing the factor ( (0) with fewer particles can have a denoising effect acting as regularizer.
DSVGD operates as U-DSVGD apart from the computation of the gradient in (14 and the man- agement of the local particle buffers. The key idea is that, instead of using the recursion (15)
to compute 14 DSVGD computes the gradient Ve log & (0) from the KDE tk (
En K(0,
based on the local particles ( I== in the buffer. At the end of each round
i, the local particles (ol:=") In=1 updated by running L' local SVGD iterations with target given (0))
by the updated local factor f (0) = (0) t ((i-1)(0),
(0). This amounts to the updates
e(' 0!"-11
E'
N [m(ol" -1) "pr o log to of
+ Vo, M(ll" 11 -f'--"") "ly,
+
(17)
for l' = I....,L'' L' and some learning rate e', where the gradient V log (0) (0)
V log o (0) +
No log l/ (i-1) (0) Vo log can be directly computed using KDE based on the available particles (0!²) J=== (updated global particles), (B(i=1)) J== (local particles) and [0!: ( (down-
loaded global particles). Finally, we note that the distillation operation can be performed after send- ing the updated global particles to the server and thus enabling pipelining of the L' local iterations
with operations at the server and other agents. DSVGD summarized in Algorithm
6 RELATED WORK
Extensions of SVGD. Since its introduction, SVGD has been ex- tended in various directions. Most related to this work Zhuo et al.) (2018), which introduces a messagg-ppsssnng SVGD solution for high-dimensional latent parameter spaces by leveraging conditional independence properties in the variational posterior; and Yoon et al. 20188), which uses SVGD as the per-task base learner in a meta- learning algorithm approximating Expectation Maximization.
KL divergence belween
posterior approximat
0.30
U-DSVGD
DSVGD
Pv1
GVI
a.2
SVGD
a.s
2000
Number of rounds
Generalized Bayesian Inference. Owing to its reliance on point estimates in the model parameter space, frequentist learning meth- ods, such as Federated Stochastic Gradient Descent (FedSGD).
Figure 3: KL divergence be- tween exact and approximate global posteriors as function of the number of rounds i LL= 200).
FedAvg and their extensions (Zhang et al.] 22200] Li et al. 2018 2020
Pathak & Wainwright 2020 Nguyen et al. 2020 Wang et al.
are limited in their capacity to combat overfitting and quantity uncertainty (Guo et al.) 2017 Mitros & Mac Namee 2019; Neal, 2012 Jospin et al. 2020; MacKay 2002). This contrasts with the gen- eralized Bayesian inference framework that produces distributional, rather than point, estimates by optimizing the free energy functional, which is a theoretically principled bound on the generalization performance (Zhang 2006; Knoblauch et al., 2019). Practical algorithms for generalized Bayesian inference can leverage computationally efficient scalable solutions based on either MC sampling or VI methods (Angelino et al. 2016 Alquier et al. 2016)).
Distributed MC Sampling. The design of algorithms for distributed Bayesian learning has been so far mostly focused on one-shot, or "embarrassingly parallel", solutions under ideal communications (Jrdan et al. 2019). These implement distributed MC "consensus" protocols, whereby samples from the global posterior are approximately synthesized by combining particles from local posteriors
(Soott et al. 2016 Liu & Thler 2014). Iterative extensions, such as Weierstrass sampling (Wang &
7

---

Under review as a conference paper at ICLR 2021
Rendell et al. 2018), impose consistency constraints across devices and iterations in a way similar to the Alternating Direction Method of Multipliers (ADMM) (Angelino et al. 2016
Dunson 2013
State-of-the arr results have been obtained via DSGLD (Ahn et al. 2014).
Distributed VI Learning. Considering first one-shot model fusion of local models, Bayesian meth- ods have been used to deal with parameter invariance and weight matching (Yurochkin et al. 2019
Claici et al. 2020). Iterative VI such as streaming variational Bias (SVB) (Broderick et al. provide a VI-based framework for the exponential family to combine local models into global ones.
2013)
PVI provides a general framework that can implement SVB, as well as online VI (Bui et al. and has been extended multi task learning in Corinzia & Buhmann (2019a).
2018
7 EXPERIMENTS
As in Liu& & Wang (2016). for all our experiments with SVGD and DSVGD, we use the Radial Basis Function (RBF) kernel k(x, zo) = exp(-Ilæ zo 2011//h). The bandwidth h is adapted to the set of particles used in each update by setting h = med?/ / log n. where med is the median of the pairwise distances between the particles in the current iterate. The Gaussian kernel K(.,)) used for the KDEs has bandwidth equal to 0.55. Unless specified otherwise, we use AdaGrad with momentum to choose the learning rates and e' for (U-)DSVGD. Throughout, we fix the temperature parameter a = in []. Finally, to ensure a fair comparison with distributed schemes, we run centralized schemes for the same total number I L of iterations across all experiments. Additional results for all experiments can be found in Appendix B in the supplementary materials, which include also additional implementation details.
Covertype Dataset
Gaussian 1D mixture toy example. We start by considering
simple one-dimensional mixture model in which the local un-
Log-lkellhood
normalized local posteriors py. (0) = Pf()
LL (0)) at
each agent k are defined as pi (0) po (0), 4) and p2(0)
DBVGD homageneous
Po (N)() - 3. 1) + (0 2)) and the prior (0) is uniform over [-6, 6]. i.c., (0) = U(0) - 6, 6). The local posteriors are shown in Fig. as dashed lines, along with the global posterior (0))(() x Tppe in (2). which represented as a shaded area. We fix the number of particles to N = 200. The approximate poste- riors obtained from the KDE over the global particles are plotted in Fig. 2 as solid lines. It can be observed that at each round, the global posterior updated by DSVGD integrates the local likelihood of the scheduled agent, while still preserving information about the likelihood of the other agent from prior iterates, until (approximate) convergence to the true global posterior Jopt: which is a nor- malized version of Copt in ( is reached. Finally, in Fig. 3 plot the KI divergence between (()) and gope (0) as a function of the number of rounds. Both U-DSVGD and DSVGD exhibit simi- lar behaviour, converging SVGD and outperforming the parametric counterparts PVI and Global Variational Inference (GVI) (uui et al. 2018).
DSVED reterogenesus
Feddug homogenecus
Fod/ug hoterogeneous
Number of rounds ?
Figure 5: Log-likelihood for Bayesian logistic regression with non-iid data distributions
(N = 6, L = L' = 200).
Bayesian logistic regression. We now consider Bayesian logistic regression for binary classification using the same setting as in GerGhmann et al. (2012). The model parameters 0 = [w, log(E)) include the regression weights w E IR.' along with the logarithm of a precision parameter& &. The prior is given as Po(w, {) Po Po(w()))m (W)))) with Pa (o(w)) = NNWW0,, 00 La) and PC( Gamma(Ela, b) with a = and b 0.01. The local training loss (0) each agent k is given as Ix = E(x. 1(xx, Ve. w). where D2 is the dataset at agent i with covariates x E Rd and label yk E 11, and the loss function 1(X2, Uk. w) is the cross-entropy. Point decisions are taken based on the maximum of the average predictive distribution. We consider the datasets Covertype and Twonorm (Gershman et al. 2012). We randomly split the training dataset into partitions of equal size among the K agents. We also include FedAvg, Stochastic Gradient Langevin Dynam- ies (SGLD) and DSGLD for comparison. We note that FedAvg is implemented here for consistency with the other schemes by scheduling a single agent at each step. In Fig. 4 we study how the accuracy evolves as function of the number of communication rounds i, or number of communica- tion rounds, across different datasets, using N = 2 and N = 6 particles. We observe that DSVGD consistently outperforms the mentioned decentralized benchmarks and that, in contrast to FedAvg and DSGLD, its performance scales well with the number K of agents. Furthermore, the number N of particles is seen to control the trade-off between the communication load, which increases with N, and the convergence speed, which improves as N grows larger. Through reduction of the number
8

---

Under review as a conference paper at ICLR 2021
Kin8nm dataset
Year Prediction dalaset
KinBnm dataset
Year Prediction dataset
D5V33
DEVDD
poyG3
Average RMSE
Fed/ig
D5.0LD DLD
Average RMSE
(2 agente)
DRGLD
Fedfrg
Average RMSE
(20 a agental
Pesheg
D60 0
Average RMSE
(20 agents)
Pediog
DSGLD
DVG2
SVSI
SVGD
SVDD
Number rounds
Number of rounds
Numbor of rounds: 1
Number of rounds
Figure 6: Average RMSE as a function of the number of communication rounds for regression using BNN with a single hidden layer of ReLUs with (left) K = 2 agents and (right) K = 20 agents (N- 20, L
200, 100 hidden neurons for the Year Prediction and 50 for Kin8nm).
MNIST datasel
Fashion MNIST datasel
MNIST dataset
Fashion MNIST datasel
Accuracy
(2 (gggesss)
Accurbat
Ceceae le Iannee
Accuracy
(20 agents)
Accuracy
agents)
newap
DEVGO
DEVi2
Tackkog
Fachug
Fac/vg
9/100
8V3D
3+00 SOLD
BV30
S0LD
Number rounds
Numbor rounds:
Number rounds
Numbor rounds
Figure 7: Multi-label classification accuracy using BNN with single hidden layer of 100 neurons as function of i, or number of communication rounds, using MNIST and Fashion MNIST with (left) K agents and (right) K 20 agents (N = 20, L L' 200).
of communication rouds, DSVGD can also reduce the overall communication load. For example, in the third plot in Fig. 4 DSVGD reaches an accuracy of 70% after communication rounds with N = 6, requiring the exchange of 30 particles. In contrast, FedAvg requires around 100 rounds to obtain the same accuracy, making the total communication load much higher than that of DSVGD.
To capture heterogeneous datasets with non i..dd. data, we now consider for different dataset parti- tions across K = 4 agents. In the homogeneous case, labels are split equally among agents, while, in the heterogeneous case, each agent stores 40% of one label and 10% of the other. DSVGD is seen in Fig. 5 to have a robust performance against heterogeneity as compared to FedAvg, whose con- vergence speed is severely affected. This result hinges on the fact that Bayesian learning provides a predictive distribution that is a more accurate estimate of the ground-truth posterior distribution. This is true irrespective of the level of "noo-iidness" Bayesian learning can account in a principled away for all competing "explanations" provided by different devices. This is in contrast to FedAvg, whose reliance on a point estimate of the parameters yields an overconfident predictive distribution that cannot properly account for the diversity of predictions provided by different devices.
Bayesian Neural Networks. We now consider regression and multi-label clas- sification with Bayesian Neural Networks (BNN) models. The experimental setup is the same as in Mlernnndez-Lobate & Adams (2015), with the only exception that the prior of the weights is set to Po(w (w) = [W A >--La) with fixed pre- cision/ À = e. We plot the average Root Mean Square Error (RMSE) for K 2
FedAvg
DSVGD
Accuracy
of rasroru
Af 100
Purfect ca brat
Parluct calloration
M4 M Confidence
6a
10
Confideriee
Bs s0 10
Figure 8: Reliability plots for classification using BNN with variable number of hidden neurons using fashion MNIST (N- = 20, / 10, L L' 200, K 20). sion over the Kin8nm and Year datasets, and accuracy for multi-label classification on the MNIST and Fashion MNIST datasets in Fig. 7 Confirming the results for logistic regression, DSVGD con- sistently outperforms the other decentralized benchmarks in terms of RMSE and accuracy, while being more robust in terms of convergence speed to an increase in the number of agents.
and K = 20 agents in Fig. 6 for regres-
Calibration. Reliability plots are a common visual tool used to quantify and visualize model cal- ibration (Guo et al.) 2017). They report the average sample accuracy as function of the confidence level of the model. Perfect calibration yields an accuracy equal to the corresponding confidence (dashed line in Fig. B). Fig. X shows the reliability plots for FedAvg and DSVGD on the Fashion MNIST dataset for the BNN setting. While increasing the number of hidden neurons negatively affects FedAvg due overfitting, DSVGD enjoys excellent calibration even for large models and is hence able to make trustworthy predictions.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Sungjin Ahn, Babak Shahbaba, and Max Welling. Distributed Stochastic Gradient MCMC. In Eric P. Xing and Tony Jebara (eds.), Proceedings the 31 International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pp- 1004--105,, Bejing, China, 22-24 Jun 2014. PMLR. URL/ttt:://pooceeeeee wtt.:eeiiiinns.....essee.e wtt../.....e.//////attaaaaa html
Pierre Alquier. James Ridgway, and Nicolas Chopin. On the properties of variational approximations of Gibbs posteriors. Journal of Machine Learning Research, 17(236):1-4) 2016. URI http: hwtpl//rrgpppppr//eeeee http//11111111111 html
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-
276, 1998.
Elaine Angelino, Matthew James Johnson, and Ryan P Adams. Patterns of scalable Bayesian infer- ence. arXiv preprint (rXiv:1602.0522) 2016.
Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statis- tics). Springer- Verlag, Berlin, Heidelberg, 2006. ISBN 0387310738.
Tamara Broderick, Nicholas Boyd, Andre Wibisono, Ashia C Wilson, and Michael I Jordan. Streaming Variational Bayes. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 26, pp. 1727-1735. Curran Associates, Inc., 2013. URL http::/ nips. ccppsccc/ppprr/ 4980 1980-st reami nn- vari a ona -aaallbbaes... .pdf
Thang D Bui, Cuong Nguyen, Siddharth Swaroop, and Richard E Turner. Partitioned Variational Inference: A unified framework encompassing federated and continual learning. arXiv preprint 1rX.v:1811.11206, 2018
Hong- You Chen and Wei-Lun Chao. FedDistill: Making Bayesian Model Ensemble Applicable to Federated Learning. 2020.
Sebastian Claici, Mikhail Yurochkin. Soumya Ghosh, and Justin Solomon. Model Fusion with Beilbackk-eibber Divergence. arXiv preprint arXiv:2007.061 06166, 2020.
Luca Corinzia and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint
1rXiv:1906.06268 2019a.
Luca Corinzia and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint arXiv: 06266.6668.88 2019b.
Bo Dai, Niao He, Hanjun Dai, and Le Song. Provable Bayesian Inference via Particle Mir- ror Descent. In Arthur Gretton and Christian C. Robert (eds.). Proceedings of the 19th In- ternational Conference on Artificial Intelligence and Statistics, volume 51 of Proceedings of Machine Leurning Research, pp. 985-994, Cadiz, Spain, 09-11 May 2016. PMLR. URL
http: / httpp//ppcppp.cppppppppppp mlr..................mmm....mm..
Morris H. DeGroot and Stephen E. Fienberg. The Comparison and Evaluation of Forecasters. Journal of the Royal Staristical Society. Series D (The Statistician), 32(1/2):12-22, 1983. ISSN 00390526, 14679884. URL http:: jstor org http://aaabeeee9//5ee/
Samuel J. Gershman, Matthew D. Hoffmann, and David M. Blei. Nonparametric Variational Infer-
ence. In Proceedings of the 29th International Coference on International Conference on Machine Learning. ICML' 12, pp. 235-242, Madison, WI, USA, 2012. Omnipress. ISBN 978195011 778845031222 12851.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On Calibration of Modern Neural Networks. In Proceedings of the 34th International Conference Machine Learning Volume 70, ICML'CC 17, pp. 1321-1330. JMLR. org, 2017.
José Miguel Hernández-Lobann and Ryan P. Adams. Probabilistic backpropagation for scalable neural networks. In Proceedings of the 32nd International Conference In- ternational Conference on Machine Learning Volume 37, ICML' 15, pp. JMLR.org, 2015.
learning of bayesian
10

---

Under review as a conference paper at ICLR 2021
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in Neural Network. arXiv preprint arXiv:1503.025:: 2015.
Michael I. Jordan, Jason D. Lee, and Yun Yang. Communication--ffficcee Distributed Statis- tical Inference. Journal of the American Statistical Association, 114(526):668 681, 2019. doi: 10. 2110:001622459921188429999999999 201 8 1429244
(tttps::://oo/ /tt.:s///ob..crrg//..000///////2
Laurent Valentin Jospin, Wray Buntine, Farid Boussaid, Hamid Laga, and Mohammed Bennamoun. Hands-on Bayesian Neural Networks-a Tutorial for Deep Learning Users. arXiv preprint (rXiv:2007.06823 2020.
Peter Kairouz, H 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:/912.04977, 2019.
Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. Generalized variational inference. star, 1050:21, 2019.
Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A Non- Noymmsymiootic Analysis for Stein Variational Gradient Deseent. arXiv preprint arXiv:2006.09997, 2020.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.
Federated Optimization in Heterogeneous Networks. arXiv preprint arXiv:1812.0612. 2018.
Frank Po-Chen Lin, Christopher G Brinton, and Nicolò Michelusi. Federated Learning with Com- munication Delay in Edge Networks. arXiv preprint arXiv:2008.0932. 2020.
Qiang Liu. Stein Variational Gradient Descent as Gradient Flow. In I. Guyon, U. Luxburg, S. Bengio, H. Wallach, Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, pp. 3115-3123. Curran Associates, Inc., 2017a. URL lattp://papers nips ctt.pprr/cpppp////////// wtw.s--grrrrrrreeeeloo...ooood flo..
Qiang Liu, Stein Variational Gradient Descent as Gradient Flow. In Advances in Neural Informa- tion Processing Systems, pp. 3115-3123, 2017b. URL URU/http://pppers..22 http:/ppwwpr.........ooop.pappaa
Qiang Liu and Alexander lhler. Distributed Estimation, Information Loss and Exponential Families. In Proceedings of the 27th International Conference on Neural Information Processing Systems- Volume 1, NIPS' 14, pp. 1098-1101 1106, Cambridge, MA, USA, 2014. MIT Press.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in neural information processing systems, pp. 2378-22888 7778-238,,222 2016.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized Stein discrepancy for goodness-of-ff tests. In International conference machine learning, pp. 276-284, 2016.
David J. C. MacKay. Information Theory, Inference & Learning Algorithms. Cambridge University Press, USA, 2002. ISBN 0521642981.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Learning of Deep Networks from Decentralized Data. In Proceed- ings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA, volume 54 of Proceedings of Machine Learning Research, pp. 1273-1282. PMLR, 2017. URL http::t httt:://proceeeingg... mir. wtt.eeediinnsmm....eeessss// memahan1 7a. html
Communication-Eficn
Diego Mesquita, Paul Blomstedt, and Samuel Kaski. Embarrassingly Parallel MCMC using Deep Invertible Transformations, volume 115 of Proceedings of Machine Learning Research, pp. 1244 1252, Tel Aviy, Israel, 22-25 Jul 2020. PMLR. URL URL.Attp://prooeeaiii .mlr. press/ v1 v115/mesquita20a... htm.0l
11

---

Under review as a conference paper at ICLR 2021
S. Mika, G. Ratsch, J. Weston, B. Scholkopf, and K. R. Mullers. Fisher discriminant analysis with kernels. Neural Networks for Signal Processing IX: Proceedings of the 1999 IEEE Signal Processing Society Workshop (Cat. No. 18.98878448) pp. 41-48, 1999.
John Mitros and Brian Mac Namee. On the Validity of Bayesian Neural Networks for Uncertainty Estimation. arXiv preprint arXiv:1912.01539, 2019.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Willie Neiswanger, Chong Wang, and Eric Xing. Embarrassingly parallel variational inference in nonconjugate models. arXiv preprint arxiv:1510.04/63, 2015.
Hung T Nguyen, Vikash Sehwag, Seyyedali Hosseinalipour, Christopher G Brinton, Mung Chiang, and H Vincent Poor. Fast-Convergent Federated Learning. arXiv preprint arXiv:2007.13137, 2020.
Alexandru Niculescu-Mizil and Rich Caruana. Predicting Good Probabilities with Supervised Learning. In Proceedings of the 22nd International Conference on Machine Learning, ICML. '05. pp. 625-632, New York. NY, USA, 2005. Association for Computing Machinery. ISBN 1595931805 doi: 10.1145/1 102351.
URL
1145/
1102351 1102430
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, d Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8026-8037. Curran Associates, Inc., 2019. URL htwp: httpt:t::pppptttatappppppe nips. 9015 1115-pytorch-ee an- www--mmmeeeeeeettteeeeeee pdf
wwtpp//eaa.iissccpppppppe..eee cc/paper/ ottp:p/ooooooooaaaooooaaaeeeteaa hac.aaee-ettteeiccccceoooooaaaae wtw.-oommmnnneeeeee-eeeeiaannn library.
Reese Pathak and Martin Wainwright. FedSplit: An algorithmic framework for fast federated optimization. arXiv preprini arXi::2005.05238, 2020.
Thomas Pinder, Christopher Nemeth, and David Leslier. Stein Variational Gaussian Processes. arXiv preprint arXiv:2009. 2020.
Mark S Pinsker. Information and information stability of random variables and processes. Holden-
Day, 1964.
Lewis J Rendell, Adam M Johansen, Anthony Lee, and Nick Whiteley. Global consensus Monte Carlo. arXiv preprint arXiv:/807.099 887.09288, 2018.
M. Sato. Online Model Selection Based on the Variational Bayes. Neural Computation, 13(7):
1649-1681, 2001.
Steven L. Scott, Alexander W. Blocker, Fernando V. Bonassi, Hugh A. Chipman, Edward George, and Robert E. McCulloch. Bayes and Big Data: The Consensus Monte Carlo Algo- rithm. International Journal of Management Science and Engineering Management, 11:78- 88, 2016. URL http:/ www. Lannfonnnnnnnneee www.ecoom/aaooo /ttp////1111...000/1111111111110 2016 42191
Aki Vehtari, Andrew Gelman, Tuomas Sivula, Pasi Jylanki, Dustin Tran, Swupnil Sahai, Paul Blom- stedt, John Cunningham, David Schiminovich, and Christian Robert. Expectation Propaga- tion as Way of Life: A Framework for Bayesian Inference on Partitioned Data. Journal of Machine Learning Research, 21(17):1-53, 2020. URL http:///mmm....ggapeeres//// 18 817 317.html
Jany Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization. arXiv preprint arXiv:2007.07481. 2020.
12

---

Under review as a conference paper at ICLR 2021
Xiangyu Wang and David B. Dunson. Parallel MCMC via Weierstrass abs/1312.4605, 2013.
Sampler. ArXiv,
Zheng Wei and Erin M Conlon. Parallel Markov chain Monte Carlo for Bayesian hierarchical models with big data, in two stages. Journal of Applied Statistics, 46(11):1917-1936 2019.
Max Welling and Yee Whye Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning. ICML' 11, pp. 681-688, Madison, WI, USA, 2011. Omnipress. ISBN 9781450306195.
Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian Model-Agnostic Meta-Learning. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Desa-BBanchh, and R. Gamett (eds.), Advances in Newral Information Processing Systems 31, pp. 7332 7342. Curran Associates, Inc., 2018. URL/ttpp:/paatee
nips cctccc/pppaac ccppappr/
7963 bayesian Saayesian-modell model agnostic meta learning.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian Nonparametric Federated Learning of Neural Networks. Kama- lika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Confer- ence on Machine Learning, volume 97 Proceedings of Machine Learning Research, pp. 7252- 7261, Long Beach, California, USA, 09-15 Jun 2019. PMLR. htt: http://eorccee......//.. mlr. wttppp/pww//...oo.//pcccccp
9a. html
Tong Zhang. Information-theorede upper and lower bounds for statistical estimation. IEEE Trans- actions on Information Theory, 52(4)::307-13211 2006.
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data. arXiv preprint MrXiv:2005.11418, 2020.
King-Xuan Zhou, Derivative reproducing properties for kernel methods in learning theory. Journal of computational and Applied Mathematics, 220(1-):456-463, 2008
Jingwei Zhuo, Chang Liu, Jiaxin Shi, Jun Zhu, Ning Chen, and Bo Zhang. Message Passing Stein Variational Gradient Descent. In International Conference on Machine Learning, pp. 6018-6027.
PMLR, 2018
13

---

Under review as a conference paper at ICLR 2021
A COMPLEMENTARY MATERIALS
A.1 ALGORITHMIC TABLES
Algorithm 2: Partitioned Variational Inference (PVI) (Bui et al. 2018)
Input: prior po (0). local loss function (L)( temperature a > 0
Output: global posterior g(0l1)
initialize ( (0) for k 1,... K; for i do do
po(0)
At scheduled agent download current global parameters (1-1) from server Agent & solves local free energy problem in to obtain new global parameters r7())
Agent & sends (7 to the server and server rj 17()
Agent & updates new approximate likelihood: la (0)m))
2(0 1(01) 12 (011))
end
return ((0) q((1r)
Algorithm 3: Stein Variational Gradient Descent (SVGD) (Liu & Wang 2016)
Input: target distribution p(0). initial particles (0!.) j (0). kernel k(., :). learning rate
Output: particles [0,]m= that approximates the target normalized distribution
for i 1,..., L do
for 1...., N do
0 - 0
N N Ez=1
[k(0fi=1) - Vo, logū(0fi-1))+ + Vo, 1(0fi-11,000
end
end
q(0) = K(0, (
return
Algorithm 4: Annconstrained-Distiibue Stein Variational Gradient Descent (U-DSVGD) 0, learning rate 0, kernels
Input: prior (o(O), local loss function (0) temperature k(---) and -).,-)
Output: global posterior q()) = EN,, K(0,0.)
initialize t (0) (0) for h 1...., K; ( (0) (d
(0); (0. J
iid po(
for
1,..., do
// New communication round: server schedules an agent
At scheduled agent k. download and memorize in local buffer current global particles ()(1-1) Ja=i
Agent k sets {6]0) = =((--)) 1n=1
for = 1,..., L do // Local iterations: agent & minimizes lccal free energy
Compute V 0, = V, log (i- - Vo log (f-") (olt-")) - L. (0!'-1) with KDE
( (0) = ΣN K(0, o and Va log ((-11) (0) computed using 15
for particle
I,
do
AP.. + N Σ, k(pll"", 0l! v
1
Vo, k(ol -0. 11 0 -0hh11)
! + 1111 + cddn
end end
11
12
Agent & sets updated global particles {o!") = 0.40 N
- and memorize them in the local buffer
Agent & sends partieles (0(.) , to the server and server sets (0., = ((() In=l
13
14 end
return (())= = >NN K(0,00 )(0)
15
14

---

Under review as a conference paper at ICLR 2021
Algorithm 5: Parallel-Distribuue Distributed Stein Variational Gradient Descent (P-DSVGD)
Input: prior Po (0). local loss functions (Lx( 2-1, lemperature kernels K(.,)) :) and k(.)))
Output: global approximate posterior q(0) N Σm=l K(0, d..)
initialize q (0) (0); (0 N Lid po()); (0t0) 0 1 I==1 and { ())()) (0) 1 for
k 1, K for i
do
Server schedules a set (3))) of agents in parallel
Agents downloads current server particles (oh (oll-)) from server
Agents obtains updated global particles {0(.) 1 Ju using 13 (o!!-1)
1 N and
{ (0(t-2)) J==
Agents carries distillation to obtain {0(.) 1X
[ encoding { using 17 and {or (0) N
Ja=l
Agents sends the obtained local particles (0()) 1N for k KC the server
Server obtains (6h (9) using 29
[on Jn=1 and [0 W for E X((
end
10 return (0) = NN- > K(0, )
15

---

Under review as a conference paper at ICLR 2021
A.2 A RELATIONSHIP BETWEEN PVI AND U-DSVGD
We show here that PVI with Gaussian variational posterior ((0 (0 (21a) of fixed covariance IL and mean 2 parametrized by natural parameter can be recovered as a special case of U-DSVGD. To elaborate, consider U-DSVGD with one particle 01 (i.c., N 1), RKHS kernel that satisfies Vok(0, 0) = 0 and k(0, 0) = (the RBF kernel is an example of such kernel) and isotropic Gaussian kernel K(0,, ) (((lo AA2La) of bandwidth À used for computing the KDE of the global posterior using the particles. The U-DSVGD particles update in (33 reduces to the following single particle update:
0- - of--1) + VVo logg (0
(0!--1), for = I....,L,
(18)
with tilted distribution
(0)
(-1e(10)
Pk (0) & (0) exp
(19)
a
The numerator in 19 can rewritten as ( = K(0, K(0, 0(1 = ((0)))(-1)) with 1/(
A 2-20( while the denominator can be rewritten as
((0)))77) 10) t/ (0 (0)m6-11))
(0)
II
(20)
SEI
with Ml
(1-1) EjEzg- 1y 176 - 115-11 This recovers the PVI update (6
A.3 RELIABILITY PLOTS
In this part we give some background on reliability plots and Maximum Calibration Error (MCE). Reliability plots Niculescu-Mizil Caruana 2005). Consider a model that outputs a prediction f(x,) and a proba- bility p(z;) of correct detection for an input æ; with true label y4. We divide the test samples into bins {Bر} each bin B, containing all indices of samples whose prediction confidence falls into
visual tool to evaluate model calibration (DeGroot & Fienberg 1983
the interval (111 1 where B is the total number of bins. Reliability plots evaluate the accuracy as function of the confidence which are defined respectively as
acc(B,)
(((2))==
B
1
and cBnf(
p(xx).
B
Perfect calibration means that the accuracy is equal to the confidence across all bins. For example, given 100 predictions, each with confidence approximately 0.7, one should expect that around 70% of these predictions be correctly classified.
To compute p(z),, we need the predictive probability p(y)(x;) for all samples E [1: This can be obtained by marginalizing the data likelihood with respect to the weights vector w. This marginal- ization generally intractable but can be approximated for both Bayesian logistic regression and Bayesian Neural Networks as detailed in Sec. A.3. and Sec. A.3.2
While reliability plots are a useful tool to visually represent the calibration of a model, is often desirable to have a single scalar measure of miscalibration. In this paper, we use the MCE that measures the worst case deviation the model calibration from perfect calibration (Guo et al. 2017). Mathematically, the MCE is defined
MCE max (acc(B,) (B,))
(21)
JE(1....BB
Additional numerical results using both reliability plots and MCE can be found in Sec. B.5
16

---

Under review as a conference paper at ICLR 2021
A.3.1 PREDICTIVE DISTRIBUTION FOR BAYESIAN LOGISTIC REGRESSION WITH SVGD AND DSVGD
In this section, we show how the predictive distribution for the Bayesian logistic regression ex- periment can be obtained when using DSVGD or SVGD. The predictive distribution provides the confidence values to be used in the calibration experiment. Given KDE of the posterior q(w) = En=i k(w, Wx) with N particles (wn)n=r the predictive probability for Bayesian lo- gistic regression can be estimated as
p(u: = 1|x,) ~0 / p(v: = 1|x,, w)g(w)dw = Σ
exp(2x² Wa w|l?) + exp(- dw.
N(222
(22)
A good approximation of (22) can be obtained by replacing the logistic sigmoid function with the
probit function (Bishop 2006 Sec. 4.5), yielding
1
ply = 1(x;) ~
(23)
N + exp( 1xxp(-x((22pup
where
14 W
o² 1
12 xxxx
(24)
-1/2
and n((²)
1 0² 2T
A.3.2 PREDICTIVE DISTRIBUTION FOR BAYESIAN NEURAL NETWORKS WITH SVGD AND DSVGD
In a manner similar to the predictive distribution for BNN can be estimated as
N(2X2x)a12 / [(x,, w) R=1
N
(=IW;, W, 1 dw,
p(u: = 1 (x)
(25)
2A
where f(xe, w) is the sigmoid output of the BNN with weights w. Using the first order Taylor approximation of the network output around the n-th particle (Bishop 2006 Sec. 5.7.1)
fx,,, w) R f(x,, w.) V fx,, w wn),
(26)
the predictive distribution now be rewritten
/UfEe. wn) V (x) w)(w wnn) exp N
-- -||w, w/12
ply: 1|xx)
dw
(222
2X2
N
fxt, w.) +
( w)w, -W w)wn)
N
R=1
N Σ
/ ((x,
(27)
where we have used the fact that f [N(w/w.. XPIa)dw 1 and wwW/W//www www AFL W..
A.4 SPACE-TIME COMPLEXITY, COMMUNICATION LOAD AND CONVERGENCE
This section offers a brief discussion on the complexity, communication load and convergence of DSVGD.
Space Complexity. DSVGD inherits the space complexity of SVGD. In particular, DSVGD re- quires the computation of the kernel matrix k(.,.) between all particles at each local iteration, which
17

---

Under review as a conference paper at ICLR 2021
can then be deleted before the next iteration. This requires O( space complexity. As pointed out by Liu & Wang (2016) and noticed in our experiments, for sufficiently small problems of practical interest for mobile embedded applications, few particles are enough to obtain state-of-the art per- formance. Furthermone, N particles of dimension d need to be saved in the local buffer, requiring O(Nd) space. Given that N is generally much lower than the number of data samples, saving the particles in the local buffer shouldn't be problematic.
Time complexity. When scheduled, an agent has perform ( L')N2) L')N operations with ( operations for the first loop (lines 5-11) and O()) operations for the second loop (lines 15-21) in Algorithm [ Furthermore, the L' distillation iterations in the second loop can be performed by the scheduled agent after it has sent its global particles to the central server. This enables the pipelining of the second loop with the operations at the server and at other agents, which can potentially reduce the wall-clock time per communication round.
Communication load. Using DSVGD, the communication load between a scheduled agent and the central server is of the order (NN)) since N particles of dimensions d need to be exchanged at each communication round. In contrast, the communication load of PVI depends on the selected parametrization. For instance, one can use PVI with a fully factorized Gaussian approximate poste- rior, which requires only 2d parameters to be shared with the server, namely mean and variance of each of the d parameters at the price of having lower accuracy.
Convergence. The two local SVGD loops produce set of global and local particles, respectively, that are convergent their respective targets as the number N of particles increases (Liu 2017a). Furthermore, as discussed, a fixed point of the set of local free energy minimization problems is guaranteed to be a local optimum for the global free energy problem (see Property 3 in Bui et al. (0018). ) This property hence carries over to DSVGD in the limit of large number of particles. How- ever, convergence fixed point is an open question for PVI, and consequently also for DSVGD.
A.5 PARALLEL-DSVGD
In this section, we present a direct extension of DSVGD in which multiple agents can be scheduled in parallel during the same communication round. In Parallel-DSVGD (P-DSVGD), each agent in the set K(() of scheduled agents at round i applies the same steps as in DSVGD except that it shares the local particles (0()) 1 with the server instead of the global ones. Then, the server distills the received local particles into set of N server-side particles (0l 1 using SVGD to obtain the next iterate of the global posterior.
As discussed in Sec. be obtained
parallel implementation requires the i-th iterate of the global posterior to
(0) = Po())
II kEKB
(0)
II (0),
(28)
k'gKta
where (2) = (0) for k' & IC0, To replicate this same behaviour while preserving the
non-parametric property of DSVGD, in P-DSVGD, each agent h E K(O) shares its local particles
k. I representing the approximate likelihood where (( (0) N ΣN' K(B, Then, to
approximate g (0) in 28 using SVGD, the server carries out L. SVGD updates as
on df
k(ol 11 2h 1) log g ( 1--1 o!! (f) for].,........... l-l...,,bb
N
(29)
For the (ith) communication round, scheduled agents ((++1) download particles for (ollt
(++1) 1) N
[ w that are treated in similar fashion as in DSVGD. The full algorithmic table for Parallel- Distributed Stein Variational Gradient Descent (P-DSVGD) provided in Algorithm Numerical results for P-DSVGD are provided in Sec. B.3 of the Appendix.
18

---

Under review as a conference paper at ICLR 2021
A.6 PROOFS
In this section, we prove Theorem ! and 2
Theorem 1, The global posterior ()pe(() in is the unique fixed point of the DVI algorithm.
Proof. Consider the general implementation of DVI, were a set IC of agents are scheduled in parallel. DVI is equivalent to the following functional mapping
llgg (i(0)
Higk 4(0)
((:(00) exp ( ( (-2Lm(0))
(O))
KER.
(((0) (0) II (()))
(1(00) - po(0)
II 4())) II 4(0)
REX
where Z Pu(() gg (e()) Iee (0)
Therefore, assuming that all devices are periodically scheduled, g(0) is a fixed point of DVI if and only if the following equality holds
(0) (0) for k I.
This condition is satisfied by ((0) = qopt (0) and by no other distribution. This concludes the proof.
We move now to Theorem 2 for U-DSVGD. We leave the analysis of the impact of the additional distillation step used by DSVGD for future work. The analysis builds on the following result from Korba et al. (020),, which is restated here using our notation.
Denote by || 1| the norm in the RKHS H defined by the positive definite kernel k(0,0) We assume that the kernel satisfies the following technical condition: there exist a constant B > 0 such that
0k(0,-) 00,
[k(0,-)llz B and Σ
B².
(30)
7
This condition is for instance satisfied by the RBF kernel with B Zhou 2008 Furthermore, we define the kernelized Stein discrepancy (Liu et al. 2016) between two distributions and q as S(p, q), and the total variation distance as ||g pllvv
(q(0) p00)dd0.
Lemma 1. (Guaranteed per-iteration decrease of the local free energy.) (Korba et al, 2020) For kernel satisfying (30), assume that, at given communication round and local iteration with agent & scheduled, we have:
(0) is upper bounded by the maximum absolute eigenvalue of the Hessian - 2 (0) constant M > 0; and
the inequality (0) (1))) < C holds for some C >>> 0.
For learning rate e S (B 1)/(BBC2) with any B 1, the decrease in the local KL divergence
from local iteration / to + satisfies the inequality
F(g)+1((0)) ((g)))))(()))))))))))))))) (0)) S
-aeS(g ((()) e7),
(31)
where 7 = ((82 + M2B²)/2.
Lemma shows that by choosing a learning rate c < min(>-1. (B-111BB111/8881 1)/(BC one can guarantee per-iteration decrease in the local-free energy, i.e., in the KL divergence between the particles'
distribution and the target tilted distribution (()) (0) that depends on the kernelized Stein discrepancy S(g (()) the iteration before the update.
19

---

Under review as a conference paper at ICLR 2021
Lemma 2. (Relationship between global and local free energy.) The global free energy F(g(()) in is related to the local free energy Fx (g(0)) in of the k-th scheduled agent as
11-1(()) (0) exp( Lm (8))
F(a(O))= 1/ (a(0)) +
Σ mif
Eyy log
(32)
Proof. The global free energy can be written as
(())
F(g(O)) = (. log
P0( exp(
a Σ
1= Tim (0))
o 12 (0)
9(0)
<
((-1)(()
(0)
Bg log
Pk ( (0) exp( Emy Lm (0))
(33)
((0)
po(() Imzer I'm
(8-1) (0)
= =BByy( log
+ +al log
(0) (xxp(-1) E Emp Lm (0))
Pk
= = +
Σ matk
Egy log
mm
exp(- Lm
where in the second equality we have used and in the third equality we have used the equality (i-11) 1)()) (0) I m==1 (()=))()), which guaranteed by the U-DSVGD update (10) and
(see Bui et al. 2018 Property 2)).
Theorem (Guaranteed per-iteration decrease of the global free energy.). The decrease in the global free energy from local iteration I to I - 1 during communication round for which agent l is sched- uled can be lower bounded as
(16) where miax sup max (os(thnn))) L. (0)) S(q, p) denotes the Kernalized Stein Dis-
PP F (0)) F1g (0)) Z (q MI (1) - e7) -2a(K - 1) V 21(g(+)1lq)))
(()
6 mffk
crepancy between distributions q and Liu el al. 2016), and is constant depending on the
RKHS kernel and the target distribution.
We know from Lemma that learning rate min(y-1, (B 1)/(BC is sufficient to ensure a per-iteration decrease in the local free energy. Given that the KL divergence in the second term in (6 generally increases with e. demonstrates that, in order to guarantee reduction of the global free energy, a smaller learning rate may be required. We also note that the KL divergence term [[" may be explicitly related to the learning rate by following Pinder et al. 2020 Sec. 8), but we do not further pursue this aspect here. We finally remark that, in the presence of K agent, the upper bound (3I) in (Korba et al., (2200) is recovered. This is because, in the presence of one agent, the global free energy reduces to the local free energy (see (32) and accordingly U-DSVGD reduces to SVGD.
Proof. We wish to obtain an upper bound on the decrease of the global free energy ((g F(qh)(() (0)) across each local SVGD iteration during communication round Using 322 the decrease in the global free energy can be written as
F(g F(g ((g = ((())
(q 1+ (0)) - F (0))
(a)
( (0)
1 (0) [exp(- Iem (0))
+o
( log
E log
eepp(= Tm (0))
(b)
(34)
We now derive upper bounds for (a) and (b). Using Lemma 1 and the definition of the local free energy in (4) we have the following upper bound (a)
(a) F(.) "+1) F( (g (0)) s S -7))))))))))))))))))))))())))(
(35)
20

---

Under review as a conference paper at ICLR 2021
while (b) can be rewritten and upper bounded by using the properties of the total variation distance as
t (0)
(b) /0= (0)+ - (0)) log
exp( Lm (0)) d0 S
2/
I1 [++]]
q |Tv.
(36)
max
Using Pinsker's inequality (Pinkerr 1964.. the term (b) can be further upper bounded
(6) - 2
max / 22 (|
(37)
Accordingly, the global energy dissipation in (34) be upper bounded in 16
21

---

Under review as a conference paper at ICLR 2021
B ADDITIONAL EXPERIMENTS
An overview of the benchmarks considered in the experiments is provided in Table
Table 1: Overview of benchmarks used in the experiments.
Algorithm
Non paramelric
Decentralized
Inference
Stein Variational Gradient Descent (S VGDD duu Wang 2016
Yes
No
VI
Stochastic Gradient Langovin Dynamics (SGLD) Twoling & renr 2011
Yes
No
MC
Distributed Stochastic Gradient Langevin Dynamics (DSG
2014
Yes
Yes
MC
Particle Mitror Desceni (PMD) [Dai el 2016
Yes
No
VI
Partitioned Variational Inference (PVT)
2018
No
Yes
VI
Global Variational Inference (GVI) [Sala 2001
No
No
VI
Non- Parametric Variational Infernce (NPV) (Gershnai 2012
No
No
VI
Federated Averaging (FedAvg) (McMilsan al A
No
Yes
Freq.
Federated Stochastic Gradient Descent (FrUSGD) [McMaas
2017
No
Yes
Freg.
Distributed Stein Variational Gradient Deseelt (DSVGD
Yes
Yes
VI
B. 1-D MIXTURE OF GAUSSIANS TOY EXAMPLE
KL divergence belween
posterfor and agprrxxiatte
0.86
U-DVVGO
divergence between
posterior and approximale
DSVCD
U.CVVDD
0.30
PV1
0.12
GVI
SVDD
0.10
GVI
0.0
0.15
0.08
0:10
0.06
0.04
0.00
10
20
30
40
Number of rounds 8
Local Iteration number Z.
(a)
(b)
Figure 9: KI. divergence between exact and approximate global posteriors (a) as function of the number of communication rounds for I. I'' 200; and (b) function of the local iterations number L for 5.
This section complementary to the 1-D mixture Gaussians experiment Sec. of the main text.
We compare DSVGD with PVI and the counterpart centralized schemes. In Fig. (a) plot the
KL divergence between the global posterior 4opt(8 and its current approximation (8) as a function
of the number of communication rounds which corresponds to the number of communication rounds for decentralized schemes. We use = 200 particles for U-DSVGI and DSVGD with L. I. 200 local iterations. The number of SVGD iterations is fixed to 800. Gaussian prior Po (o) ( 1) is assumed in lieu of the uniform prior considered in Fig. to facilitate the implementation of PVI and conventional centralized GVI which was done following Bui al. (2018 Property 4). More specifically, we use Gaussian approximate likelihoods, i.e., (flm)
"n -1
2ns with
natural parameters 7) and < O. We observe that DSVGD has similar convergence speed as PVI, while having superior performance thanks to the reduced bias of non-parametric models. Furthermore, DSVGD exhibits the same performance U-DSVGD with the advantage of having memory requirements that do not scale with the number of iterations. Finally, both U-DSVGD and DSVGD converge to the performance of (centralized) SVGD as the number of rounds increases.
In Fig. 9bb 5 rounds for the decentralized schemes. It observed that non- non-parametric schhmes-nameds a sufficiently large number of local iterations in order to outperform the parametric strategies PVI and GVI.
we plot the same KL divergence as function of the number of local iterations L. We
SVGD and (U-)DSVGD-reque)
22

---

Under review as a conference paper at ICLR 2021
(a)
(b)
(c)
(d)
Figure 10: Performance comparison of a GVI, (b) PVI, (c) SVGD and (d) DSVGD for multi-
variate Gaussian mixture model. Solid contour lines correspond to the approximate posterior while dashed contour lines to the exact posterior (N 200,
200 and b 0.1).
B.2 2-D MIXTURE OF GAUSSIANS TOY EXAMPLE
We now consider the following 2-D mixture of Gaussians model:
(1
NH (N En) M( E²)) and (2( N(p: Eo), Es) where
[ 2
No = [0,0]; Eg =
2 4
0.226
0.1652
111.71 -1.11- E
0.1652
0.6779
2 0.5
142 0,0]; Eg =
0.5
3
0.5
[1, 0];
Eg
0.5
3
We plot in Fig. 10 the approximate posterior (()) (blue solid contour lines) and the exact posterior Tope( (0) (red dashed contour lines) for PVI, GVI, SVGD and DSVGD. We see that, as the 1-D case and in contrast parametric methods PVI and GVI, non-parametric methods SVGD and DSVGD are able to capture the different modes of the posterior, obtaining lower values for the KL divergence between the approximate and posterior.
B.3 BAYESIAN LOGISTIC REGRESSION
This section provides additional results for the Bayesian logistic regression experiment in Sec. of
the main text. In Fig. [T] we compare the performance of DSVGD (bottom row), and U-DSVGD
(top row) both with SVGD and NPV (Gershman et al. 2012) using the model described in Sec. 7
We use binary classification datasets summarized Appendix C as used in Liu & Wang 2016
and Gershman et al. (2012). We assumed N 100 particles. To ensure fairesss, iterations for SVGD. while U-DSVGD and DSVGD are executed with two agents with half of the dataset split randomly each agent. We I rounds and L L' 200 local iterations. In
used L 800
Fig. [T] we plot the accuracy and the log- ligeliioddood of the four algorithms. We observe that both U-DSVGD and DSVGD perform similarly SVGD and NPV most datasets, while allowing
23

---

Under review as a conference paper at ICLR 2021
NPV accuracy
SVGD log sikiihood
og- og-lkk toood
28
-1.
-0
U-DSVGD acauracy
U-DSVGD log-likeilhood
SVGD Hiklhood
Og Uhe food
NPW accuracy
08
08-
24
24
-4
DSVGD accuracy
DSVGD lgglkkalhood
Figure 11: Binary classification with Bayesian logistic regression results using the setting in Ger- shman et al. (2012): accuracy and log-likelihood for U-DSVGD (upper row) and DSVGD (bottom row), along with NPV and SVGD, for various datasets.
0.75
Accuracy
0.70
U.DeVED
Accuracy
0.00
DSV/GD
1540
SVGD
(1
SGLD
DSVGD (f
DSGLD
DSVGD (1-10)
Fadiyp
SWGD
080
0.00
Number of 16 37
Number of particess
(a)
(b)
Figure 12: Accuracy as function of the number of particles N for Bayesian logistic regression on the Covertype dataset: (a) comparison with various benchmarks summarized in Table [ and (b) performance for different number rounds /. L 2000 iterations were used for centralized schemes while I x L 10 x 200 total local iterations were used for decentralized schemes.
distributed implementation. We note that NPV requires computation of the Hessian matrix which is relatively impractical to compute.
We plot in Fig. 12(a) the accuracy as function of the number of particles N. DSGLD is executed with two agents, where N/2 chains per agent are ran for a trajectory of length and 500 rounds, which we have found to work best. We found that SVGD, DSVGD and U-DS VGD exhibit the same performance, which is superior to Particle Mirror Descent (PMD) and similar to SGLD and DSGLD when the number of particles increases. Fig. [2 plots the accuracy for DSVGD for the same setting for different number of communication rounds. We can see that, by increasing the number obtain similar accuracy
of particles, i.e., the communication load, of particles but with a higher number of communication rounds. For example, N = 8 with communication rounds achieves similar performance N 4 with I = 10 communication rounds.
for lower number
Fig. 13 is a complementary figure for Fig. in the main text. shows that similar conclusions based accuracy can be made when using the log- likelihooood
24

---

Under review as a conference paper at ICLR 2021
Fig. 14 shows the accuracy of DSVGD for different datasets function of the total number L of local iterations. We fix N = 6, I = 10, L = L' = 200 for U-DSVGD, DSGLD and DSVGD while L 2000 for SVGD and SGLD. We observe that U-DSVGD and DSVGD have similar performance to SVGD and that they consistently outperform other schemes for sufficiently high L.
Fig. in the main text. We note that the slightly noisy behaviour of DSVGD with K = 20 agents is attributed to the small local dataset sizes resulting from splitting the original small datasets.
Fig. T6 is complementary
Finally, Fig. 15 compares the accuracy of P-DSVGD with FedAvg and DSGLD with K = 100
agents and a proportion of 0.2 randomly scheduled agents per communication round. We that P-DSVGD exhibits similar behaviour and gain over other schemes similarly to DSVGD.
Covertypa dataset
Twonorm detssel
Ringnomm dataset
Image de beset
Lag-Skellhodd
20/22
DEVGO
03050
2010213
Number of munds
Number roundss
Number rounds
Number rounds
Log-Skellkond
28
43
001/20
52.0
D5 4
204
EOAOE 18
2000017
2010
Number of rounds: 1
Number of rounds
Number of rounds
Number of rounds
Figure 13: Bayesian logistic regression log-likelihood with K = and K 20 agents using
the setting in Gershman et al. (2012) comparing DSVGD to distributed (DSGLD) and centralized (SVGD and SGLD) schemes as function of the number of communication rounds i. We use N = particles and fix L L' 200. FedAvg has been removed as it has likeelilhhooo lower than -1 in all cases and to allow us to focus on relevant values for DSVGD.
Coverty pe dafaset
Twonorm dataset
Ringnorm dataset
Image dataset
Accuncy aganta
1SVID
9050
00025
Local borations number
Local iteration number
Local Itrations' rumber
es
Local terations number i."
(20 agento)
Accuracy
in
06600
(9100
axto
Local Borations number
Lodal iterationn" number L
Lodal terations to
rumber L"
Loddl itrationa numBer 1!
Bayesian logistic regression accuracy for K = 2 (top row) and K = 20 (bottom row) agents using the setting Gershman al. (2012) comparing U-DSVGD and DSVGD distributed (DSGLD) and centralized (SVGD and SGLD) schemes as function of the local iterations number L. We fix N particles, I (top row) and 20 (bottom row).
Figure 14:
25

---

Under review as a conference paper at ICLR 2021
Covertypa dalnee
Twonorn datanel
Ringnomm dabseet
Imaga didtnet
Accumoy
7-052C0
7-55/00
7-00/20
7000000
D5GLD
OC0LD
330
D6GLD
OC0LD
Number rounds
Numbar rounds
Nambor rounds:
Number rounds
Figure 15: Accuracy for Bayesian logistic regression with P-DSVGD, FedAvg, DSGLD and SVGD using different datasets with K 100 agents and a proportion of C 0.2 randomly scheduled agents. SVGD was executed for L = C' x 100 x 4000 iterations while we fix L = L' L, 200 total local iterations for the remaining schemes, We use = 6 particles.
Ringnorm dataset
Image dataset
Vahicle Sensors dataset
Accurnoy agentts
Fadtis
BC0LD
Fadiig
D50L0
Fathag
D00LD
Number of rounds
Number of rounds
Number of rundss
Accuracy (20 agents)
BSYGB
D50LD
Fidios
CO0L0
Number of FOuncOS
Number of rounds
Number of rounds
Figure 16: Accuracy for Bayesian logistic regression with K = 2 (top row) and K 20 (bottom row) agents under the setting in Gershman et a (012) as function of the number of communication rounds i, or number of communication rounds 6 particles, L L' 200).
B.4 BAYESIAN NEURAL NETWORKS FOR REGRESSION AND CLASSIFICATION
This part contains additional results on regression and multilabel classification experiments using Bayesian Neural Networks. Figures and 18 are complementary to Figures 6 and 7 in the main text and validate our conclusions using additional datasets for regression and the log-likelihood metric for multi-label classification.
Propulsion datas
delasel
DSVGD
B5/GB
RMGE Avarage RMSE 12 (ggents)
Feeling
Average RMSE
12 agennt))
Feakr)
DSSLD
DS0L
SVDD
SVGD
Numbar rounds
Numbar rounds CCPP dalaset
Naval Propulsion dataset
DSVDD
DSVDD
Average RMBE [20 agantss)
Feshog
Average RMBE
(20 sgannss
Fockin
Number rounds
Number rounds
Figure 17: Average Root Mean Square Error (RMSE) as function of the number of communication rounds i, or number of communication rounds, for regression using Bayesian neural networks with a single hidden layer of ReLUs under the setting HeHernnnezz-Lbaar & Adams (201)). with K (top row) and K = 20 (bottom row) agents. (N = 20, I. = 200 and 50 hidden neurons).
26

---

Under review as a conference paper at ICLR 2021
MNIST dataset
Fashion MNIST dataset
Log-llkilltod (2 sggunts))
Log Log-likellhood (2 sgente)
SGLD
$G.0
DSCLD
DSGLD
Number of rounds ;
Number of rounds
Lag-ilkathoos (20 agants)
seellleelheddd
(20 sgentt))
DS/G3
35/CD
Pedfvg
Tesking
BYCD
SYGD
50.0
DSGLD
2522.0
Number of rounds
Number of rounds 4
Figure 18: Log-likelihood for multi-label classification using Bayesian neural networks with single hidden layer of 100 neurons as function of the number of communication rounds i, number of communication rounds, using MNIST and Fashion MNIST with K (top row) and K 20
(bottom row) agents (N 20, L = L' = 200.
27

---

Under review as a conference paper at ICLR 2021
B.5
RELIABILITY PLOTS AND MAXIMUM CALIBRATION ERROR
Covortype
Twororn
MNGT detaned
Fashion MNIST
Fillyg Accurosy
Rauaoos 5 adeutios
mouuer 130
Confionee
Confidric
Confideics
Confdaric
Confidirios
Ritaica 180
realos 12
reuile 100
100
Comridenc
Corridence
Combderce
Confidance
Cordidencs
Deve0 Accurans
10
10 Parfect celrnton
courss
Fefeet calonion
calonin
Conidecca
Confidecce
Consdence
Corfidence
14
Confidanioo
Figure 19: Reliability plots for classification using Bayesian neural networks for a variable number of hidden neurons with FedAvg (top row), SVGD (middle row) and DSVGD (bottom row). We use N = 20 particles (I = 10, = L' 200 and K = 20 agents).
Covertype
Twonorm
Image dataset
Acsuracy
Accuracy
MCE
SVGD
SVGD DEVEO
SVGD
DSV3D
DSVGD
Feotwg
Fed/ig
Foshing
Number particles N
Number particles
Number of particles N
Figure 20: Accuracy and Maximum Calibration Error (MCE) as function of the number of particles N for Bayesian neural networks. We fix I 10, L L' 200 and K 20 agents in both figures.
This section provides additional results on the calibration experiment conducted in Sec. of the main text using additional datasets. In Fig. 19] we show the reliability plots for SVGD, DSVGD and FedAvg with K = 20 agents across various datasets and for different number of neurons in the hidden layer. We first note that DSVGD retains the same calibration level as SVGD across all datasets. Furthermore, while increasing the number of hidden neurons negatively affects FedAvg due to overfitting, it does not affect the trustworthiness of the predictions for the Bayesian counterparts. This is general property for Bayesian methods that contrast with frequentist approaches, for which increasing the number of parameters improves accuracy at the price of miscalibration (Guo al., 2017].
Fig. 20 plots the accuracy and MCE function of the number of particles N. While increasing N improves the accuracy (as also shown in Fig. [2] for SVGD and DSVGD, the MCE is unaffected and is lower than the MCE value for FedAvg.
28

---

Under review as a conference paper at ICLR 2021
C IMPLEMENTATION DETAILS
C.I DATASETS, BENCHIARKS AND HYPERPARAMTERS DETAILS
Datasets. We summarize in Table 2 the main parameters used across different datasets that are invariant across all experiments. The covertype datasef and the remaining binary classfication datasets that are selected from the Gunnar Raetsch's Benchmark datasets" as compiled by Mika ctt al. (199)) are used directly without normalization as Liu & Wang (2016) except for the vehicle sensors datase which is normalized by removing the mean of each feature and dividing by their standard deviations. Regression datasets are normalized by removing the mean of each feature and dividing by their standard deviations, and multi-label classification dataset II are normalized by multiplying each pixel value by 0.99/ 0.99/255 255 and adding 0.01 such that every pixel value after normal- ization belongs to the interval [0.01, 1J. All performance metrics used averaged over the number of trials. In each trial, unless specified otherwise, we permute the datasets and randomly split them across different agents.
The hyperparameters used are summarizedd in Table 3 These apply for all schemes except for DSGLD and SGLD, where the learning rates annealed and are respectively equal to ag (0.5 + 8. L + 1) 1)-0.5 and (0.5 + 1)-0.55 to ensure that they go from the order of 0.01 to 0.0001 as advised by [Weliig Teh (21T)) Og fixed according to the values in Table
Hyperparameters.
DSGLD implementation. DSGLD implemented by splitting the N particles among the K agents. More specifically, when scheduled, each agent runs [N/K] Markov chains. We assumed that response delay in addition to the trajectory length of the chains (Ahn et al. 2014) to be equal among all workers and unchanged throughout the learning process.
FedAvg implementation. FedAvg is implemented as in McMahan et al. (017)) with the only dif- ference that the server schedules a single agent at a time. Each scheduled agent performs L SGD iterations to minimize its local loss.
PVI and GVI implementation. PVI and GVI are implemented using a Gaussian parametrization for both the posterior and the prior. The natural parameters are updated via the closed form update Bui et al. (2018) Property 4).
Scheduling. Unless specified otherwise, ever, any scheduler can be used as long as it schedules one agent per communication round.
round robin scheduler to schedule agents. How-
Table 2: Overview of datasets and parameters used in the experiments. Datasets in bold are used in the experiments section of the main text.
Dataset Name
Size
Task
batchsize
trials
Train/test split
Covertype
581, 012 55
Binary classilication
100
50
817%
Twonorm
7, 7,400 20
Binary classification
10
50
80%%
Ringnorm
7,400
Rinary classification
80%,
Image
2,088 18
Binary classification
10
50
80%,
Breast Cancer
263
Rinary classification
80% /20%
Diabetis
768 x 8
Binary classification
10
50
80%; 200%
Binary classification
30
80%
Heart
270 x 13
Binary classification
10
50
80%; 200%
Waveform
0,08
Binary classification
80%
Vehicle Scnsors
2010 x 23
Binary classification
10
50
80%, 200%
Kin8am
1,192
Regression
90% (10%
Naval Propulsion
11, 91,931110 x 16
Regression
100
50
90% 110%
Combined cycle power plant (CCPP)
60.688
Regression
100
50
90% 10%%
Year Prediction
515, 315. 90
Rogression
1000
20
90% 110%
MNIST
60,000 000 785
Multi-label classification
100
20
86%,
Fushion MNIST
60,000 x 785
Multi label classification
100
20
86%, 14%
http:
wtt.:/ccpp...c.c.cee http://hwwa.a.aaproeee.eeeaaaaaa html
Bttp::/yyan..eecu..oom
com
29

---

Under review as a conference paper at ICLR 2021
Table 3: Summary of hyperparameters used across various experiments.
Hyperparameter
Regression
Binary Classification
Multi-label Classification
Ada Learning rate
0.001
0.05
0.001
smoothing fuge factor)
10
10-
10
Momentum
0.9
0.9
0.9
KDE bandwidth
0.55
0.53
0.35
Table 4: Learning rate for DSGLD and SGLD used across various datasets.
Hyperparameter
Year
MNIST
F-MNIST
Other
DSGLD de SGLD
0.0005
0.01 0.001
0.00015
0.001
0.01
C.2 SOFTWARE DETAILS
We implement all experiments in PyTorch (Paskke al. 2019) Version 10.3.1. Our experiments and code are based on the original SVGD experiments and code available at: httpps: //gitubb. com/DartML/S St inn-arratt can be easily obtained by running SVGD twice at each scheduled agent and suitably adjusting its target distribution. Our code is attached with the supplementary materials.
Sonal-Gradia -ntadiendd----a-------------e-ee entiintt--eeeee Descent) More specifically, DSVGI
All learning rates for por-mereetr particle-based benchmark schemes used are scaled by factor of 1/N match our learning rate and ensure fair comparison.
30


Reviews:
[{"review_id": "pOOSlgIgjG", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Promising but incomplete federated learning algorithm\n\nThis paper proposes a federated version of the Stein Variational Gradient Descent (SVGD) method. The general approach to perform federated learning is based on a previously published method called Partitioned Variational Inference (PVI). This work takes the PVI approach and adapts it to the SVGD framework. \n\nThe paper is in general well-written and easy to follow. Main ideas are clearly highlighted, and the technical parts are well structured and provide enough details.  The study problem is of great relevance because most of the data today is generated and stored in a distributed way. The presented approach is sound and builds on top of well-established methods. \n\nStrong points:\n\n- The paper addresses a very relevant problem by combining two well-founded approaches. \n\n- The use of particle-based variational inference methods in the context of federated learning is worth exploring. \n\n- The presented approach is rigorous and well evaluated. \n\n- The presented approach can train models with similar prediction performance than standard centralized approaches, and it's also able to produce well-calibrated predictions. \n\nWeak points:\n\n- The presented approach has limited practical use because of the current restrictions it imposes (i.e. update one agent at a time). \n\n- The convergence of the approach seems quite dubious once the current constrains of \"one agent updated at a time\" is lifted. \n\n- The implementation of this method for federated learning of large deep neural networks cast doubts in the feasibility of the approach due to the high overhead of sending/receiving multiple set of weights. \n\n\nI can not recommend the acceptation of this work for the following reasons:\n\n- The originality of method is low because it directly builds on top of two well-established approaches PVI and SVGD. However, the combination of these two approaches is not straightforward and shows how particle-based approximation methods can be also used in this challenging setting. \n\n- In my opinion, there is a relevant limitation to this approach which, although acknowledged by authors, is not properly discussed: Federated updates can only be done by one of the agents at a time, which implies that this particular algorithm is of limited practical use. One of the key points in federated learning is the possibility to exploit distributed computing infrastructure such as our mobile phones. So, updating one agent at a time practically makes unfeasible this possibility.\n\n- Another limitation is the convergence of the algorithm or, at least, an iterative improvement of the \"global free energy\". This current version of the paper guarantees that the global free energy is decreased at every round, mainly because the algorithm only updates one agent at a time and it essentially works as the standard SVGD. The proof of convergence is directly borrowed from (Korba et al., 2020). But I have strong doubts that this approach can provide this guarantee once the \"updating one agent at a time\" constraint is lifted. Because the PVI framework, which is the basis of this approach, can not guarantee a decrease of the global energy or convergence at every updating round. This is not properly discussed in the paper.  \n\n- Another relevant limitation, which is inherent to the  SVGD method, is the number of particles to be used. Each particle, in the context of the deep neural networks, corresponds to the whole set of weights of the network. So, the transmission of several set of weights can lead to very significant communication costs/delays. This is not properly/explicitly discussed in the paper. \n\n\n\nMinor comments:\nEq (17) is out of margin\n\nPost-Rebutal: I really thank the authors for their efforts  following my comments. I think they have really addressed my concerns. I therefore raise the score of the paper and recommend it for acceptance. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "dGtaXEaBWvf", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "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.\n \nThe problem of federated Bayesian learning is important, especially in the time when communication and data privacy arise public attention. The algorithm proposed is interesting and has shown benefits in the experiments.  I have a few comments as follows:\n\n1. The organization of the paper needs to be revised. Compared with PVI and SVGD, the proposed algorithm seems too long and too complicated, and the algorithm is not shown in the body but in Appendix. I think it might be better to decompose the proposed algorithms into several components, e.g., the server part and the agent part, and packing important updating steps as a procedure.\n\n2. The motivation is not clearly stated in this paper. The comparison with PVI should be extended with more details. For example, it seems to me that the authors use a different way to optimize the local free energy functional, compared with PVI.  If my understanding is wrong please correct me. But if so, what's the motivation of doing so? Why we should consider SVGD instead of natural gradient for this subproblem?\n\n3. I would also like to see more theoretical understanding of the proposed algorithm. For example, the convergence analysis for (U-)SVGD, even for the most simplified case, and compares with the existing work. Also the density evolution and the fixed-points analysis are also good to include. I do see some analysis in the appendix, but I think some results are better to be present in the main text.\n\n4. The experiments are conducted on synthetic or small datasets. I think the authors should include experiments on larger datasets and/or more complicated models.\n\n5. I'm not sure if \"Global iteration index\" is a common term but in federated learning (e.g., the FedAvg paper), it's usually called \"communication rounds\".\n\n6. A quick question: the proposed algorithm and PVI only selects ONE agent per communication round. What if more agents can be selected in a single round, like FedAvg does?", "extra_scores": {}, "extra_text": {}}, {"review_id": "Vr5jdJUD6zW", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper proposes distributed SVGD, which maintains N particles both on the server and on the client. The communication between the server and the client is conducted by uploading/downloading these N particles. The learning of local client is formulated as inferring corresponding tilted distribution. Experiments are conducted on synthetic Gaussian 1D mixture, Bayesian logistic regression on Covertype and Twonorm dataset and Bayesian NN on the UCI dataset.\n\nThe idea of federated Bayesian learning is important and well motivated. However, using DSVGD for this purpose is not well supported in the paper. Non-iid data on each client, communication cost bottleneck and limited computational resources (either storage or computation) are three characteristics of federated learning. However, this paper pays little attention to them. The performance of SVGD relies on N (Liu, 2017) and more particles are needed for higher dimensional problems (e.g., NN). This property is unsuitable for federated learning. Since DSVGD needs to transfer N particles between the server and the client in each round, and  it needs to store and compute N particle in each client. The former increases the communication cost and the latter increases the burden of client. Besides, in the experiments, training dataset is randomly split into partition of equal size among the K agents, which follows the iid setting. Non-iid dataset partition is needed to evaluate thoroughly the performance of DSVGD for federated learning. \n\nSome minor points:\n1. Sec 3 SVGD: original SVGD use the particles directly as an approximation instead of using KDE\n2. Eq. (14): the notation of $q^{(j)}$ and $q^{(i)}$ is confusing", "extra_scores": {}, "extra_text": {}}, {"review_id": "6wfUXoaRf1D", "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": "PAPER SUMMARY\n\nThis paper introduces a new approach to probabilistic federated learning, which builds on the previous PVI work of (Bui, 2018). \n\nThe proposed approach follows the same recipe in PVI where local agents learn their own model posteriors from private data, and communicate their posterior representations to a server, which aggregate local posterior representations into a universal representation. Local agents then download the aggregated posterior and offset it with their current posterior. The offsetted posterior is in turn used as the new local prior to re-run the corresponding local posterior approximation (via a generalized form of variational inference). New local posterior estimates are subsequently communicated to the server and so on. \n\nHowever, unlike PVI, the proposed method aims to replace the parametric representation of posterior with a non-parametric particle representation developed by the prior SVGD work of (Liu & Wang, 2016). This necessitate the development of a distributed particle aggregation algorithm in Section 4, which is the key contribution of this work. This development is also motivated by two practical desiderata of federated learning: (a) a good trade-off between communication load (per iteration) and no. of communication iterations; and (b) well-calibrated predictions that are more trustworthy.\n\nFollowing the above summary, I will give my opinions regarding several aspects of this paper below.\n\nNOVELTY & SIGNIFICANCE\n\nOn the high level of idea, this paper presents an interesting perspective on a practical federated learning system: communication trade-off & trustworthy prediction. These are definitely important problems in the direction of making federated learning more efficient and robust. This is the novel angle that I like about this paper.\n\nIts technical development, on the other hand, is leaning a bit more on the incremental side as the entire system is pretty much the same as that of PVI with the exception that a new particle representation is considered instead of PVI's parametric representation (in the statistical form of an exponential distribution). \n\nA common pattern here is that both representations allow universal posterior information to factorize additively across local devices (in the respective forms of local posterior representation). In both cases, this leads to a variant of a distributed sum problem where each local party has some running estimate of some piece of local information & the goal is to communicate asynchronously so that each can refine its local estimate and eventually, recover the correct sum of information. \n\nIn the case of SVGD, however, the exact local update would require buffering all previous particle representations (i.e., past estimates) to date so that the downloaded posterior can be accurately offsetted to act as a prior for the local model (i.e. independent of local data). This necessitates the development of a distillation scheme in Section 4.2 which is, to me, the key technical contribution here. In addition, the theoretical analysis on the U-DVSGD's per-iteration decrease for the KL divergence is also an interesting contribution.\n\nOn this note, it seems the authors have deferred the demonstration of how well the KDE distillation approximate the original particle representation to various places in the appendix. Perhaps putting some of those back into the main text would be better (if space permits). \n\nOn the practical aspect of this paper (i.e. communication load & trustworthy prediction), while the demonstration is sufficient against point-estimate method such as FedAvg and DSGLD, there is no comparison against other non-parametric probabilistic methods such as PVI (Bui, 2018) and/or PNFM (Yurochkin, 2019). Given that the difference between PVI and DSVGD is a matter of posterior representation, comparison against PVI is probably necessary to showcase that the particle representation yields better calibrated predictions.\n\nAlso, the probabilistic non-parametric federated learning work of (Yurochkin, 2019) also allows multiple rounds of communication (although it can also be used as a one-shot model fusion of pre-trained local models) so it would be good to also compare both the communication load & the prediction caliberation against this work.\n\nTECHNICAL SOUNDNESS\n\nI have made high-level check of the derivations and have not found any technical issues.\n\nCLARITY\n\nThe paper is very well-written, especially the part that summarizes the background on SVGD and PVI.\n\nREVIEW SUMMARY\n\nIn short, this paper presents an interesting perspective on non-parametric probabilistic federated learning via particle representation of posterior. The technical development is sufficiently novel with demonstrated practical advantages against FedAvg and DSGLD. These practical advantages however were not demonstrated against existing probabilistic non-parametric federated learning works such as PVI & PNFM -- this is perhaps strange given that DVSGD builds on PVI and is mostly different only in terms of posterior representation.\n\n--- post-rebuttal feedback ---\n\nThe authors have addressed most of my concerns. My rating for this paper therefore remains on the positive side. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work presents a distributed SVGD (DSVGD) algorithm as a new non-parametric Bayesian framework for federated learning. The reviewers concerned with the practical advantages of the proposed method, including the communication cost and the constraint of updating one agent per time. The authors rebuttal helped addressing some of the concerns, including proposing a new Parallel-DSVGD algorithm. This is very much appreciated. However, given the significant modification needed over the original version, we think it is better for the authors to further improve the work and submit to the next conference. 
Decision: Reject