Paper: Regularized KL-Divergence for well-defined function space variational inference in BNNs (ICLR-2024)

Reviewer comment (Novelty):
No significant contribution. As far as I could understand, the only contribution of this paper is to bring the regularized KL divergence which was well developed by (Quang 2019), and use it as a substitute for the ordinary KL divergence in the framework of tractable function space VI (TFVI, Rudner et al., 2022). Other than this, I fail to see any contribution, and even the combination of those two methods is implemented quite straightforwardly, without any issue to consider during that process.

Author response:
> 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.

Reviewer comment (Experiments):
Limited experiments. While I appreciate the experiments on the regression tasks, they are relatively small-scale tasks, and only the small BNNs (MLPs mostly) are tested. It is hard to judge the effectiveness of the proposed method without scaling, for instance, the image classification task (CIFAR-10 or CIFAR-100, at least) solved with ResNet, as typically done in the literature.

Author response:
> 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.

Reviewer comment (Theory):
Once we replace the KL divergence with regularized KL divergence, the resulting objective becomes something that is different from ELBO, so we end up optimizing an objective that is not particularly a lower bound on the marginal likelihood. There is a class of inference algorithms (generalized Bayesian inference) generalizing the standard Bayesian inference procedure and thus variational inference by extending the likelihood or KL regularization term with more flexible functions, and for such cases the utility of the extended objective can be described in the perspective of generalizaton error. However, in the current form, the alternative objective with the regularized KL does not explain anything about its optimum. I think there should be some intuitive justification for this.

Author response:
> 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.

Reviewer comment (Novelty):
Compared to the tractable function space VI of [1], it seems that the objective in Eq. (9) exhibits only minor differences, i.e., jitter term ($\gamma M I_M$) in each covariance term in Eq. (11). Based on my understanding, adding the 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. Therefore, the proposed objective itself does not seem novel in sense of training objective for VI.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Novelty):
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]?

Author response:
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.

Reviewer comment (Evaluation):
In comparison to [1], what specific difference in the proposed method leads to the improved performance in Table 1?

Author response:
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.