Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 5: Marginally below acceptance threshold
review_content: 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.
 
The 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:

1. 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.

2. 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?

3. 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.

4. The experiments are conducted on synthetic or small datasets. I think the authors should include experiments on larger datasets and/or more complicated models.

5. 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".

6. 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?
review_title: Official Blind Review #1

final_score: rating: 5: Marginally below acceptance threshold
review_content: 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.

The 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. 

Some minor points:
1. Sec 3 SVGD: original SVGD use the particles directly as an approximation instead of using KDE
2. Eq. (14): the notation of $q^{(j)}$ and $q^{(i)}$ is confusing
review_title: Seem less correlated to federated learning

final_score: rating: 6: Marginally above acceptance threshold
review_content: PAPER SUMMARY

This paper introduces a new approach to probabilistic federated learning, which builds on the previous PVI work of (Bui, 2018). 

The 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. 

However, 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.

Following the above summary, I will give my opinions regarding several aspects of this paper below.

NOVELTY & SIGNIFICANCE

On 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.

Its 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). 

A 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. 

In 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.

On 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). 

On 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.

Also, 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.

TECHNICAL SOUNDNESS

I have made high-level check of the derivations and have not found any technical issues.

CLARITY

The paper is very well-written, especially the part that summarizes the background on SVGD and PVI.

REVIEW SUMMARY

In 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.

--- post-rebuttal feedback ---

The authors have addressed most of my concerns. My rating for this paper therefore remains on the positive side. 
review_title: An interesting paper with potential practical advantages over existing works but demonstration of these advantages against some existing works seem to be missing. 

final_score: rating: 6: Marginally above acceptance threshold
review_content: Promising but incomplete federated learning algorithm

This 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. 

The 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. 

Strong points:

- The paper addresses a very relevant problem by combining two well-founded approaches. 

- The use of particle-based variational inference methods in the context of federated learning is worth exploring. 

- The presented approach is rigorous and well evaluated. 

- The presented approach can train models with similar prediction performance than standard centralized approaches, and it's also able to produce well-calibrated predictions. 

Weak points:

- The presented approach has limited practical use because of the current restrictions it imposes (i.e. update one agent at a time). 

- The convergence of the approach seems quite dubious once the current constrains of "one agent updated at a time" is lifted. 

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


I can not recommend the acceptation of this work for the following reasons:

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

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

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

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



Minor comments:
Eq (17) is out of margin

Post-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. 

review_title: Promising but incomplete federated learning algorithm

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