Title: FUNCTION SPACE PARTICLE OPTIMIZATION FOR BAYESIAN NEURAL NETWORKS

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

1. INTRODUCTION:

Bayesian nerual networks (BNNs) provide a principled approach to reasoning about the epistemic uncertainty-uncertainty in model prediction due to the lack of knowledge. Recent work has demonstrated the potential of BNNs in safety-critical applications like medicine and autonomous driving, deep reinforcement learning, and defense against adversarial samples (see e.g. Ghosh et al., 2018; Zhang et al., 2018b; Feng et al., 2018; Smith & Gal, 2018) .

Modeling with BNNs involves placing priors on neural network weights, and performing posterior inference with the observed data. However, posterior inference for BNNs is challenging, due to the multi-modal and high dimensional nature of the posterior. Variational inference (VI) is a commonly used technique for practical approximate inference. Traditional VI methods approximate the true posterior with oversimplified distribution families like factorized Gaussians, which can severely limit the approximation quality and induce pathologies such as over-pruning (Trippe & Turner, 2018) . These limitations have motivated the recent development of implicit VI methods (Li & Turner, 2018; Shi et al., 2018b) , which allow the use of flexible approximate distributions without a tractable density. However, most of the implicit inference methods require to learn a "generator network" that maps a simple distribution to approximate the target posterior. Inclusion of such a generator network can introduce extra complexity, and may become infeasible when the number of parameters is very large, as in the case for BNNs.

Compared with those generator-based methods, particle-optimization-based variational inference (POVI) methods constitute a simpler but more efficient class of implicit VI methods. In an algorithmic perspective, POVI methods iteratively update a set of particles, so that the corresponding empirical probability measure approximates the target posterior well. Formally, these methods consider the space of probabilistic measures equipped with different metrics, and simulate a gradient flow that converges to the target distribution. Examples of POVI methods include Stein variational gradient descent (SVGD; Liu & Wang, 2016) , gradient flows in the 2-Wasserstein space (Chen et al., 2018) , and accelerated first-order methods in the 2-Wasserstein space (Liu et al., 2018) .

While POVI methods have shown promise in a variety of challenging inference problems, their performance in BNNs is still far from ideal, as with a limited number of particles, it is hard to characterize the highly complex weight-space posterior. The first problem is the curse of dimensionality: Zhuo et al. (2018) and Wang et al. (2018) show that for SVGD with a RBF kernel, particles can collapse to the maximum a posteriori (MAP) estimate as the dimension of parameter increases. One may hope to alleviate such a problem by switching to other POVI methods that could be more suitable for BNN inference; however, this is not the case: BNN is over-parameterized, and there exist a large number of local modes in the weight-space posterior that are distant from each other, yet corresponding to the same regression function. Thus a possible particle approximation is to place each particle in a different mode. In prediction, such an approximate posterior will not perform better than a single point estimate. In other words, good approximations for the weight-space posterior do not necessarily perform well in prediction.

To address the above issue, we propose to perform POVI directly for the posterior of regression functions, i.e. the function-space posterior, instead for the weight-space posterior. In our algorithm, particles correspond to regression functions. We address the infinite dimensionality of function space, by approximating the function particles by weight-space parameters, and presenting a minibatch version of particle update. Extensive experiments show that our method avoids the degenerate behavior of weight-space POVI methods, and leads to significant improvements on several tasks, including prediction, model robustness, and exploration in reinforcement learning.

The rest of our paper is organized as follows. We first briefly review BNNs and POVI in Section 2. In Section 3 we present our algorithm for function-space POVI. We compare our method with existing work in Section 4, and finally demonstrate our method's effectiveness in Section 5.

2. BACKGROUND:

Bayesian Neural Networks (BNNs) Consider a supervised learning task. Let X = {x i } N i=1 denote the training inputs and Y = {y i } N i=1 denote the corresponding outputs, with x i ∈ X and y i ∈ Y, respectively. Let f (•; θ) : X → R F denote a mapping function parameterized by a neural network, where F will be clear according to the task. Then, we can define a conditional distribution p(y|x, θ) by leveraging the flexibility of function f (x; θ). For example, for real-valued regression where Y = R, we could set F = 1 and define the conditional distribution as p(y|x, θ) = N (y|f (x; θ), σ 2 ), where σ 2 is the variance of observation noise; for a classification problem with K classes, we could set F = K and let p(y|x, θ) = Multinomial(y|softmax(f (x; θ))). A BNN model further defines a prior p(θ) over the weights θ. Given the training data (X, Y), one then infers the posterior distribution p(θ|X, Y) ∝ p(θ)p(Y|θ, X), and for a test data point x test , y is predicted to have the distribution p(y test |x test , X, Y) = p(y test |x test , θ)p(dθ|X, Y).

Posterior inference for BNNs is generally difficult due to the high-dimensionality of θ. The overparameterized nature of BNNs further exacerbates the problem: for over-parameterized models, there exist multiple θ that correspond to the same likelihood function p(y|x, θ). One could easily obtain an exponential number of such θ, by reordering the weights in the network. Each of the θ can be a mode of the posterior, which makes approximate inference particularly challenging for BNNs.

Particle-Optimization based Variational Inference (POVI) Variational inference aims to find an approximation of the true posterior. POVI methods (Liu & Wang, 2016; Chen et al., 2018) view the approximate inference task as minimizing some energy functionals over probability measures, which obtain their minimum at the true posterior. The optimization problem is then solved by simulating a corresponding gradient flow in certain metric spaces, i.e. to simulate a PDE of the form

∂ t q t = -∇ • (v • q t ),

where q t is the approximate posterior at time t, and v is the gradient flow depending on the choice of metric and energy functional. As q t can be arbitrarily flexible, it cannot be maintained exactly in simulation. Instead, POVI methods approximate it with a set of particles {θ (i) } n i=1 , i.e. q t (θ) ≈

1 n n i=1 δ(θ -θ (i)

t ), and simulate the gradient flow with a discretized version of the ODE dθ

(i) t /dt = -v(θ (i) t ).

In other words, in each iteration, we update the particles with 

EQUATION

K ij := k(θ (i) , θ (j)

) is the gram matrix. We omit the subscript for brevity. -v(θ (i) ) SVGD (Liu & Wang, 2016) (Chen et al., 2018) sum of -v in SVGD and w-SGLD-B GFSF (Liu et al., 2018) ∇

Method:

1 n n j=1 K ij ∇ θ (j) log p(θ (j) |x) + ∇ θ (j) K ij w-SGLD-B (Chen et al., 2018) ∇ θ (i) log p(θ (i) |x) + n j=1 ∇ θ (j) K ji / n k=1 K jk + n j=1 ∇ θ (j) K ji / n k=1 K ik pi-SGLD

θ (i) log p(θ (i) |x) + n j=1 (K -1 ) ij ∇ θ (j) K ij

where is the step-size at the -th iteration. Table 1 summarizes the common choices of the gradient flow v for various POVI methods. We can see that in all cases, v consists of a (possibly smoothed) log posterior gradient term, which pushes particles towards high-density regions in the posterior; and a repulsive force term (e.g. ∇ θ (j) K ij for SVGD), which prevents particles from collapsing into a single MAP estimate.

While the flexibility of POVI methods is unlimited in theory, the use of finite particles can make them un-robust in practice, especially when applied to high-dimensional and over-parameterized models. The problem of high dimensionality is investigated in Zhuo et al. (2018) and Wang et al. (2018) . Here we give an intuitive explanation of the over-parameterization problem1 : in an overparameterized model like BNN, the target posterior has a large number of modes that are sufficiently distant from each other, yet corresponding to the same regression function f . A possible convergence point for POVI with finite particles is thus to occupy all these modes, as such a configuration has a small 2-Wasserstein distance (Ambrosio et al., 2008) to the true posterior. In this case, prediction using these particles will not improve over using the MAP estimate. Such a degeneracy is actually observed in practice; see Section 5.1 and Appendix A.1.

3. FUNCTION SPACE PARTICLE OPTIMIZATION:

To address the above degeneracy issue of existing POVI methods, we present a new perspective as well as a simple yet efficient algorithm to perform posterior inference in the space of regression functions, rather than in the space of weights.

Our method is built on the insight that when we model with BNNs, there exists a map from the network weights θ to a corresponding regression function f , θ → f (•; θ), and the prior on θ implicitly defines a prior measure on the space of f , denoted as p(f ). Furthermore, the conditional distribution p(y|x, θ) also corresponds to a conditional distribution of p(y|x, f ). Therefore, posterior inference for network weights can be viewed as posterior inference for the regression function f .

A nice property of the function space inference is that it does not suffer from the overparameterization problem. However, it is hard to implement, as the function space is infinitedimensional, and the prior on it is implicitly defined. We will present a simple yet effective solution to this problem. In the sequel, we will use boldfaced symbols x (or y) to denote a subset of samples from X (or Y)2 . We denote the approximate posterior measure as q(f ). For any finite subset x, we use f (x) to denote the vector-valued evaluations of the regression function on x, and define p(f (x)), q(f (x)) accordingly. We will use the notation [n] := {1, 2, . . . , n}.

3.1. FUNCTION SPACE PARTICLE OPTIMIZATION ON A FINITE INPUT SPACE:

For the clarity of presentation, we start with a simple setting, where X is a finite set and the gradient for the (log) function-space prior, ∇ f (x) log p(f (x)) is available for any x. These assumptions will be relaxed in the subsequent section. In this case, we can treat the function values f (X ) as the parameter to be inferred, and apply POVI in this space. Namely, we maintain n particles f 1 (X ), . . . , f n (X ); and in step , update each particle with

EQUATION

This algorithm is sound when f (X ) is finite-dimensional, as theories on the consistency of POVI (e.g. Liu, 2017) directly apply. For this reason, we refer to this algorithm as the exact version of function-space POVI, even though the posterior is still approximated by particles.

However, even in the finite-X case, this algorithm can be inefficient for large X . We address this issue by approximating function-space particles in a weight space, and presenting a mini-batch version of the update rule. As we shall see, these techniques naturally generalize to the case when X is infinite.

3.1.1. PARAMETRIC APPROXIMATION TO PARTICLE FUNCTIONS:

Instead of explicitly maintaining f (x) for all x ∈ X , we can represent a function f by a parameterized neural network. Although any flexible network can be used, here we choose the original network with parameters θ of the BNN model, which can faithfully represent any function in the support of the function prior. Note that although we now turn to deal with weights θ, our method is significantly different from the existing POVI methods, as explained below in Remark 3.2 and Appendix D. Formally, our method maintains n weight-space particles θ i (i ∈ [n]) at iteration , and defines the update rule as follows:

EQUATION

where we use the shorthand

f i (•) := f (•; θ i ).

As the weights correspond to n regression functions, the rule (3) essentially updates the particles of f . In the following remarks, we relate (3) to the "exact rule" (2), and discuss its implications. Remark 3.1. ((3) as a single step of GD) The update rule (3) essentially is a one-step gradient descent (GD) to minimize the squared distance between f i +1 (X ) and f i (X ) -v[f i (X )] (the exact function-space update (2)) under the parametric representation of f . Similar strategies have been successfully used in deep reinforcement learning (Mnih et al., 2015) . Also note that (3) is easy to implement, as it closely relates to the familiar back-propagation (BP) procedure for the maximum likelihood estimation (MLE). Namely, the GD for MLE corresponds to the update θ +1 ← θ -∂f (X ;θ ) ∂θ ∇ f (X ;θ ) log p(Y|X, f (X ; θ )), where ∇ f (X ;θ ) log p(Y|X, f (X ; θ )) is commonly referred to as the "error signal of top-layer network activation" in BP. In our algorithm, this term is replaced with v. Recall that v in commonly used POVI algorithms is the sum of a possibly smoothed log posterior gradient, which is similar to ∇ f log p(Y|X, f ) used in MLE training, and the repulsive force (RF) term. Thus our algorithm can be seen as BP with a modified top-layer error signal. Remark 3.2. (Relation between (3), ensemble training, and weight-space POVI) The widely used ensemble training method (Opitz & Maclin, 1999) obtains n MAP estimates separately via GD. As stated above, our algorithm can be seen as a BP procedure with a modified top-layer error signal, thus it is closely related to ensemble training. The main difference is that our algorithm adds a RF term to the error signal, which pushes the prediction of each particle away from that of others.

As an example, consider function-space SVGD with RBF kernels. The function-space RF term is

1 n j ∇ f j k(f i , f j ) ∝ j =i (f i -f j )k(f i , f j ) (see Appendix D)

, which drives f i away from f j . Our algorithm thus enhances ensemble training, in which the predictions of all particles converge to the MAP and could suffer from overfitting. The relation to ensemble training also exists in weight-space POVI methods (Liu & Wang, 2016) . However, the RF in those methods is determined by a weight-space kernel. As discussed in Section 2, commonly used weight-space kernels cannot characterize the distance between model predictions in over-parameterized models like BNNs. In contrary, the function-space RF in our method directly accounts for the difference between model predictions, and is far more efficient. We will present empirical evidence in Section 5.1. Derivations supporting this remark are included in Appendix D. Remark 3.3. Finally, the update rule (2) corresponds to a gradient flow in a Wasserstein space of measures of functions, while it is not clear if similar results exist for rule (3). As (3) is simple to implement, and has the desirable properties discussed in Remark 3.2, we consider (3) as a reasonable approximation to the exact gradient-flow based algorithm. As we show empirically in Appendix C, it does not impact convergence; a thorough theoretical treatment is left for future work.

3.1.2. MINI-BATCH VERSION OF THE PARTICLE UPDATE:

One shortcoming of the above procedure is that it still needs to iterate over the whole set X when calculating the update rule (3), which can be inefficient for large X . We further improve the efficiency by presenting a mini-batch version, i.e., in each iteration, we draw a mini-batch with B elements, x ∼ µ, for an arbitrary distribution µ supported on X B ; we then replace the update rule (3) with evaluations on x, i.e.

EQUATION

Justification of (4) As above, we have related the full-input-space update rule (3) to (2), a gradient flow (GF) in the Wasserstein space of measures of functions (Liu et al., 2017; Chen et al., 2018; Liu et al., 2018) . Similarly, for any fixed x, the update (4) corresponds to a GF in the same space, denoted as

∂ t q = -∇ • (v x • q t ).

Simulating (4) with x sampled from µ can be seen as a stochastic approximation to simulating the "averaged GF",

∂ t q = -∇ • E x∼µ (v x • q t ).

It is known that under certain assumptions, stochastic approximation to GF simulation does not impact convergence3 , so it suffices to justify the simulation of the "averaged GF". For any fixed and finite x, it is known that the GF Chen et al. (2018) for definition of E x ).

∂ t q = -∇ • (v x • q t ) minimizes a geodesically convex energy, denoted as E x [q(f )], whose minimizer q(f ) must satisfy q(f (x)) = p(f (x)|X, Y) (see

Thus the averaged GF minimizes the averaged energy et al., 2008) . We discuss the implication of using E x∼µ E x [q(f )] as the variational objective below.

E x∼µ E x [q(f )] (Ambrosio

If the posterior process can be uniquely determined by almost all4 B-dimensional marginals {p(f (x)|X, Y) : x ∈ supp(µ)}, it will clearly become the unique minimizer of the averaged energy, and simulation of the averaged gradient flow yields the true posterior. As a concrete example, if the posterior is a Gaussian process, a sufficient condition is to set B ≥ 2, and µ to be the product measure of measures with support X , and measures with support X.5 

Even if not all marginals of the variational posterior q(f ) converge to those of the true posterior, the averaged energy serves as a good variational objective on its own. This is because it measures the average approximation error of the lower-dimensional marginals of the posterior, weighted by µ; and in most applications, we are only interested in this quantity. For example, in supervised learning tasks with i.i.d. training and test samples, we need to minimize the average approximation error of posterior predictive mean and variance. In this case, it is sufficient to obtain a good approximation of all two-dimensional marginal distributions, weighted by the input distribution; and we can set the full-support component in µ to be the kernel density estimation (KDE) of the training set.

We reiterate that the above condition on µ does not imply µ must be "uniform" on X in any sense, or each of its B marginals must be identically distributed.

Computation of (4) To implement (4) we need to compute v[f i (x)]. As shown in Table 1 , it requires the specification of a kernel and access to (the gradient of) the log posterior density, both on the B-dimensional space spanned by f (x); it also requires the specification of µ. For kernels, any positive definite kernels can be used. In our experiments, we choose the RBF kernel with the median heuristic for bandwidth, as is standard in POVI implementations (Liu & Wang, 2016) .

The log posterior gradient consists of the gradient of log prior and that of log likelihood. As in this subsection, we assume ∇ x log p(f (x)) is known, we only consider the log likelihood. As is standard in large-scale inference, we approximate it using mini-batches, i.e. to approximate it with (a scaled version of) log p(y b |x b , f i (x b )), where (x b , y b ) is a mini-batch of the training set. The requirement to sample (x b , y b ) is implemented by specifying an appropriate form of µ: we define µ in such a way that a sample from µ consists of B < B samples x b from the training set, and B -B i.i.d. samples from a continuous distribution ν over X . This is a valid choice, as stated before. Now we can use the training-set part of samples to compute the log likelihood. Finally, the continuous component ν can be chosen as the KDE of X, when the test set is identically distributed as the training set, or incorporate distributional assumptions of the test set otherwise. For example, for unsupervised domain adaptation (Ganin & Lempitsky, 2015) , we can use the KDE of the unlabeled test-domain samples.

Summing up, we present a simple yet efficient function-space POVI procedure, as outlined in Algorithm 1. As we will show empirically in Appendix C, our algorithm converges robustly in practice.

Algorithm 

∼ ν. Denote x = x b ∪{x i : i ∈ [B -B ]}. 5: For each i ∈ [n], calculate the function space POVI update v[f i (x)]. 6: For each i ∈ [n], calculate θ i

+1 according to (4).

7::

Set ← + 1. 8: end for

3.2. GENERAL SCENARIOS:

We now relax the assumptions in Section 3.1 to make our setting more practical in real applications. Below we address the infinity of X and the lack of function-space prior gradient in turn.

Infinite Set X While we assume X is a finite set to make Section 3.1.1 more easily understood, our algorithm works no matter X is finite or not: as our algorithm works with mini-batches, when X is infinite, we can also sample x from X and apply the whole procedure.

Function-Space Prior Gradient The function space prior for BNN is implicitly defined, and we do not have access to its exact gradient. While we could in principle utilize gradient estimators for implicit models (Li & Turner, 2018; Shi et al., 2018a) , we opt to use a more scalable workaround in implementation, which is to approximate the prior measure with a Gaussian process (GP). More specifically, given input x, we draw samples θ1 , . . . , θk from p(θ) and construct a multivariate normal distribution that matches the first two moments of p(f (

x)) = 1 k k j=1 δ(f (x) -f (x; θj )).

We expect this approximation to be accurate for BNNs, because under assumptions like Gaussian prior on weights, as each layer becomes infinitely wide, the prior measure determined by BNNs will converge to a GP with a composite kernel (de G. Matthews et al., 2018; Garriga-Alonso et al., 2018; Lee et al., 2018; Novak et al., 2019) .

A small batch size is needed to reduce the sample size k, as otherwise the covariance estimate in GP will have a high variance. While our procedure works for fairly small B (e.g. B ≥ 2 for a GP posterior), we choose to use separate batches of samples to estimate the gradient of the log prior and log likelihood. In this way, a much larger batch size could be used for the log likelihood estimate.

4. RELATED WORK:

Our algorithm addresses the problem of over-parameterization, or equivalently, non-identifiability 6 . A classical idea addressing non-identifiability is to introduce alternative metrics in the weight space, so that parameters corresponding to similar statistical models are closer under that metric. The typical choice is the Fisher information metric, which has been utilized to improve Markov Chain Monte Carlo methods (Girolami & Calderhead, 2011) , variational inference (Zhang et al., 2018a) and gradient descent (Amari, 1997) . While such methods explore locally non-identifiable parameter regions more rapidly than their weight-space counterparts (Amari, 2016) , they still suffer from global non-identifiability, which frequently occurs in models like Bayesian neural networks. Our work takes one step further: by defining metrics in the function space, we address local and global non-identifiability simultaneously.

Closely related to our work is the variational implicit process (VIP; Ma et al., 2018) , which shares the idea of function space inference. VIP addresses inference and model learning simultaneously; however, their inference procedure did not address the challenge of inference in complex models: the inference algorithm in VIP draws S prior functions from p(f ), and fits a Bayesian linear regression model using these functions as features. As S is limited by the computational budget, such an approximation family will have problem scaling to more complex models. We present comparisons to VIP in Appendix A.2.3.

Approximate inference for BNN is a rich field. Under the VI framework, apart from the implicit VI methods mentioned in Section 1, Louizos & Welling (2017) proposed a hierarchical variational model, which approximates p(θ|x, y) with q(θ) = q(θ|z)q(z)dz, where z represents layer-wise multiplicative noise, parameterized by normalizing flows. While this approach improves upon plain single-level variational models, its flexibility is limited by a oversimplified choice of q(θ|z). Such a trade-off between approximation quality and computational efficiency is inevitable for weight-space VI procedures. Another line of work use stochastic gradient Markov Chain Monte Carlo (SG-MCMC) for approximate inference (Li et al., 2016; Chen et al., 2014) . While SG-MCMC converges to the true posterior asymptotically, within finite time it produces correlated samples, and has been shown to be less particle-efficient than the deterministic POVI procedures (Liu & Wang, 2016; Chen et al., 2018) . Finally, there are other computationally efficient approaches to uncertainty estimation, e.g. Monte-Carlo dropout (Gal & Ghahramani, 2016) , batch normalization (Hron et al., 2018) , and efficient implementations of factorized Gaussian approximation (e.g., Blundell et al., 2015; Zhang et al., 2018a; Khan et al., 2018) .

5. EVALUATION:

In this section, we evaluate our method on a variety of tasks. First, we present a qualitative evaluation on a synthetic regression dataset. We then evaluate the predictive performance on several standard regression and classification datasets. Finally, we assess the uncertainty quality of our method on two tasks: defense against adversarial attacks, and contextual bandits.

We compare with strong baselines. For our method, we only present results implemented with SVGD for brevity (abbreviated as "f-SVGD"); results using other POVI methods are similar, and can be found in Appendix A.1 and A.2.2. Unless otherwise stated, baseline results are directly taken from the original papers, and comparisons are carried out under the same settings.

Code for the experiments will be available at https://github.com/thu-ml/fpovi. The implementation is based on ZhuSuan (Shi et al., 2017) .

5.1. SYNTHETIC DATA AND THE OVER-PARAMETERIZATION PROBLEM:

To evaluate the approximation quality of our method qualitatively, and to demonstrate the curseof-dimensionality problem encountered by weight space POVI methods, we first experiment on a simulated dataset. We follow the simulation setup in Sun et al. (2017) : for input, we randomly generate 12 data points from Uniform(0, 0.6) and 8 from Uniform(0.8, 1). The output y n for input x n is modeled as y n = x n + n + sin(4(x n + n )) + sin(13(x n + n )), where n ∼ N (0, 0.0009). The model is a feed-forward network with 2 hidden layers and ReLU activation; each hidden layer has 50 units. We use 50 particles for weight space SVGD and our method, and use Hamiltonian same model, i.e., they correspond to the same likelihood function. Global non-identifiability means there exists such θ , but not necessarily in all neighborhoods. Monte Carlo (HMC) to approximate the ground truth posterior. We plot 95% credible intervals for prediction and mean estimate, representing epistemic and aleatoric7 uncertainties respectively. Fig. 1 shows the results. We can see our method provides a reasonable approximation for epistemic uncertainty, roughly consistent with HMC; on the other hand, weight-space POVI methods severely underestimate uncertainty. Furthermore, we found that such pathology exists in all weight-space POVI methods, and amplifies as model complexity increases; eventually, all weight-space methods yield degenerated posteriors concentrating on a single function. We thus conjecture it is caused by the over-parameterization problem in weight space. See Appendix A.1 for related experiments.

5.2. PREDICTIVE PERFORMANCE:

Following previous work on Bayesian neural networks (e.g. Hernández-Lobato & Adams, 2015), we evaluate the predictive performance of our method on two sets of real-world datasets: a number of UCI datasets for real-valued regression, and the MNIST dataset for classification.

5.2.1. UCI REGRESSION DATASET:

On the UCI datasets, our experiment setup is close to Hernández-Lobato & Adams (2015) . The model is a single-layer neural network with ReLU activation and 50 hidden units, except for a larger dataset, Protein, in which we use 100 units. The only difference to Hernández-Lobato & Adams ( 2015) is that we impose an inverse-Gamma prior on the observation noise, which is also used in e.g. Shi et al. (2018b) and Liu & Wang (2016) . Detailed experiment setup are included in Appendix A.2.1, and full data for our method in Appendix A.2.2.

We compare with the original weight-space Stein variational gradient descent (SVGD), and two strong baselines in BNN inference: kernel implicit variational inference (KIVI, Shi et al., 2018b) , and variational dropout with α-divergences (Li & Gal, 2017) . The results are summarized in Fig. 2 . We can see that our method has superior performance in almost all datasets.

In addition, we compare with another two state-of-the-art methods for BNN inference: multiplicative normalizing flows and Monte-Carlo batch normalization. As the experiment setup is slightly different following (Azizpour et al., 2018) , we report the results in Appendix A.2.3. In most cases, our method also compares favorably to these baselines. Following previous work such as Blundell et al. (2015) , we report results on the MNIST handwriting digit dataset.We use a feed-forward network with two hidden layers, 400 units in each layer, and ReLU activation, and place a standard normal prior on the network weights. We choose this setting so the results are comparable with previous work.

We compare our results with vanilla SGD, Bayes-by-Backprop (Blundell et al. (2015) ), and KIVI.

For our method, we use a mini-batch size of 100, learning rate of 2 × 10 -4 and train for 1,000 epochs. We hold out the last 10,000 examples in training set for model selection. The results are summarized in Table 2 . We can see that our method outperform all baselines. 

5.3. ROBUSTNESS AGAINST ADVERSARIAL EXAMPLES:

Deep networks are vulnerable to adversarial noise, with many efficient algorithms to craft such noise (cf. e.g. Dong et al., 2018) , while defending against such noise is till a challenge (e.g. Pang et al., 2018b; a) . It is hypothesized that Bayesian models are more robust against adversarial examples due to their ability to handle epistemic uncertainty (Rawat et al., 2017; Smith & Gal, 2018) . This hypothesis is supported by Li & Gal (2017) , in a relatively easier setup with feed-forward networks on MNIST; to our knowledge, no results are reported using more flexible approximate inference techniques. In this section, we evaluate the robustness of our method on a setting compatible to previous work, as well as a more realistic setting with ResNet-32 on the CIFAR-10 dataset. We briefly introduce the experiment setup here; detailed settings are included in Appendix A.3.

On the MNIST dataset, we follow the setup in Li & Gal (2017) , and experiment with a feed-forward network. We use the iterative fast gradient sign method (I-FGSM) to construct targeted white-box attack samples. In each iteration, we limit the ∞ norm of the perturbation to 0.01 (pixel values are normalized to the range of [0, 1]). We compare our method with vanilla SVGD, and MC Dropout.

On the CIFAR-10 dataset, we use the ResNet-32 architecture (He et al., 2016a) . As dropout requires modification of the model architecture, we only compare with the single MAP estimate, and an ensemble model. We use 8 particles for our method and the ensemble baseline. We use the FGSM method to construct white-box untargeted attack samples.

Fig. 3 shows the results for both experiments. We can see our method improves robustness significantly, both when compared to previous approximate BNN models, and baselines in the more realistic setting.

5.4. IMPROVED EXPLORATION IN CONTEXTUAL BANDIT:

Finally, we evaluate the approximation quality of our method on several contextual bandit problems, Contextual bandit is a standard reinforcement learning problem. It is an arguably harder task than supervised learning for BNN approximation methods, as it requires the agent to balance between exploitation and exploration, and decisions based on poorly estimated uncertainty will lead to catastrophic performance through a feedback loop (Riquelme et al., 2018) . Problem background and experiment details are presented in Appendix A.4.

We consider the Thompson sampling algorithm with Bayesian neural networks. We use a feedforward network with 2 hidden layers and 100 ReLU units in each layer. Baselines include other approximate inference methods including Bayes-by-Backprop and vanilla SVGD, as well as other uncertainty estimation procedures including Gaussian process and frequentist bootstrap. We use the mushroom and wheel bandits from Riquelme et al. (2018) .

The cumulative regret is summarized in Table 3 . We can see that our method provides competitive performance compared to the baselines, and outperforming all baselines by a large margin in the wheel bandit, in which high-quality uncertainty estimate is especially needed. 

6. CONCLUSION:

We present a flexible approximate inference method for Bayesian regression models, building upon particle-optimization based variational inference procedures. The newly proposed method performs POVI on function spaces, which is scalable and easy to implement and overcomes the degeneracy problem in direct applications of POVI procedures. Extensive experiments demonstrate the effectiveness of our proposal.

A EXPERIMENT DETAILS AND ADDITIONAL RESULTS:

A.1 SYNTHETIC DATA

Comparison with Other POVI Methods:

We present posterior approximations obtained by weight-space and function-space versions of other POVI methods. The simulation setup is the same as in Section 5.1. As shown in Figure 4 , function-space methods provide improvement in all cases, avoid the degenerate behavior of weight-space methods.

Experiments with Increasing Model Complexity To obtain a better understanding of the degenerate behavior of weight-space POVI methods, we repeat the experiment with increasingly complex models. Specifically, we repeat the experiment while varying the number of hidden units in each layer from 5 to 100. Other settings are the same as Section 5.1. The posteriors are shown in Figure 5 . We can see that weight space methods provide accurate posterior estimates when the number of weights is small, and degenerate gradually as model complexity increases, eventually all particles collapse into a single function. On the contrary, function space methods produce stable approximations all the time. Figure 5 : Posterior approximations on increasingly complex models. L denotes the number of hidden units in each layer. We can see a clear trend of degeneration for weight-space method.

A.2 UCI DATASETS A.2.1 EXPERIMENT DETAILS IN 5.2.1

For our method in all datasets, we use the AdaM optimizer with a learning rate of 0.004. For datasets with fewer than 1000 samples, we use a batch size of 100 and train for 500 epochs. For the larger datasets, we set the batch size to 1000, and train for 1000 epochs. We use a 90-10 random traintest split, repeated for 20 times, except for Protein in which we use 5 replicas. For our method and weight-space POVI methods, we use 20 particles. We use the RBF kernel, with the bandwidth chosen by the median trick (Liu & Wang, 2016) .

For our method, we approximate the function-space prior with GP, and separate the mini-batch used for prior gradient and other parts in v, as discussed in the main text. We construct the multivariate normal prior approximation with 40 draws and a batch size of 4.

A.2.2 FULL RESULTS WITH DIFFERENT POVI PROCEDURES:

In this section, we provide full results on UCI datasets, using weight-space and function-space versions of different POVI procedures. The experiment setup is the same as Section 5.2.1. The predictive RMSE and NLL are presented in Table 4 and 5, respectively. We can see that for all POVI methods, function space variants provide substantial improvement over their weight-space counterparts. 

A.2.3 COMPARISON WITH OTHER METHODS:

In this section, we provide comparison to a few other strong baselines in BNN inference.

MNF and MCBN:

We first present comparisons with multiplicative normalizing flow (MNF; Louizos & Welling, 2017) and Monte-Carlo batch normalization (MCBN; Azizpour et al., 2018) . We make the experiment setups consistent with Azizpour et al. (2018) , and cite the baseline results from their work. The model is a BNN with 2 hidden layer, each with 50 units for the smaller datasets, and 100 for the protein dataset. We use 10% data for test, and hold out an additional fraction of 10% data to determine the optimal number of epochs. The hyperparameters and optimization scheme for our model is the same as in Section 5.2.1. The results are presented in Table 6 . We can see that in most cases, our algorithm performs better than the baselines.

VIP We present comparison to the variational implicit process (VIP; Ma et al., 2018) . We follow the experiment setup in their work, and cite results from it. The model is a BNN with 2 hidden layer, each with 10 units. For our method, we follow the setup in Section 5.2.1. We note that this is not entirely a fair comparison, as the method in Ma et al. (2018) also enables model hyper-parameter learning, while for our method we keep the hyper-parameters fixed. However, as shown in Table 7 , our method still compares even or favorably to them, outperforming them in NLL in 5 out of 9 datasets. This supports our hypothesis in Section 4 that the linear combination posterior used in Ma et al. (2018) limits their flexibility. 

A.2.4 FURTHER EXPERIMENTS:

As suggested by the reviewers, we add further experiments to evaluate our method.

Convergence Properties of Our Method We demonstrate the stability of our procedure, by presenting the negative log-likelihood on held-out data during a typical run. The experiment setup follows (Azizpour et al., 2018) , and the result is plotted in Figure 6 . As shown in the figure, the training process of our procedure is fairly stable.

Benchmark on a Narrow Architecture To balance our discussion on weight-space POVI methods, we present a benchmark on a narrower BNN model. More specifically, we use the experiment setup in (Ma et al., 2018) , which uses 10 hidden units in each layer. Other settings are summarized in the previous subsection. We compare f-SVGD with SVGD, and include the results for mean-field VI and MC dropout with α-divergence in (Ma et al., 2018) for reference. The results are presented in Table 8 . We can see that weight-space POVI is a strong baseline under priors with simpler architectures. However, as the issue of over-parameterization still presents, the function-space method still outperform it by a significant margin. 

A.3 ADVERSARIAL EXAMPLES:

MNIST Experiment Details We follow the setup in Li & Gal (2017) . We use a feed-forward network with ReLU activation and 3 hidden layers, each with 1000 units. For both SVGD and f-SVGD, we use the AdaM optimizer with learning rate 5 × 10 -4 ; we use a batch size of 1000 and train for 1000 epochs. The attack method is the targeted iterated FGSM with ∞ norm constraint, i.e. for t = 1, . . . , T , set

x (t) adv := x (t-1) adv + • sgn(∇ log E q(f ) [p(y = 0|x (t-1) adv , f )]),

where = 0.01, and x (0) adv denotes the clean input. We remove clean images labeled with "0". For our method and weight-space SVGD, we use 20 particles, and scale the output logits to have a prior variance of 10. We use an additional 40 particles to generate 8-dimensional prior samples.

A few notes on applying our method to classification problems: it is important to use the normalized logits as f , so that the model is fully identified; also as for classification problems, f (x) is already high-dimensional for a single data point x, one should down-sample f (x) within each x. We choose to always include the true label in the down-sampled version, for labeled samples, and share negative labels across different x.

Test accuracy and log likelihood on clean samples are reported in Table 9 . He et al. (2016b) , and uses the same training scheme. We use 8 particles for our method and the ensemble method. For our method, we use 32 particles to generate 6-dimensional prior samples.

The ResNet architecture consists of batch normalization layers, which needs to be computed with the full batch of input. We approximate it with a down-sampled batch, so that more prior samples can be produced more efficiently. In our implementation, we down-sample the input batch to 1/4 of its original size.

The test accuracy and log likelihood on clean samples are reported in Table 10 . Our method outperforms the single point estimate, but is slightly worse than the ensemble prediction. Performance

Figure 7 : Visualization of the wheel bandit, taken from (Riquelme et al., 2018) . Best viewed in color.

Wheel Bandit Wheel bandit is a synthetic problem that highlights the need for exploration. Let δ ∈ (0, 1) be an "exploration parameter". Context X is sampled uniformly at random in the unit circle in R 2 . There are k = 5 possible actions: the first action results in a constant reward 1 ∼ N (µ 1 , σ 2 ); the reward corresponding to other actions is determined by X:

• For contexts inside the blue circle in Figure 7 , i.e. for X s.t. X ≤ δ, the other four actions all result in a suboptimal reward N (µ 2 , σ 2 ) for µ 2 < µ 1 . • Otherwise, one of the four contexts becomes optimal depend on the quarter X is in. The optimal action results in a reward of N (µ 3 , σ 2 ) for µ 3 µ 1 , and other actions still have the reward N (µ 2 , σ 2 ).

As the probability that X corresponds to a high reward is 1 -δ 2 , the need for exploration increases as δ increases, and it is expected that algorithm with poorly calibrated uncertainty will stuck in choosing the suboptimal action a 1 in these regions. Such a hypothesis is confirmed in (Riquelme et al., 2018) , making this bandit a particularly challenging problem. In our experiments, we use 50000 contexts, and set δ = 0.95.

Experiment Setup:

The model for neural networks is a feed-forward network with 2 hidden layers, each with 100 units. The hyper-parameters for all models are tuned on the mushroom bandit, and kept the same for the wheel bandit. For the baselines, we use the hyper-parameters provided in (Riquelme et al., 2018) . The main difference from (Riquelme et al., 2018 ) is that we use 20 replicas for bootstrap. For our method, we also use 20 particles.

B ON ALTERNATIVE KERNELS FOR WEIGHT-SPACE POVI:

Many POVI methods use a kernel to make the gradient flow well-defined for discrete distributions. A natural question is whether we could design the kernel carefully to alleviate the problem of overparameterization and high dimensionality. This idea is tempting: intuitively, for most POVI methods listed in Table 1 , the kernel defines a repulsive force term, which push a particle away from its "most similar" neighbors. A better choice of the kernel makes the similarity measure more sensible. In this section, we will list candidates for such a kernel, and show that they do not work in general. Evaluation Setup For kernels considered in this section, we evaluate their empirical performance on the synthetic dataset and the UCI regression datasets, following the same setup as in the main text. We test on two POVI methods: SVGD and w-SGLD. As results are similar, we only report the results for SVGD for brevity. The results are presented in Figure 8 and Table 11 .

The "Function Value" Kernel Similar to our proposed method, one could define a weight-space kernel on function values, so it directly measures the difference of regression functions; i.e.

EQUATION

where k is an ordinary kernel on R B (e.g. the RBF kernel), x ∈ X B , and µ is an arbitrary measure supported on X B . We first remark that k f is not positive definite (p.d.) due to over-parameterization, and there is no guarantee that weight-space inference with k f will converge to the true posterior in the asymptotic limit8 . Furthermore, it does not improve over the RBF kernel empirically: predictive performance does not improve, and it drastically overestimates the epistemic uncertainty on the synthetic data.

To understand the failure of this kernel, take SVGD as an example: the update rule θ

(i) +1 ← θ (i) - v(θ (i) ) is defined with -v(θ (i) ) = 1 n n j=1   Kij∇ θ (j) log p(θ (j) |x) grad ij + ∇ θ (j) K ij rfij    .

Finite-sample SVGD is usually understood as follows (Liu & Wang, 2016) : the first term follows a smoothed gradient direction, which pushes the particles towards high-probability region in posterior;

the second term acts as a repulsive force, which prevents particles from collapse together. However, for inference with non positive definite kernels such as k f , it is less clear if these terms still play the same role:

1. When there are particles corresponding to the same regression function, their gradient for log posterior will be averaged in grad. This is clearly undesirable, as these particles can be distant to each other in the weight space, so their gradient contains little learning signal for each other.

2. While for stationary kernels, the repulsive force

rf ij = 1 n j ∇ θ (j) K ij = - 1 n j ∇ θ (i) K ij

drives particle i away from other particles, this equality does not hold for k f which is non-stationary. Furthermore, in such summation over ∇ θ

(j) K ij , as ∇ θ (j) K ij = ∇ f (x;θ (j) ) K ij ∇ θ (j) f (x; θ (j) ), assuming ∇ θ (j) f (x; θ (j)

) is of the same scale for j, ∇ θ (j) f (x; θ (j) ) will contribute to particle i's repulsive force most if the function-space repulsive force, ∇ f (x;θ (j) ) K ij is the largest. However, unlike in identifiable models, there is no guarantee that this condition imply θ (j) is sufficiently close to θ (i) , and mixing their gradient could be detrimental for the learning process.

As similar terms also exist in other POVI methods, such an argument is independent to the choice of POVI methods. Also note that both parts of this argument depends on the over-parameterization property of the model, and the degeneracy (not being p.d.) of the kernel.

The Activation Kernel Another kernel that seems appropriate for BNN inference can be defined using network layer activations. Specifically, let h(x; θ) be the activations of all layers in the NN parameterized by θ, when fed with a batch of inputs x. The kernel is defined as

k a (θ (1) , θ (2) ) := E x∼µ [k(h(x; θ (1) ), h(x; θ (2) )))].

This kernel is positive definite if the batch size B is sufficiently large, but it does not work either: predictive performance is worse than weight-space SVGD with RBF kernel, and as shown in Figure 8 , it also suffers from the collapsed prediction issue. Intuitively, our argument on over-parameterization in Section 2 should apply to all positive definite kernels.

In conclusion, constructing sensible kernels for BNN in weight-space POVI is non-trivial; on the other hand, our proposed algorithm provides a more natural solution, and works better in practice.

C IMPACT OF PARAMETRIC AND STOCHASTIC APPROXIMATIONS:

Our algorithm consists of several approximations. Assuming the function-space prior gradient is known, there remains two approximations, namely the parametric approximation to function particles (Section 3.1.1), and the stochastic approximation to particle update (Section 3.1.2). In this section, we examine the impact of them empirically, by comparing them to an exact implementation of the function-space POVI algorithm in a toy setting.

Figure 9 : Estimated KL (q p) w.r.t. training iterations, for f-SVGD using different sets of approximations. A baseline value is presented to help understanding the scale of KL divergence in this experiment; see below for details.

Experiment Setup To make exact simulation of the function-space POVI procedure possible, we consider a 1-D regression problem on a finite-dimensional function space, with a closed form GP prior. Specifically, we set X := X train ∪ X test := {-2, -1.8, -1.6, . . . , +2} ∪ {1.7, 1.9, 2.1}. The training targets are generated by y i ∼ N (sin(x i ), 0.1 2 ). The function-space prior is a GP prior with zero mean and a RBF kernel, which has a bandwidth of 0.5. The true posterior can be computed in closed form (Rasmussen, 2004):

f (X test )|X train , Y train ∼ N (µ, Σ), where µ := K tr (K rr + σ 2 I) -1 (Y train ), Σ := K tt -K tr (K rr + σ 2 I) -1 K tr , and K tt , K tr , K rr denote the gram matrices k(X test , X test ), k(X test , X train ) and k(X train , X train ), respectively. We consider three versions of f-SVGD implementation:

• An "exact" implementation, which treats f (X ) as the parameter space and directly applies SVGD. It is tractable in this experiment, as f (X ) is finite dimensional. The only approximation errors are due to discretization and particle approximation. It is known that as step-size approach 0, and number of particles approaches infinity, this method recovers the true posterior exactly (Liu et al., 2017) . • f-SVGD with the parametric approximation, i.e. the algorithm described in Section 3.1.1.

• f-SVGD with the parametric approximation, and the mini-batch approximation, i.e. the algorithm described in Section 3.1.2.

To understand the scale of the KL divergence used here, we also report its value using a baseline posterior approximation q, defined as the GP posterior conditioned on a down-sampled training set, Xtrain := {-2, -1.6, . . . , +2}.

We use 1, 000 particles and a step size of 10 -3 . We run each version of f-SVGD for 20,000 iterations, and report the exclusive KL divergence of test predictions, KL (q[f (X test )] p[f (X test )]), where q(f (•)) is approximated with a multivariate normal distribution.

The result is presented in Figure 9 . We can see that even with 1,000 particles, our introduced approximations has a negligible impact on the solution quality, compared to the (well-understood) discretization and particle approximation in the original SVGD algorithm. Furthermore, the final algorithm in Section 3.1.2 has no convergence issues.

D THE RELATION BETWEEN FUNCTION-SPACE AND WEIGHT-SPACE POVI, AND FREQUENTIST ENSEMBLE TRAINING:

Function-space POVI, weight-space POVI, and ensembled gradient descent share a similar form: they all maintain K particles in the parameter space; in each iteration, they all compute a update vector for each particle according to some rule, and add it to the parameter. In this section, we discuss this connection in more detail, by comparing the update rule for each algorithm.

Consider a real-valued regression problem, where the likelihood model is p(y|f (x)) := N (y|f (x), σ 2 y ). We choose SVGD as the base POVI algorithm. The B-dimensional kernel for f-SVGD is the RBF kernel with bandwidth σ 2 k . For all algorithms, denote the particles maintained at iteration as θ 1 , . . . , θ K . The three algorithms can all be written as

θ i +1 := θ i + u i ,

for different choices of u. We drop the subscript below, as the context is clear. Denote J i :=

∂f i (x b ) ∂θ i

, where x b , y b is the mini-batch from training set drawn in step ; denote f i (•) := f (•; θ i ).

The Ensemble Method The ensemble algorithm computes K maximum-a-posteriori estimates independently. The update rule is thus

EQUATION

≈ N B ∇ θ i log p(y b |θ i , x b ) + ∇ θ i log p(θ i ) (8) = J i N Bσ 2 y (y b -f i (x b ))

BP-ed error signal

EQUATION

where B is the batch size.

f-SVGD The update rule of f-SVGD can be derived by plugging in v from Table 1 into (4). With a RBF kernel and the Gaussian likelihood model, its form is

EQUATION

)

Denote S i := j K ij , Kij := K ij /S i , we turn (12) into

EQUATION

The similarity between ( 12) and ( 9) is now clear: the first term in ( 14) corresponds to the backpropagated error signal in (9); the second term in ( 14), the function-space prior gradient, plays the same role as the parametric prior gradient in (9); A newly added term, fRF ij directly pushes f i (x) away from f j (x), ensuring the particles eventually constitutes a posterior approximation, instead of

:

People familiar with SVGD could argue that this issue can be mitigated by choosing a reasonable kernel function in v, e.g. a kernel defined on f (•; θ). We remark that similar kernels do not work in practice. We provide detailed experiments and discussion in Appendix B.

Note x does not need to be a subset of X.

Simulation of gradient flows corresponds to optimization in the same space, so results about stochastic optimization on Riemannian manifold apply (e.g. Theorem 1 inOmbach & Tarłowski (2012)). Stochastic optimization on Riemannian manifolds is not as well-studied as their Euclidean counterparts, and assumptions introduced in existing work do not apply in our case (e.g. many results requires compactness(Ombach & Tarłowski, 2012;Bonnabel et al., 2013;Tripuraneni et al., 2018)). However, they should serve as adequate motivations of our method.

The concise condition is that for any Xs ⊂ supp(µ) s.t. µ(Xs) = 1, {p(f (x)|X, Y) : x ∈ Xs} uniquely determines the posterior process.

Recall µ is a measure on X B , i.e. a sample from µ contain B points in the input space.

Non-identifiability means there are multiple parameters that correspond to the same model. Local nonidentifiability means for any neighborhood U of parameter θ, there exists θ ∈ U s.t. θ and θ represent the

predictive uncertainty due to the noise in the data generating process.

E.g.Liu & Wang (2016) requires a p.d. kernel for such guarantees.

Figures and Tables:

-batch x b , y b from the training set, and x1...B-B i.i.d.

Figure 1: Approximate posterior obtained by different methods. Dots indicate observations, solid line indicates predicted mean, light shaded area corresponds to the predictive credible interval, and dark shaded area corresponds to the credible interval for mean estimate.

Figure 2: Average test RMSE and predictive negative log-likelihood, on UCI regression datasets. Smaller (lower) is better. Best viewed in color.

Figure 3: Accuracy on adversarial examples.

Figure 4: Posterior approximations obtained by weight space (up) and function space (down) variants of other POVI methods.

Figure 6: Heldout NLL as a function of training time (in epochs) during a sample run on the Concrete dataset.

Figure 8: Posterior approximations with weight-space SVGD and alternative kernels. k f corresponds to the function-value kernel, and k a the activation kernel. We include the results for f-SVGD and HMC for reference.

Common choices of the gradient flow v in POVI methods, where k denotes a kernel, and

Table content: <table/>

1 Function Space POVI for Bayesian Neural Network 1: Input: (Possibly approximated) function-space prior p(f (x)) for any finite x; training set (X, Y); a continuous distribution ν supported on X (e.g. the KDE of X); a choice of v from Table 1; batch size B, B ; and a set of initial particles {θ i 0 } n i=1 .

Table content: <table/>

Test error on the MNIST dataset. Boldface indicates the best result.

Table content: <table><tr><td>Method</td><td colspan="4">BBB (Gaussian Prior) BBB (Scale Mixture Prior) KIVI f-SVGD</td></tr><tr><td>Test Error</td><td>1.82%</td><td>1.36%</td><td>1.29%</td><td>1.21%</td></tr><tr><td colspan="2">5.2.2 MNIST CLASSIFICATION DATASET</td><td/><td/><td/></tr></table>

Cumulative regret in different bandits. Results are averaged over 10 trials.

Table content: <table><tr><td/><td>BBB</td><td>GP</td><td>Bootstrap</td><td>f-SVGD</td></tr><tr><td colspan="5">Mushroom 19.15 ± 5.98 16.75 ± 1.63 2.71 ± 0.22 4.39 ± 0.39</td></tr><tr><td>Wheel</td><td colspan="4">55.77 ± 8.29 60.80 ± 4.40 42.16 ± 7.80 7.54 ± 0.41</td></tr></table>

Average test RMSE on UCI datasets. Bold indicates statistically significant best results (p < 0.05 with t-test).

Table content: <table><tr><td>Dataset</td><td>SVGD</td><td>Weight Space w-SGLD</td><td>pi-SGLD</td><td>SVGD</td><td>Function Space (Ours) w-SGLD</td><td>pi-SGLD</td></tr><tr><td>Boston</td><td>2.96 ± 0.10</td><td>2.84 ± 0.15</td><td>2.84 ± 0.15</td><td>2.61 ± 0.12</td><td>2.62 ± 0.12</td><td>2.62 ± 0.12</td></tr><tr><td>Concrete</td><td>5.32 ± 0.10</td><td>5.51 ± 0.10</td><td>5.49 ± 0.10</td><td>4.73 ± 0.13</td><td>4.78 ± 0.14</td><td>4.77 ± 0.13</td></tr><tr><td>Kin8nm</td><td>0.09 ± 0.00</td><td>0.07 ± 0.00</td><td>0.07 ± 0.00</td><td>0.07 ± 0.00</td><td>0.07 ± 0.00</td><td>0.07 ± 0.00</td></tr><tr><td>Naval</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td></tr><tr><td>Power</td><td>3.94 ± 0.03</td><td>3.97 ± 0.03</td><td>3.97 ± 0.03</td><td>3.80 ± 0.03</td><td>3.83 ± 0.03</td><td>3.83 ± 0.03</td></tr><tr><td>Protein</td><td>4.61 ± 0.01</td><td>4.40 ± 0.02</td><td>4.40 ± 0.02</td><td>3.96 ± 0.03</td><td>4.02 ± 0.03</td><td>4.02 ± 0.03</td></tr><tr><td>Winered</td><td>0.61 ± 0.01</td><td>0.63 ± 0.01</td><td>0.63 ± 0.01</td><td>0.61 ± 0.01</td><td>0.61 ± 0.01</td><td>0.61 ± 0.01</td></tr><tr><td>Yacht</td><td>0.86 ± 0.05</td><td>0.95 ± 0.07</td><td>0.93 ± 0.07</td><td>0.63 ± 0.06</td><td>0.63 ± 0.06</td><td>0.63 ± 0.06</td></tr></table>

Average test NLL on UCI datasets. Bold indicates best results.

Table content: <table><tr><td>Dataset</td><td>SVGD</td><td>Weight Space w-SGLD</td><td>pi-SGLD</td><td>SVGD</td><td>Function Space (Ours) w-SGLD</td><td>pi-SGLD</td></tr><tr><td>Boston</td><td>2.50 ± 0.03</td><td>2.52 ± 0.07</td><td>2.52 ± 0.07</td><td>2.47 ± 0.08</td><td>2.52 ± 0.09</td><td>2.51 ± 0.08</td></tr><tr><td>Concrete</td><td>3.08 ± 0.02</td><td>3.15 ± 0.03</td><td>3.15 ± 0.03</td><td>2.96 ± 0.04</td><td>3.02 ± 0.04</td><td>3.01 ± 0.04</td></tr><tr><td>Kin8nm</td><td>-0.98 ± 0.01</td><td>-1.20 ± 0.01</td><td>-1.20 ± 0.01</td><td>-1.24 ± 0.00</td><td>-1.25 ± 0.00</td><td>-1.25 ± 0.00</td></tr><tr><td>Naval</td><td>-4.09 ± 0.01</td><td>-6.33 ± 0.02</td><td>-6.39 ± 0.02</td><td>-5.92 ± 0.03</td><td>-6.43 ± 0.04</td><td>-6.45 ± 0.05</td></tr><tr><td>Power</td><td>2.79 ± 0.01</td><td>2.80 ± 0.01</td><td>2.80 ± 0.01</td><td>2.76 ± 0.01</td><td>2.76 ± 0.01</td><td>2.76 ± 0.01</td></tr><tr><td>Protein</td><td>2.95 ± 0.00</td><td>2.90 ± 0.00</td><td>2.90 ± 0.00</td><td>2.79 ± 0.01</td><td>2.81 ± 0.01</td><td>2.81 ± 0.01</td></tr><tr><td>Winered</td><td>0.93 ± 0.01</td><td>0.97 ± 0.01</td><td>0.97 ± 0.01</td><td>0.92 ± 0.02</td><td>0.93 ± 0.02</td><td>0.94 ± 0.02</td></tr><tr><td>Yacht</td><td>1.23 ± 0.04</td><td>1.39 ± 0.07</td><td>1.37 ± 0.07</td><td>1.01 ± 0.06</td><td>0.99 ± 0.06</td><td>1.00 ± 0.06</td></tr></table>

Average test RMSE and NLL following the setup ofAzizpour et al. (2018). Bold indicates best results.

Table content: <table><tr><td>Dataset</td><td>MCBN</td><td>Test RMSE MNF</td><td>f-SVGD</td><td>MCBN</td><td>Test NLL MNF</td><td>f-SVGD</td></tr><tr><td>Boston</td><td>2.75 ± 0.05</td><td>2.98 ± 0.06</td><td>2.95 ± 0.08</td><td>2.38 ± 0.02</td><td>2.51 ± 0.06</td><td>2.57 ± 0.07</td></tr><tr><td>Concrete</td><td>4.78 ± 0.09</td><td>6.57 ± 0.04</td><td>5.03 ± 0.10</td><td>3.45 ± 0.11</td><td>3.35 ± 0.04</td><td>3.16 ± 0.05</td></tr><tr><td>Kin8nm</td><td>0.07 ± 0.00</td><td>0.09 ± 0.00</td><td>0.06 ± 0.00</td><td>-1.21 ± 0.01</td><td>-1.04 ± 0.00</td><td>-1.33 ± 0.00</td></tr><tr><td>Power</td><td>3.74 ± 0.01</td><td>4.19 ± 0.01</td><td>3.50 ± 0.01</td><td>2.75 ± 0.00</td><td>2.86 ± 0.01</td><td>2.67 ± 0.00</td></tr><tr><td>Protein</td><td>3.66 ± 0.01</td><td>4.10 ± 0.01</td><td>3.36 ± 0.03</td><td>2.73 ± 0.00</td><td>2.83 ± 0.01</td><td>2.56 ± 0.01</td></tr><tr><td>Winered</td><td>0.62 ± 0.00</td><td>0.61 ± 0.00</td><td>0.63 ± 0.01</td><td>0.95 ± 0.01</td><td>0.93 ± 0.00</td><td>0.98 ± 0.01</td></tr><tr><td>Yacht</td><td>1.23 ± 0.05</td><td>2.13 ± 0.05</td><td>0.85 ± 0.06</td><td>1.39 ± 0.03</td><td>1.96 ± 0.05</td><td>1.03 ± 0.03</td></tr></table>

Average test RMSE and NLL on UCI datasets, following the setup inMa et al. (2018). Bold indicates best results.

Table content: <table><tr><td/><td>NLL</td><td/><td colspan="2">RMSE</td></tr><tr><td/><td>VIP-BNN</td><td>f-SVGD</td><td>VIP-BNN</td><td>f-SVGD</td></tr><tr><td>Boston</td><td>2.45 ± 0.04</td><td>2.33 ± 0.05</td><td>2.88 ± 0.14</td><td>2.58 ± 0.12</td></tr><tr><td>Concrete</td><td>3.02 ± 0.02</td><td>2.93 ± 0.02</td><td>4.81 ± 0.13</td><td>4.63 ± 0.12</td></tr><tr><td>Energy</td><td>0.60 ± 0.03</td><td>0.69 ± 0.03</td><td>0.45 ± 0.01</td><td>0.43 ± 0.01</td></tr><tr><td>Kin8nm</td><td>-1.12 ± 0.01</td><td>-1.11 ± 0.01</td><td>0.07 ± 0.00</td><td>0.07 ± 0.00</td></tr><tr><td>Power</td><td>2.92 ± 0.00</td><td>2.76 ± 0.00</td><td>4.11 ± 0.05</td><td>3.78 ± 0.02</td></tr><tr><td>Protein</td><td>2.87 ± 0.00</td><td>2.85 ± 0.00</td><td>4.25 ± 0.07</td><td>3.95 ± 0.02</td></tr><tr><td>Winered</td><td>0.97 ± 0.02</td><td>0.89 ± 0.01</td><td>0.64 ± 0.01</td><td>0.61 ± 0.01</td></tr><tr><td>Yacht</td><td>-0.02 ± 0.07</td><td>1.17 ± 0.01</td><td>0.32 ± 0.06</td><td>0.67 ± 0.07</td></tr><tr><td>Naval</td><td>-5.62 ± 0.04</td><td>-5.41 ± 0.10</td><td>0.00 ± 0.00</td><td>0.00 ± 0.00</td></tr></table>

Comparison with other baselines under a narrower network prior on some UCI datasets. boldface indicates best results.

Table content: <table><tr><td/><td/><td>NLL</td><td/><td/><td/><td>RMSE</td><td/><td/></tr><tr><td/><td>MFVI</td><td>alpha-dropout</td><td>SVGD</td><td>f-SVGD</td><td>MFVI</td><td>alpha-dropout</td><td>SVGD</td><td>f-SVGD</td></tr><tr><td>Boston</td><td>2.76 (0.04)</td><td>2.45 (0.02)</td><td>2.42 (0.07)</td><td>2.33 (0.05)</td><td>3.85 (0.22)</td><td>3.06 (0.09)</td><td>2.77 (0.20)</td><td>2.58 (0.12)</td></tr><tr><td>Kin8mn</td><td>-0.81 (0.01)</td><td>-0.92 (0.02)</td><td>-1.11 (0.01)</td><td>-1.11 (0.00)</td><td>0.10 (0.00)</td><td>0.10 (0.00)</td><td>0.08 (0.00)</td><td>0.08 (0.00)</td></tr><tr><td>Power</td><td>2.83 (0.01)</td><td>2.81 (0.00)</td><td>2.79 (0.01)</td><td>2.76 (0.00)</td><td>4.11 (0.04)</td><td>4.08 (0.00)</td><td>3.95 (0.02)</td><td>3.78 (0.02)</td></tr><tr><td>Protein</td><td>3.00 (0.00)</td><td>2.90 (0.00)</td><td>2.87 (0.00)</td><td>2.85 (0.00)</td><td>4.88 (0.04)</td><td>4.46 (0.00)</td><td>4.29 (0.02)</td><td>3.95 (0.02)</td></tr></table>

MNIST: Test accuracy and NLL on clean samples

Table content: <table><tr><td/><td colspan="3">MC-dropout SVGD f-SVGD</td></tr><tr><td>accuracy</td><td>0.983</td><td>0.970</td><td>0.984</td></tr><tr><td>average NLL</td><td>0.075</td><td>0.109</td><td>0.065</td></tr><tr><td colspan="4">CIFAR Experiment Details We use the ResNet-32 architecture, defined in</td></tr></table>

Test NLL on UCI datasets for SVGD with alternative kernels.

Table content: <table><tr><td>RBF</td><td>Weight-Space SVGD ka</td><td>k</td><td>Function-Space SVGD</td></tr></table>

Peer review (iclr2019):
 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.