Regularized KL-Divergence for well-defined function space variational inference in BNNs

We thank the reviewer for their time and informative feedback. We are encouraged that they found our method well motivated, simple and that it provides good experimental results. We are also glad they thought the background and issues with function-space BNNs were well presented.

The review has led us to realize the need for some clarifications in the paper, and we will upload a revised version in the coming days. In the meantime, please find our answers and comments below.

Although the idea of introducing generalized KL divergence may be reasonable, the final objective is without much difference from the previous one.

We realize that the differences to previous work require some clarification in our paper. We invite the reviewer to read items 1 and 3 in the overall response (“Importance of function-space priors and fundamental differences to TFSVI and FVI“) which highlights how our method differs from its closest baselines TFSVI (Rudner et al., 2022) and FVI (Sun et al., 2019).

In brief, our method allows for informative priors in function space whereas TFSVI requires priors to be specified in weight space. This may not be obvious since TFSVI pushes the weight-space prior forward into function space, but the choice of priors in TFSVI is still limited to distributions in weight space where meaningful belief specification is difficult e.g. see Figure 4 in the appendix. For details, see item 1 in the overall response.

Compared to FVI, our method uses the linearized BNN and avoids score function estimators (see Issue (i) in Section 2.2), which have been reported to make FVI difficult to use in practice (Ma et Hernández-Lobato 2021). From a more theoretical point of view, we show that, within the linearized BNN, the KL-divergence estimator of FVI (which introduces noise in a somewhat ad-hoc way) can be well-motivated and interpreted as an estimator of a well-defined regularized KL-divergence. See item 3 in our overall response for more details.

During the practical implementation of previous functional BNNs, we usually add a diagonal identity matrix to ensure the non-singular the covariance matrix in (2) or (11)

We agree that existing methods also introduce such jitter, and that it is somewhat subtle how the role of jitter differs across methods. We invite the reviewer to read items 2 and 3 in our overall response, which discusses these differences between our method, TFSVI, and FVI.

In summary, the jitter introduced in the implementation of TFSVI only addresses numerical issues that would not arise in the absence of rounding errors (see Issue (iii) in Section 2.2). By contrast, the jitter in our method does not solve numerical issues but instead is fundamentally necessary as we calculate a divergence between measures with different support, which would be infinite without the jitter.

FVI also adds jitter to the variational posterior and prior when estimating the KL divergence. This jitter serves a similar purpose as in our method. However, since FVI does not linearize the BNN, it does not have access to an explicit variational measure in function space. Therefore, in order to take jitter into account, FVI has to draw Gaussian noise. By contrast in the linearized BNN, the marginalization over jitter can be done analytically as the variational posterior is a Gaussian measure.

In the experiments, there are no results to show the cases where the non-singular covariance matrix impacts greatly the results.

We are not sure to understand the reviewer’s question. If the reviewer asks why we present results for TFSVI and FVI despite explaining that these methods do not work without jitter (see Section 2.2 issue (iii)), it is because our implementations are based on the code provided by their respective authors which add jitter in the KL divergence estimators. If the reviewer asks for additional experiments where we use our method, TFSVI or FVI without jitter, then we would like to point out that training such models would be impossible as we would immediately run into singularities.

Our experiments show that our method is better than FVI and TFSVI at capturing prior beliefs (see, e.g., Figure 2, and Figure 4 in the appendix), and that it leads to better predictive performance and out-of-distribution detection (Tables 1 and 3).

The authors claim that “VI is too restrictive for BNNs with informative function space priors”

After consideration, we agree with the reviewer that this point is a bit weak as a contribution. Indeed, our argument for this claim builds entirely on the results by Burt et al. (2020). We will remove this claim for the paper. The main contribution of our paper is not theoretical, but instead of a practical and well-defined method that empirically performs well.

We hope that our response clarifies how our contribution differs from prior work, and that we were able to convince you that our simple method could be helpful to members of the community.

We thank the reviewer for their time and informative feedback. We are encouraged that they thought that the regularized KL divergence was a good contribution for the community and that they found our method to perform well on regression tasks. We are also glad they thought the paper to be well written and easy to follow.

The review has led us to realize the need for some clarifications in the paper that we will make in the coming days. In the meantime, please find our answers and comments below.

the only contribution [...] is to bring the regularized KL divergence [...] as a substitute for the [...] KL divergence in the framework of [...] TFVI (Rudner et al., 2022).

We kindly refer to the reviewer to item 1 and 3 which highlights how our method differs from TFSVI and FVI (Sun et al., 2019) in the overall response (“Importance of function-space priors and fundamental differences to TFSVI and FVI“).

In summary, our main contribution is a simple method for inference in BNNs with function-space priors. Unlike our method, TFSVI defines prior beliefs in weight space using standard BNN priors, but pushes it to function-space to compute the KL divergence. Thus priors in TFSVI are still limited to distributions in weight space where meaningful belief specification is difficult e.g. see Figure 4 in the appendix. For details, see item 1 in the overall response.

Compared to FVI, our method uses the linearized BNN and avoids score function estimators (see Issue (i) in Section 2.2), which have been reported to make FVI difficult to use (Ma et Hernández-Lobato 2021). From a more theoretical point of view, we show that, within the linearized BNN, the KL-divergence estimator of FVI (which introduces noise in a somewhat ad-hoc way) can be well-motivated and interpreted as an estimator of a well-defined regularized KL-divergence. See item 3 in our overall response for more details.

While I appreciate the experiments on the regression tasks, they are relatively small-scale tasks, and only the small BNNs (MLPs mostly) are tested.

Indeed, the BNNs we consider are relatively small as is typical in the literature on function space VI for BNNs¹. Computing the regularized KL divergence can be expensive for large models. Thus, in its current state, we believe that our method is best suited for models in a regime in-between analytically solvable and deep learning models, which find many practical applications involving decision making, such as Bayesian optimization or bandits. At the same time, we believe that one of the strengths of our method is its simplicity, which would make it a solid starting point for future research, including ways to make it more scalable.

¹ note that Rudner et al. (2022) reports results using a Resnet-18 on a Nvidia V-100 GPU with 32GB of memory. It should be possible to fit this model with our method using similar resources. However our method offers little advantages for classification tasks, as discussed in the next answer below.

It is hard to judge the effectiveness of the proposed method without [...] the image classification task (CIFAR-10 or CIFAR-100, at least) solved with ResNet, as typically done in the literature.

We thank the reviewer for raising this point on which we did not expand on in the paper. Our main contribution is the use of GP priors for BNNs to incorporate informative knowledge about the functions generated by the BNN. Prior beliefs in the classification setting are defined in logit-space before the output activation which are much harder to specify and puts into doubt our original motivation. This is why we preferred to restrict our method to regression tasks. However, automatic prior selection methods with our method might be beneficial for classification and would be an interesting direction for future work.

There is a class of inference algorithms (generalized Bayesian inference) [...] can be described in the perspective of generalization error. However, [...] the [...] objective with the regularized KL does not explain anything about its optimum. I think there should be some intuitive justification for this.

Thank you for raising this point. We are not sure if we understand the question. As discussed at the beginning of Section 3, our method builds on generalized VI (Kloblauch et al., 2018), which views Bayesian inference as regularized empirical risk minimization by considering the ELBO as a regularized expected log-likelihood objective. While we summarize the main idea behind generalized VI at the beginning of Section 3, we think that adding more details about generalized VI would be beyond the scope of our paper, and we instead refer readers to the original paper. We hope this answers the question and we are happy to follow up if we misunderstood it.

We hope that our response clarifies how our contribution differs from prior work, and that we were able to convince you that our simple method could be helpful to members of the community.

We thank the reviewer for their useful and informative feedback on our paper. Please find our responses below.

• “I could not identify significant differences between the tractable function space VI of [1] and the proposed method. What is the primary distinction in comparison to [1]?”
• “Why is TFSVI categorized under weight space priors?”
• “In comparison to [1], what specific difference in the proposed method leads to the improved performance in Table 1?”

We kindly refer the reviewer to point 1 in our overall response (“Importance of function-space priors and fundamental differences to TFSVI and FVI”). In brief, our method allows specifying prior beliefs in function space whereas TFSVI specifies the prior in weight space, but pushes it forward to function space to calculate the KL-divergence. However, such a push-forward only changes the representation of the prior and does not lift the limitations of expressing meaningful prior beliefs in weight-space since neural networks weights are not interpretable. These differences explain the improved performance of our method in Tables 1 and 3, and the qualitative improvements in Figure 2, and Figure 4 in the appendix. In these experiments, we chose the best weight-space prior that we could find for TFSVI by cross validation.

“jitter term has been commonly used in implementation to handle the numerical issue when training the model with the Gaussian KL divergence as KL objective”

Thank you for bringing this up. The different ways in which TFSVI, FVI (Sun et al., 2019), and our method introduce jitter is indeed somewhat subtle, but these differences are crucial. We discuss the differences in paragraph issue (iii) of Section 2.2, but we realize based on reviewer feedback that this issue requires further explanation. We kindly refer to points 2 and 3 in our overall response above, which we hope clarify these differences (we will add a similar clarification to the paper over the next few days). We provide a synopsis below.

TFSVI introduces jitter to avoid numerical issues. Since TFSVI defines both prior and variational distributions in weight space, their respective push-forwards into function space have the same support, and the KL-divergence between them would be technically finite if it wasn’t for numerical errors. By contrast, the jitter in our method does not resolve numerical issues but instead is fundamental to obtaining a finite objective. Indeed, in our method the prior and variational posterior have different support since the prior is specified in function space while the variational posterior is still limited to the push-forward of a variational distribution in weight space.

Compared to our method, FVI does not linearize the BNN and only implicitly defines the variational measure in function-space which makes the method substantially more complicated to use in practice (Ma et Hernández-Lobato 2021). Thus, “adding jitter” in FVI means drawing actual realizations of Gaussian noise. By contrast, we use the linearized BNN which yields a Gaussian variational posterior for which we can marginalize over the jitter analytically. It turns out that applying the estimator of FVI to our situation is equivalent to calculating the estimator of the regularized KL-divergence, which has an explicit form that we can directly optimize.

“Compared to the tractable function space VI of [1] (TFSVI), it seems that the objective in Eq. (9) exhibits only minor differences”

We agree that the implementation of our method is similar to TFSVI. But we see this as a strength rather than a shortcoming. In fact, we believe that our method has the potential of becoming a go-to solution for function-space BNNs precisely because it can be implemented easily and yet evidently offers improvements over TFSVI and FVI.

“While the tractable function space VI of [1] has been demonstrated on both classification and regression tasks, this work has been demonstrated only for regression experiment setting.”

We thank the reviewer for raising this point that we did not clarify in the paper. As mentioned above, the main difference between our method and the one in [1] is the use of GP priors in function space, which allow us to incorporate knowledge about the properties of the functions generated by the BNN. Prior beliefs in the classification setting are specified in the space of logits before the output activation, which is less interpretable than for regression and puts into question our initial motivation for informative priors. For this reason, we prefer to concentrate on the regression setting. However, using automatic prior selection along with our method might provide benefits over standard BNNs in classification tasks and would be an interesting direction for future work.

We hope that our response clarifies how our contribution differs from prior work, and that we were able to convince you that our simple method could be helpful to members of the community.

Thank you for responding to my questions. Your responses have helped me better understand the distinction between the proposed method and TFSVI, as claimed by the authors. However, I am somewhat skeptical that using Gaussian processes to construct a function-space prior can result in an interpretable function-space prior for a given dataset. Based on my understanding, the inductive bias of the GP prior could be interpretable for certain types of covariance functions like the RBF kernel, Matern, and Periodic kernel. Choosing the proper kernel function for GP seems non-trivial and depends on the dataset. Thus, I believe that relying on a GP prior to build an interpretable function-space prior might still have some limitations.

Personally, I think that providing a specific method for forming the interpretable function-space prior would further enhance the novelty of this work, especially considering the distinction between the proposed method and TFSVI.

With these reasons, I will maintain my score for the current content of the manuscript.