Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: The paper proposes a new probabilistic regression and classification framework called "generalized variational inference in function space" (GVI-FS), which involves taking the standard function-space ELBO formulation and regularizing $\mathbb{E}_\mathbb{Q} [\log p(y|F)]$ with the Wasserstein-2 metric, $\mathcal{W}_2(\mathbb{Q}^F, \mathbb{P}^F)$, instead of the standard KL-divergence, $KL(\mathbb{Q}^F \left| \right| \mathbb{P}^F)$. This substitution is motivated by the generalized variational inference framework of Knoblauch et al. (2019) and the difficulties of working with KL divergences. The paper proposes to use Gaussian measures (GMs) to parameterize the predictive functions, although GPs with additional assumptions provide an equivalent approach. The paper proposes two specific variants of this model, a model in which both the predictive mean and the predictive covariance are parameterized by a sparse variational GP and another in which the predictive mean is parameterized by a neural network and the predictive covariance is parameterized by a sparse variational GP. The paper then presents results from regression, classification, and predictive variance-based out-of-distribution detection, comparing to existing probabilistic methods.
strengths_and_weaknesses: Strengths

1. The work is technically sound.

2. The work tackles problems of broad interest: classification, regression, and OOD detection using function-space methods. 

3. The work demonstrates substantially improved performance compared to the baseline models tested.


Weaknesses

1. The motivation for the proposed approach in terms of "generalized" variational inference is very weak. Specifically, I find the generalized variational inference framework of Knoblauch et al. (2019) to be so general as to be mostly divorced from the intentions and benefits of Bayesian inference, other than being an explicitly probabilistic method. Here are two alternative interpretations/motivations I find more convincing.

1a) The proposed method is GP regression/classification with Wasserstein-2 regularization. This is very simple and, at least to a first approximation, accurate.

1b) Another connection relates to Bayesian inference more directly. Here is a sketch. Consider the objective

$$\max_Q \mathbb{E}_{F \sim Q} [\log p(y|F) - \log q(F)]  - \lambda \mathcal{W}_2(Q,P)$$

where $\lambda \in \mathbb{R}^+$. This is a close variant of Eq. 12, only differing in the $ \mathbb{E} [- \log q(F)]$ entropy term and the $lambda$ multiplier. With an unrestricted variational family, the optimal approximate posterior in the unregularized $\lambda=0$ case is $q(F) = p(F|y)$, the true posterior.

Consider Langevin diffusion for this unregularized case. We start with a single particle $F_0$ and update in continuous time $t \geq 0$ according to

$$\text{d}F_t = -\nabla V(F_t) \text{d}t + \sqrt{2} \text{d}B_t $$

where $V(\cdot) = - \log p(y|\cdot)$ and $B_t$ denotes Brownian motion. The stationary distribution is $\propto p(F|y) \propto \exp(-V)$, the true posterior. The work of Jordan, Kinderlehrer, & Otto (1998) shows that the marginal law of this diffusion process obeys gradient flow of the functional $KL(\cdot,\pi)$ with respect to the Wasserstein-2 metric.

This implies the resulting marginal law of $F_t$ follows gradient flow of $KL(q(F), p(F|y)) = \mathbb{E}_{F \sim Q} [\log p(y|F) - \log q(F)]$, the first term in the original objective.

Now, back to the original objective, we can interpret $\lambda$ as a Lagrange multiplier, giving the dual formulation

$$\max_Q \mathbb{E}_{F \sim Q} [\log p(y|F) - \log q(F)] \quad \text{s.t. } \mathcal{W}_2(Q,P) \leq C$$

for some $C = C(\lambda) > 0$. this gives the original objective a nice interpretation. If we start the Langevin diffusion process with our variational posterior equal to the prior, $Q_0(F) = P(F)$ and stop the process once the diffusion process has covered a distance of $C$ w.r.t the Wasserstein-2 metric, then we will have approximately optimized our objective. As $C \leftarrow \infty$ (corresponding to $\lambda \leftarrow 0$), the approximate posterior $Q(F)$ approaches $p(F|y)$, the solution to the full Bayesian inference problem. This objective can then be interpreted as Bayesian inference with early stopping.

One wrinkle is that $Q(F)$ is restricted to be a GP. See the recent work of Lambert et al. (2022), "Variational inference via Wasserstein gradient flows" for a good discussion of this issue.

2. The experiments section should contain an additional comparison with a DNN with the same architecture and objective as DNN-SVGP, but with both predictive mean and uncertainty output by the network. This will give the reader a sense of how much the uncertainty handling by the sparse GP is helping performance. An additional comparison without the Wasserstein-2 regularization would also be informative. Additionally, a sparse GP method in addition to the last column in Table 1 (Exact GP) would be helpful for interpreting how good these results are.

3. (minor) l.296-8, VAEs use neural networks to parameterize variational posteriors. I would not consider this approach fundamentally different.

4. (minor) I don't believe $\Sigma$ in Eq. 18 is defined.

5. (minor) The supplied code is almost completely uncommented, making it impractical read.

questions: 1. In the experiments section, the performance numbers for baseline models are taken from previous work. Are the model architectures and training procedures for the proposed methods comparable? I didn't notice any notes on this issue.

2. An approximate Wasserstein-2 metric is given (Eq. 15). How good is this approximation in practice? In what situations is it more or less accurate. Some content along these lines would help readers understand the pros and cons of method.

3. (minor) It would be interesting to rely on the GP to influence both the predictive mean and uncertainty in the DNN-SVGP case. This could be achieved by using the NN to output "pseudo-observations" with uncertainty, which are then conditioned using the GP mean and covariance to produce a predictive mean.
limitations: Yes
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The authors propose to use the Wasserstein distance between Gaussian measures on infinite-dimensional function spaces to perform generalized variational inference over distributions over functions, with applications to Gaussian processes and Bayesian neural networks.
strengths_and_weaknesses: Strengths:

- The paper is well written and is mathematically rigorous.
- The idea of using Wasserstein inference in function spaces in this way seems novel.
- It is plausible that the proposed method might have advantages over standard variational GPs.

Weaknesses:

- The performance in the experiments seems suspiciously low and a few ablation studies seem to be missing.
- Given that the title advertises "Bayesian Deep Learning", it is unclear how one would exactly apply this to Bayesian neural networks.
questions: Major:

- If we can think of the $X_S$ like inducing points in sparse GPs, it seems like their selection should impact the performance and the estimated distances in the proposed method. Could you comment on how they are chosen? Is there any theoretical guidance on what would be a good choice?
- In the proposed GWInet, it seems that the model ultimately just uses a deterministic NN for the mean function and then wraps a normal GP around that. This seems like a rather complicated method to ultimately "just" perform GP inference. How would one use the proposed method to perform actual BNN inference in function space (that is, using a neural network with stochastic weights), so that one could actually use the learned features to capture the uncertainties instead of having to rely on a GP kernel? This seems crucial, since BNNs outperform GPs in many practical applications.
- State-of-the-art performance on FashionMNIST is 97% and on CIFAR10 around 95%, so the reported performances in Tab. 2 all seem suspiciously low. Could you comment on why that is? I think using baselines that get closer to the performances we actually see in practice would make the comparison stronger and more trustworthy.

Minor:

- l. 75: Regarding variational BNNs in function space, functional SVGD [1] and functional repulsive ensembles [2] seem to be relevant methods to mention.
- Tab. 1: I understand that both SVGP and DNN-SVGP are fit using GWI, right? If so, it seems like a comparison against a standard SVGP implementation would be interesting to see.
- Overall, it's unclear to me whether the proposed method is faster or slower than the baselines. Could you report runtimes (at least rough estimates)?


[1] https://arxiv.org/abs/2106.10760

[2] https://proceedings.neurips.cc/paper/2021/hash/1c63926ebcabda26b5cdb31b5cc91efb-Abstract.html
limitations: There seems to be no negative societal impact, but the authors could do a slightly better job at addressing the technical limitations, especially regarding runtimes, performance, and how to integrate BNNs (see above).
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: The paper presents a new methodology to apply variational inference in function space. Starting from the classic ELBO, the authors replace the KL divergence for measures with the Wasserstein distance and, by choosing Gaussian measures for the prior and the approximate posterior, the authors show how to approximate this quantity to have a tractable loss to optimize. Also, the authors present two ways to parameterize the variational posterior, one related to sparse Gaussian processes and the second based on deep neural networks. Finally, the paper is concluded with some experiments, including 1D toy regression, the classic UCI regression benchmark and some image classification problems (with OOD detection).
strengths_and_weaknesses: ## Strengths

- The paper is technically very sound and promises to improve upon standard methods for weight-space and function-space inference of neural networks (and not only).
- The paper is generally easy to follow and it is sufficiently self-contained 

## Weaknesses

- The original assumption of working with Gaussian posteriors could be a limiting factor, even if we are working in function-space. For example, a recent work [Pleiss and Cunningham, 2021] on the behaviour of wide parametric (and non) models seems to suggest that the limiting Gaussian posterior performs worse than the BNN/DGP posterior. This is common for all variational inference methods in function space with Gaussian posterior, but I would still appreciate a comment from the Authors on this. 
- The experimental evaluation is a bit limited (with nonetheless interesting results). My only concern is on the comparison with other methods, given that those numbers have been copied from various papers. Did you use the same models, same architectures, and same setup?


### Minor
- The reference section can be improved (e.g. miss-capitalizations and arxiv citations when proceedings are available).
- In the experiment, the second parameterization is sometimes referred to as GWI-DNN-SVGP and other times as GWI-Net. Maybe a uniform notation is easier to follow.

Geoff Pleiss, John Patrick Cunningham. The Limitations of Large Width in Neural Networks: A Deep Gaussian Process Perspective. NeurIPS 2021
questions: - For me, the example shown in Figure 1 is not very illustrative. Fig 1 shows only one of the proposed models with 3 different function estimation problems. For example, it would be interesting to see the comparison between GWI-SVGP and GWI-DNN-SVGP and SVGP alone (SVGP is particularly interesting to see, given that it should be the same model as GWI-SVGP trained on a different loss)
- How are you treating the prior parameters? Are they fixed, cross-validated or optimized? And what about the inducing points? Are these just a sub-sample of the training set or are they optimized?
limitations: Some additional comments on the limitations would have been welcomed (especially regarding the practical implementation). 
No comments on potential negative societal impact, as expected for this kind of work.
review_title: null

Meta-review: metareview: Technically solid paper that introduces and benchmarks a novel inference framework, with application to inference in GPs. All reviewers recommend to accept, after a decent amount of discussion in which reviewers raised their scores in response to a fairly significant round of updates to the manuscript itself. Recommend to accept, despite some questions regarding overall impact.
confidence: Certain
Decision: Accept