FUNCTION SPACE PARTICLE OPTIMIZATION FOR BAYESIAN NEURAL NETWORKS
While Bayesian neural networks (BNNs) have drawn increasing attention, their posterior inference remains challenging, due to the high-dimensional and overparameterized nature. Recently, several highly flexible and scalable variational inference procedures based on the idea of particle optimization have been proposed. These methods directly optimize a set of particles to approximate the target posterior. However, their application to BNNs often yields sub-optimal performance, as they have a particular failure mode on over-parameterized models. In this paper, we propose to solve this issue by performing particle optimization directly in the space of regression functions. We demonstrate through extensive experiments that our method successfully overcomes this issue, and outperforms strong baselines in a variety of tasks including prediction, defense against adversarial examples, and reinforcement learning.
Based on the revision, I am willing to raise the score from 5 to 7.
========================================== 
The authors address the problems of variational inference in over-parameterized models and the problem of the collapse of particle-optimization-based variational inference methods (POVI). The authors propose to solve these problems by performing POVI in the space of functions instead of the weight space and propose a heuristic approximation to POVI in function spaces.
Pros:
1) I believe that this work is of great importance to the Bayesian deep learning community, and may cause a paradigm shift in this area.
2) The method performs well in practice, and alleviates the over-parameterization problem, as shown in Appendix A.
3) It seems scalable and easy to implement (and is similar to SVGD in this regard), however, some necessary details are omitted.
Cons:
1) The paper is structured nicely, but the central part of the paper, Section 3, is written poorly; many necessary details are omitted.
2) The use of proposed approximations is not justified
In order to be able to perform POVI in function-space, the authors use 4 different approximations in succession. The authors do not check the impact of those approximations empirically, and only assess the performance of the final procedure. I believe it would be beneficial to see the impact of those approximations on simple toy tasks where function-space POVI can be performed directly. Only two approximations are well-motivated (mini-batching and approximation of the prior distribution), whereas the translation of the function-space update and the choice of mu (the distribution, from which we sample mini-batches) are stated without any details.
Major concerns:
1) As far as I understand, one can see the translation of the function-space update to the weight-space update (2) as one step of SGD for the minimization of the MSE \sum_x (f(x; \theta^i) - f^i_l(x) - \eps v(f^i_l)(x))^2, where the sum is taken over the whole space X if it is finite, or over the current mini-batch otherwise. The learning rate of such update is fixed at 1. This should be clearly stated in the paper, as for now the update (2) is given without any explanation.
2) I am concerned with the theoretical justification paragraph for the update rule (3) (mini-batching). It is clear that if each marginal is matched exactly, the full posterior is also exactly matched. However, it would usually not be possible to match all marginals using parametric approximations for f(x). Moreover, it is not clear why would updates (3) even converge at all or converge to the desired point, as it is essentially the update for an optimization problem (minimization of the MSE done by SGD with a fixed learning rate), nested into a simulation problem (function-space POVI). This paragraph provides a nice intuition to why the procedure works, but theoretical justification would require more rigor.
3) Another approximation that is left unnoted is the choice of mu (the distribution over mini-batches). It seems to me from the definition of function-space POVI that we need to use the uniform distribution over the whole object space X (or, if we do not do mini-batching, we need to use the full space X). However, the choice of X seems arbitrary. For example, for MNIST data we may consider all real-values 28x28 matrices, where all elements lie on the segment [0,1]. Or, we could use the full space R^28x28. Or, we could use only the support of the empirical distribution. I have several concerns here:
3.1) If the particles are parametric, the solution may greatly depend on the choice of X. As the empirical distribution has a finite support, it would be dominated by other points unless the data points are reweighted. And as the likelihood does not depend on the out-of-dataset samples, all particles f^i would collapse into prior, completely ignoring the training data.
3.2) If the prior is non-parametric, f(x) for all out-of-dataset objects x would collapse to the prior, whereas the f(x) for all the training objects would perfectly match the training data. Therefore we would not be able to make non-trivial predictions for the objects that are not contained in the training set unless the function-space kernel of the function-space prior somehow prevents it. This poses a question: how can we ensure the ability of our particles to interpolate and extrapolate without making them parametric? Even in the parametric case, if we have no additional regularization and flexible enough models, they could overfit and have a similar problem.
These two concerns may be wrong, as I did not fully understand how the function-space prior distribution works, and how the function-space kernel is defined (see concern 4).
4) Finally, it is not stated how the kernels for function-space POVI are defined. Therefore, it is not clear how to implement the proposed technique, and how to reproduce the results. Also, without the full expression for the weight-space update, it is difficult to relate the proposed procedure to the plain weight-space POVI with the function value kernel, discussed in Appendix B.
Minor comments:
1) It is hard to see the initial accuracy of different models from Figure 3 (accuracy without adversarial examples). Also, what is the test log-likelihood of these models?
2) It seems that sign in line 5 on page 4 should be '-'
I believe that this could be a very strong paper. Unfortunately, the paper lacks a lot of important details, and I do not think that it is ready for publication in its current form.