Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The authors propose Function space variational inference for deep Bayesian classification, an inference framework which extends ideas around evidential learning for classification by applying function space Variational Inference and marginalizing over latent functions.
Concretely, this leads to a method where the estimated Dirichlet is not directly regressed from a forward pass of a function, but is estimated as a parameter that would match the empirical distribution over sampled logits given multiple samples from the function posterior (using a nested optimization step), and hence can be plugged into the fVI framework.
The authors show this is a flexible way to model uncertainty for classification and can be combined with diverse models (ensembles, MCDropout and various Bayesian models).
The method produces good results.
main_review: Pros:

1. The paper is rigorously written, scholarship and attribution of related work is generally good, and the mathematical definitions for all used concepts and background material are abundantly available in the main paper and the appendix. This is quality material, I congratulate the authors for doing the work.

2. The key idea ,while very close to the space of papers described in appendix Section A, is a useful and interesting extension to try to utilize Bayesian estimates of the forward passes to obtain the Dirichlet parameter. While this is a narrow extension and does not introduce a better model, it is technically well justified and distinct as an inference approach.

3. In most cases the approach also seems to be empirically beneficial.

4. The authors have an unusual approach to implicitly estimate the quantity of interest using a nested Newton step, I found this quite interesting and appreciate the link to Tom Minka's old note on the Dirichlet.

Cons:
1. Primarily I want to note that the paper did not need and does not really benefit from the implicit processes connection. The key trick might actually be easier to convey and explain without that distraction in order to present the paper better. It would have been sufficient to mention in the related work section that this can be seen as an instance of this framework. Occasionally clarity suffers when trying to blend too many ideas into one paper. This paper here basically needs fVI and the rest can be derived and explained fine, Definition 3.1 does not seem to carry its weight in terms of prominence to help the reader understand the material.

2. In the otherwise well-designed experiments, I would like to see a clear demonstration of the performance of the models described in appendix Section A. It is valuable and interesting that these models have been explained as special cases of this framework here, but I would still like to see which performance can be achieved with them vis-a-vis with the fukll Bayesian estimator propsoed here to get a full idea of rthe progress this framework offers. As the evaluation currently stands, I understand that this works compared to weight space VI, but not how it works compared to the methods in Sec. A with various Bayesian models. As such, it is hard to evaluate how much benefit the novelty of this paper offers compared to the existing evidential learning-style models if the key trick with the inner Newton step to estimate the Dirichlet parameter were not applied, but for instance the models still were used in a Bayesian way.

3.
In the experiments, the authors repeatedly point out that in CIFAR 100 weight based models do better due to the strong regularisation effect of the Dirichlet prior used here when having this many classes. I appreciate that insight and hypothesis, but could we please test it? How would the performance in this case differ if the prior were to be set to be weaker (but still uniform), i.e. 1/lambda in order to be less informative. According to the authors' arguments, this should improve performance. Could we test that?

4. The authors mention that weight priors cannot obtain estimates of OOD such as the ones shown in Figure 3 middle and right. I would argue that in Hierarchical Gaussian Process Priors For Neural Network Weights by Karaletsos & Bui in Neurips 2020 explicit weight priors were used with similar results, with the key insight that these priors were data-dependent and the sampled weights as such local variables to each datapoint. There is an interesting relationship to this property here, as the Dirichlet estimate also is local to each datapoint. So maybe a softer statement could be that global weight priors may struggle with such generalizations as opposed to local priors? 

5. The authors have a mathematical hack in their method which they acknowledge: they disconnect the gradient when separately performing their inner optimization for the Dirichlet parameter and as such do not directly optimize the joint objective. I would be curious to see how the method performs if that gradient were to be carried forward using bi-level optimization (here's a review https://arxiv.org/abs/2101.11517#), as in the current inference we probably do not see the method evaluated fully. I appreciate that the authors claim it still works, but it would be instructive to get a better sense of this.


summary_of_the_review: I like the premise of this paper and it is clear that this is a work of depth and quality. However, I feel there are some details missing that I would hope can be fixed to make the work clearer and to address some of its weaknesses. 
If the authors can address the points I made about the paper (probably the first is hard to do now given the narrative of the manuscript, but the rest may be feasible) and the method is clarified to have clearly identifiable empirical benefits over the baselines in Sec. A, I would be happy to revise my opinion positively and strongly support publication.

Update post-review:
I thank the authors for their measured responses. My sense is this work could still benefit from some more iterations, such as the disconnected gradient and the somewhat dissatisfactory empirical comparisons. As such, while I still like the paper, I cannot clearly recommend acceptance because I would have these questions open as it stands. I suggest working on fixing these aspects to have a more complete version of this manuscript and hope to see it again as a reviewer when that is the case so I can strongly support it.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper introduces a function-space variational inference algorithm ("fVI") that uses a Dirichlet predictive prior, and approximates the output of a (weight-space-Bayesian) neural network as a Dirichlet distribution. Inference is achieved by viewing the network as a variational impolicit process, and optimising the functional ELBO from Sun et al., 2019. This view allows them to impose priors in function-space to many existing algorithms, such as MAP, MC Dropout, ensembles, Radial BNNs, and so on. The authors show results on a toy problem, and focus on out-of-distribution performance on larger problems such as MNIST, CIFAR10 and CIFAR100, as well as performance under a fast gradient sign method adversarial attack. The fVI variants almost always perform better when the test distribution is not the same as the training distribution. 
main_review: Firstly, this work is relevant for ICLR: designing function-space priors and inference for Bayesian neural networks is a hot topic right now. I also thought that this paper was written well, with good motivation and explanations throughout. 

I like the overall idea of taking multiple outputs from a Bayesian neural network, and instead of integrating them out to get a categorical distribution, using the samples to estimate a Dirichlet distribution (and using Minka's algorithm to estimate the parameters of the Dirichlet). 

This work is closely related to prior works such as Joo et al., 2020 and others. Appendix A summarises these works and I thought nicely summarises the differences between this method and these closest prior works. However, although theoretically there are differences between these past works and the overall framework of fVI that this paper introduces, it does not seem to me like there are significant differences in practice on large datasets. More concretely:

1. The authors say that their framework allows for index sets that are not just the training data. However, in all but one of their experiments, only the training data is used. It is nice to see the rotated MNIST example with a larger index set, however, this is quite a limited experiment (as it is only MNIST, and only tested in a very specific OOD regime). I would have liked to see more experiments where the authors used data augmentation or unlabelled data as part of the index set. Or, used some sampling similar to Sun et al., 2019, where they sample from some interval to get data not in the training set. By not having such experiments, the paper becomes similar to past work. 

2. The other difference is that fVI allows for Bayesian neural networks with M>1 samples in order to estimate the Dirichlet distribution. I like this idea. However, in practice, I was disappointed that M>1 only for ensembles and Rank1 models, and not also for MC dropout and Radial BNNs. 

Other comments / concerns: 

3. Why were z_min and z_max chosen as K and N respectively (Appendix F)? And why was z estimated to be N when M=1 (why not estimate z using Minka's algorithm)? From what I understand, as N becomes very large, the distribution would approach the same categorical that people use. N is quite large for many datasets already, so effectively this procedure is adding a very small amount of noise around what would otherwise be the categorical. I understand this reduces training cost. But it would be nice to see results for eg Radial BNNs with larger M, and also see how the estimated z compares with z=N (or z=K). (Appendix H is very short and I would like to see this discussed in more detail, including more in the main text.)

4. fVI versions of algorithms are almost always doing worse in terms of loglik and ECE on the training data. I had hoped for better 'uncertainty' performance on the in-distribution set as well. This indicates to me that the fVI version still over-estimates in-distribution data (eg MAP vs MAP fVI). The real benefit is only coming when the test distribution is not the same as the training distribution. Do the authors have some insight as to why it is not helping in-distribution? In Section 4.4 the authors say, "due to the predictive prior, the ECE was significantly higher for corruption 0". I did not understand this argument, can the authors please elaborate?

5. Did the authors try different prior parameters (rather than beta=1)? It seems like this might affect results considerably, and I would be interested to see this more (I liked the third and fourth columns of Figure 3). For example, on the toy dataset (Figure 3), what happens when using beta = [5,5] or even something like beta = [0.8,0.8]. In fact, in general, I think the ideal solution for Figure 3 for regions outside the input-data-space, would be to *slightly* predict the closest class (eg the top left corner should unconfidently predict red, and the bottom right should unconfidently predict blue). Perhaps different beta would allow for this.

6. At the end of Section 3, it would be good to add a sentence moving some of Appendix B into the main text, by mentioning the authors down-weight the KL term (and that this might be justified by referencing Wenzel et al., 2020 and Aitchison, 2021). 

Very minor points:

7. Why is label smoothing applied (Appendix F)? Is this necessary for some reason, eg for numerical stability? Why?

8. Some more work the authors could refer to: (i) Burt et al., 2021, "Understanding VI in function-space", (ii) recent SG-MCMC methods for (weight-space) Bayesian neural networks, (iii) recent natural-gradient variational inference methods for (weight-space) Bayesian neural networks.
summary_of_the_review: Overall, I like the approach of the paper, and think there is potential. Therefore, although I am borderline, I am leaning towards accept. My main review outlines my biggest issues (especially points 1-4), these tend to be about the practical implementation not aligning with what I view as the biggest benefits of the theory (and the biggest differences to past works), as well as questions about experimental results.

------------
Post-rebuttal: Please see discussion. I am sticking with my score for now, despite some of my initial concerns remaining (and additional valide points raised by Reviewer rF5E).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper considers a functional variational inference approach to Bayesian deep learning for multi-class classification.  This paradigm enables the use of a Dirichlet predictive prior in function-space, which would otherwise be cumbersome under the classical weight-space treatment. The authors approximate the posterior in function space through the broad family of variational implicit processes, which are able to subsume popular BNN paradigms such as deep ensembles and Monte Carlo dropout. Experiments are provided to support the claim that this combination of function-space regularization with a Dirichlet predictive prior improves adversarial robustness and uncertainty quantification over its weight-space counterparts.
main_review: ## Strengths

- **Significance.** This is a timely contribution that seeks to address significant limitations of Bayesian deep learning, namely those concerning the calibration of uncertainty and robustness to adversarial attacks.
- **Clarity.** The paper is generally well-written and mostly self-contained. I found it clear and easy to follow. I particularly appreciated the treatment of the prior works related to this paper, which I found to be extensive while also being succinct and clear.
- **Reproducibility.** It's excellent to see the depth of low-level implementation details that are provided in the appendix. This should greatly aid subsequent works in reproducing the results reported in this paper. It was also encouraging to see that the code is included in the supplementary material and that it seems to be clean and well-structured.

## Weaknesses

- **Novelty.** Being a combination of implicit variational processes, functional variational inference, and basic concepts from probabilistic multi-class classification,  the novelty of the contribution is somewhat limited. I find the claim that this work provides "a unifying view of prior work which use the Dirichlet distribution and function-space regularization" to be a slight over-statement. A method that combines concepts from two areas is not necessary a unifying theory thereof. One significant technical challenge that is tackled is the MC estimation and minimization of the KL divergence, which is addressed using a quasi-Newton method.

## Questions

- Section 4.4 considers image classification at five increasing levels of corruption. It is understood that each of these numbers corresponds to a particular setting of brightness, contrast, saturation, etc. However, what these settings are remains unclear. Is this detailed anywhere?
- Do you actually make use of the ability to incorporate class-biased predictive prior that the proposed approach unlocks in any experiments other than the qualitative results shown in Figure 3? More generally, are there concrete problems in which having strong informative predictive prior for multiclass classification is actually helpful, and where conventional Bayesian approaches, through posterior inference with its data-fit component are unable to adapt to the class-biases. It would be compelling to see quantitative results that support the claims concerning class-biased predictive priors.
- I am not fully convinced about the results shown in Section 4.4, in particular in Figure 5 for CIFAR100 where the accuracy and log-likelihoods are not distinctly better across corruption levels, particularly with Ensembles. Furthermore, the ECE is only lower for the highest levels of corruption. The explanation offered is that the method underfits, which is supported by the relatively low log-likelihood and accuracy. Obviously one can easily attain good ECE with OOD data by deliberately underfitting, which is exactly why the log-likelihood and accuracy are important to consider as well. Therefore, I don't understand why the following comment from the authors isn't actually done? "This suggests the function-space prior should be fine-tuned (i.e. empirical Bayes) if superior ECE is desired in-distribution." When would this not be desired?

## Miscellaneous Issues

- Pg. 2: "which use Gaussian weight priors and posteriors" → "that uses [...]"
- Pg. 4: "evidence lower bound (ELBO)" - the ELBO and, more generally, variational inference, long predates Hoffman et al. 2013.
- Pg. 4: "Further details [...], which we consider, [...]" - which → that; extraneous commas
- Pg. 8: "LLH" this abbreviation was never defined.
- The abbreviation "OOD" has even been explicitly defined. Likewise, the abbreviation "ECE" has never been defined, nor is any explanation given as to what it is. Yet it is one of the pivotal metrics considered in this line of work.
summary_of_the_review: Given the technical quality of the paper, I am leaning toward recommending acceptance of the paper. However, given some concerns around novelty and empirical evaluations detailed in my review, I am disinclined to recommend a clear acceptance.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This paper proposes a function-space variational inference method for classification tasks. The method is based on the recently proposed functional variational inference (fVI), which directly approximates the posterior over functions instead of the model parameters. The authors use a variational implicit process as a parameterized posterior and adopt Dirichlet predictive priors for the categorical predictive distributions. Since the predictive posterior is defined implicitly, the authors further propose an iterative approach to estimate the Dirichlet predictive posterior, and end up with a wake-sleep style inference procedure. Finally, the proposed method is verified on several standard classification tasks.
main_review: Summary:

See "Summary Of The Paper".

---

Reasons for score:

The idea is interesting and the paper writing is clear. However, my major concern is about the theoretical soundness of the proposed method. See my detailed comments below.

---

Pros:

- The paper is well-organized and easy to follow.
- The discussion on related work is comprehensive and systematic.
- Substantial experimental results are provided on several different important tasks (especially in the Appendix).

---

Concerns:

My main concern is about the theoretical soundness of the proposed method. In Sec. 3.2 Eq. (4), the author mentioned that _a more tractable approximation replace the supermum with an expectation_, $D_{KL}[q(f|\theta)\|p(f)]\approx E_{s\sim S} D_{KL}[q(f_s|\theta)\|p(f_s)]$. However, in the original fVI method (Sun et al., 2019), the sampling-based method samples **measurement sets**, i.e., $D_{KL}[q(f|\theta)\|p(f)]\approx E_{S\sim c} D_{KL}[q(f_S|\theta)\|p(f_S)]$, rather than **data points in the measurement set**. Note that there is a crucial difference: we will be estimating the KL divergence between the *joint* distributions induced by the variational and the prior if we sample measurement sets. In Eq. (4), however, the dependence among predictions (in both prior and variational distribution) at different input $x$ is completely ignored. This perhaps leads to meaningless results since dependence is the core of functions. Example:

Suppose that in a 10-class classification problem, we have 1000 data points located nearly at $x_A$, and the class labels are evenly distributed (that is, 100 data with label 1, 100 data with label 2,..., 100 data with label 10). Assume we also have 10 data points located at $x_B$ (1 data with label 1, 1 data with label 2,..., 1 data with label 10). Intuitively the ideal posterior predictive Dirichlet mean at $x_A$ and $x_B$ should be $(1/10, \cdots, 1/10)$, but the predictive Dirichlet precision should be different since we see more data at $x_A$. However, according to Eq. (1) and Eq. (6), the number of likelihood terms and the KL term are always equal at both $x_A$ and $x_B$, leading to equal posterior predictive Dirichlet distributions.


My second concern is about Sec. 3.3. I'm not sure whether it makes sense to estimate the distributions of $f^{(m)}_x$ with Dirichlet distribution. In the last line of Page 3, the author mentioned that *sampling the weights and evaluating $g(x,\cdot)$ is equivalent to sampling the Dirichlet predictive*. Is there any intuition that predictive distribution induced by the variational implicit process is approximately Dirichlet?




summary_of_the_review: In summary, I raise several concerns about the theoretical soundness of the proposed method, which currently prevents me from rating this paper higher. I will consider raising my evaluation if the authors address my concerns or point out my misunderstandings.
review_title: null

Meta-review: A variational function-space prior is proposed, resulting in a variational Dirichlet posterior. After rebuttal, reviewers still had many remaining questions or concerns about the paper. For instance, rF5E outlines several concerns, many relating to factorization assumptions. Reviewer 7nPR also provides several suggestions. I will not repeat them here, but do encourage the authors to look closely at these questions and suggestions. At this particular time the paper is not strongly resonating with reviewers, but could be updated so that the value of the contributions is more obvious.
Decision: Reject