Paper: Improved Function Space Variational Inference with Informative Priors (ICLR-2024)

Reviewer comment (Theory):
Since the prior is changing during the training, how to ensure the convergence of the procedure?

Author response:
Our function-space prior does not vary during training. Please refer to our response in the **Procedure of the proposed algorithm** above and Algorithm 1 on page 12 for details.

Reviewer comment (Writing):
Since the method is specially designed for the classification task, I suggest the author to revise the title and introduction accordingly to highlight the classification task.

Author response:
We have extended our method for the regression task and included regression results. Please see our response in the **Extension for regression tasks**. We will reflect this in the final manuscript.

Reviewer comment (Writing):
What is the role of (12)?

Author response:
The function-space prior requires auxiliary inputs, as discussed in our response regarding the **distinct property of the function-space prior as an input-dependent prior**. However, selecting additional datasets for the auxiliary inputs could be cumbersome if those datasets are not readily available. Therefore, for each training input $x_i$, we consider the $h(\cdot)$ from Eq. (12) to be used for constructing the function-space prior. It's important to note that the $h(\cdot)$ from Eq. (12) using $x_i$ is different from $h(x_i)$.  Therefor, this results in a function-space prior derived from $h(\cdot)$ of Eq. (12), which is distinct from the one derived from $h(x_i)$.

Reviewer comment (Writing):
Some symbols are not defined, like N^q

Author response:
The $N^q$ represents the number of the labels corresponding to $q$ for the training set, i.e., $N^q= \vert \ \\{y_n \ | \, y_n=q, y_n \in \mathcal{D}\\} \ \vert$, for the training set $\mathcal{D}=\\{(x_n, y_n)\\}_{n=1}^{N}$ and $y_n \in \\{1,..,Q\\}$. For other errors, we will correct them in the final manuscript.

Reviewer comment (Theory):
The problem setup does not make much sense to me. The introduction says: **“We build an informative function space prior by using the empirical Bayes approach along with the parameters of hidden values and the last layer weight parameters which are obtained iterations during early stages of training”**. I wonder how the approach is then different from having an uninformative prior after all. If the information comes from the data, it is technically not a prior. It appears that the paper makes its main point by differentiating from cases where the priors are just so strong that they unnecessarily restrict the model capacity.

Author response:
> The proposed function-space prior is adaptively constructed based on auxiliary inputs, as described in **The Distinct Property of the Function-Space Prior as an Input-Dependent Prior**. We devise the function-space prior in such a way that if the auxiliary inputs appear more similar to the in-distribution dataset, the proposed function-space prior would exhibit higher a peaked probability due to its construction, as outlined in the initial phase of the **Procedure of the proposed algorithm**. The regularization of this prior would adaptively influence the model parameters depending on the auxiliary input (for example, $A$: the subset of CIFAR 100 sharing the label of CIFAR 10, and $B$: the remaining datasets of CIFAR 100). This is definitely distinct to the uninformative prior applying the same regularization to both $A$ and $B$.

> Regarding the comment about the prior using the dataset, we partially agree with it because the empirical Bayesian method [1] sets the parameters of the prior distribution by training on the datasets. We adopt a similar approach to establish the function-space prior.

> [1] https://en.wikipedia.org/wiki/Empirical_Bayes_method

Reviewer comment (Writing):
The paper significantly lacks clarity. Apart from having extremely many typographical errors, it has statements without sharp enough meanings. Among the many, one example is: **“Denote auxiliary inputs, which are far from training points and are placed closely with the training sets, respectively”**. I have no idea what it means for a training point to be close to a training set. Likewise: **“h(.) from the q-th component empirical parameters of hidden feature …”** What is an empirical parameter? I also have no clue about what is going on in pages 5 and 6 after spending considerable time trying to read them. Even the purpose of all these complications such as introducing adversarial hidden features do not look to me justified.

Author response:
> In light of the feedback from other reviewers, we acknowledge that our manuscript seems difficult to understand. We kindly request you to confirm the responses from all reviewers above. Then, we recommend revisiting sections 2, 3, and 4 for a clearer understanding. Also, we will enhance the clarity of our manuscript for easier understanding.

Reviewer comment (Theory):
It is a clear weakness that after motivating Bayesian inference with lots of effort, the paper ends up using it only on the penultimate layer. Those layers are typically linear, where even closed-form Bayesian linear regression would work and the learned model weights would give a degree of interpretability. Why should one use function-space Bayesian inference if the prior will come from the weights learned in another data set and only the penultimate layer will be Bayesianized?

Author response:
> Although our proposed approach models only the last layer to be the Bayesian layer, the regularization via the function-space KL divergence of Eq. (6) has an impact on the overall parameter learning, including the weight parameters of the $L-1$ deterministic layers $\\{f^{l} \\}_{l=1}^{L-1}$. 

> Note that from a Bayesian perspective, the weight parameters of  L-1 deterministic layers $\\{f^{l}\\}_{l=1}^{L-1}$ can be also understood as the mean parameters of the variational distribution for the BNN having a fixed variance (not trainable) with a very small value.

Reviewer comment (Evaluation):
I do not think the reported results demonstrate the benefit of the proposed approach clearly enough. All models in all experiments perform very closely to each other. The results reported in Figure 2 are mixed: (a) and © are favorable for the central message of the paper while (b) and (d) are just the opposite.

Author response:
> We would like to emphasize that the primary message from Figure 2 is that the proposed method results in more reliable uncertainty estimation for both the in-training set (lower NLL) and out-of-training test set (higher AUROC) across various regularization parameters λ. Therefore, it demonstrates that the proposed method effectively mitigates the trade-off issue associated with the existing input-independent function-space prior.

Reviewer comment (Experiments):
The take-home message given in the last sentence of Section 5.1 is obvious and comes from the nature of using an arbitrary regularizer. I wonder why one needs even an experiment for that.

Author response:
> We agree that controlling the regularization hyperparameter $\lambda$ can alleviate the trade-off issue for the input-independent prior. However, we propose a more effective solution by incorporating the input-dependent function-space prior. Note that the NLL and AUROC of the proposed method (R-FVI) consistently outperform the baseline (T-FVI) for all varying regularization hyperparameters in Figure 2 (a) and (b).

Reviewer comment (Theory):
In Section 4, although the authors motivated the paper from the view of Bayesian Model Averaging, and claimed that *“we may design a function space prior that does not explicitly encourage generating high-entropy predictions for each predictive probability but the average prediction (via BMA) would still have high-entropy when encountered with an OOD input”*, it is not clear how the proposed prior can achieve this.

Author response:
### Q1. Regarding our claim that the proposed function-space prior allows the average prediction (via BMA) to have high-entropy for an OOD input, it is not clear  how the proposed function-space prior achieves this.

> We request you to confirm our response regarding the **How the proposed methods achieve the BMA property using the inductive bias** in the general review.

Reviewer comment (Novelty):
The authors ignored a very related work from Izmailov et al. (2021). The narrative of this work also relies on Bayesian model averaging. This work also considers designing a novel prior that is robust to out-of-distribution data by using the empirical Bayes approach. The authors should cite, discuss, and compare experimentally with this work.

Author response:
### Q3. The authors ignored a very related work from Izmailov et al. (2021) relying on Bayesian model averaging and a novel prior. 

> Based on our understanding, in the referenced study [1], the authors addressed the issue of poor uncertainty estimation in BNNs on the covariate shift dataset when the BNNs, using a zero-mean Gaussian prior distribution, are trained using SGHMC (oracle Bayesian inference yielding an almost true posterior). To address this limitation, they identified shortcomings in the posterior distribution arising from the zero-mean Gaussian prior and proposed a novel empirical covariance for the prior distribution of weight parameters.

>  We acknowledge some similarities with [1], where both approaches use training datasets to construct the prior distribution for BNNs with the aim of improving uncertainty estimation. However, our work specifically focuses on resolving the trade-off issue caused by the input-independent function-space prior. Therefore, we refine the function-space prior and variational distribution to achieve reliable uncertainty estimation for both the in-training distribution and out-of-distribution dataset. This objective distinguishes our work from [1], which primarily concentrated on addressing the covariate shift dataset. We will elaborate on this discussion in the main manuscript.

>  [1] Dangers of Bayesian Model Averaging under Covariate Shift - NeurIPS 21

Reviewer comment (Reproducibility):
In Equation (8), what is $N^q$ how to define it? Do we have to compute the means and shared covariance over the dataset?

Author response:
R5 and R6: ### Q5. What is $N^q$, how to define it?, 

> The $N^q$ represents the number of the labels corresponding to $q$ for the training set, i.e., $N^q= \vert \ \\{y_n \ | \ y_n=q, y_n \in \mathcal{D} \\} \ \vert$, for the training set $\mathcal{D} =\\{(x_n, y_n)\\}_{n=1}^{N}$ and $y_n \in \\{1,..,Q\\}$.

### Q6. Do we have to compute the means and shared covariance over the dataset for Eq. (8)?

> We can utilize a subset of the full training set for computing the mean and shared covariance. However, this introduces another consideration—specifically, how to choose the subset of the training set. If the subset is selected in a way that introduces bias, it can give rise to additional issues for the function-space prior and the latent variable $z$ in the variational distribution. Consequently, we opt for using the entire training set. Since the means and shared covariances are updated in a batch manner, there is no need for additional iterations to compute the mean and shared covariance.

Reviewer comment (Writing):
In Equation (9): what is $T$? It should be consistent with the sentence before the Equation (8).

Author response:
### Q7. what is $T$ ?,  It should be consistent with the sentence before the Equation (8).

> The $T$ represents the number of pre-determined iterations. For example, if $\mathcal{T}=\\{145, 150, 155, 160 \\}$, then $T=4$. We apologize for the error in this notation. We will clarify it in the final manuscript.

Reviewer comment (Evaluation):
From Figures 3c and 3f, it seems that the proposed method always induces a much higher disagreement ratio on both in- and out-of-distribution data compared to the uniform prior. Is this good? On in-distribution data, the entropy should be low.

Author response:
### Q8. From Figures 3c and 3f, it seems that the proposed method always induces a much higher disagreement ratio on both in- and out-of-distribution data compared to the uniform prior. Is this good? On in-distribution data, the entropy should be low.

> We acknowledge that in Figure 3c and 3f, our proposed method (R-FVI) exhibits a higher disagreement ratio on both in-distribution and out-of-distribution data compared to the uniform function-space prior (T-FVI).  

> However, the disagreement ratio of R-FVI on in-distribution data is significantly smaller than that on out-of-distribution data  in Figure 3c. In addition, in Figure 3.b (in-distribution data), the entropy of R-FVI becomes smaller compared to T-FVI as the number ($J$) of sample functions increases. Conversely, in Figure 3.e (out-of-distribution data), the entropy of R-FVI is larger than that of T-FVI. This suggests that the sample functions having a higher disagreement ratio can induce the smaller predictive entropy on in-distribution data.

Reviewer comment (Writing):
In the paragraph "Context inputs from adversarial hidden features". How do we define $r$?

Author response:
### Q9. What is $\gamma$ ?, in the paragraph "Context inputs from adversarial hidden features". How do we define it ?

> The $\gamma$ denotes the radius for the neighborhood set of $h(x_i)$, i.e.,
$\\{h \ | \   ||h - h(x_i)||_2 \leq r \\}$ for a training input $x_i$ .

> We determine the value of $\gamma$ through validation. Specifically, we consider $\gamma$ within the set $\\{0.01, 0.05, 0.1, 0.2\\}$ and select the value that yields the best negative log-likelihood (NLL) from the validation set. Additional details can be found in Table 7 on Page 16.

Reviewer comment (Theory):
The proposed method relies on a last-layer approximation, which is not thoroughly discussed enough. While this approach can be employed with the right arguments, the authors do not make the efforts necessary to justify this choice or the consequences it may entail in the proposed technique.

Author response:
We explain why the last-layer approximation was considered in **Reasons for taking the Last-layer BNN** in the responses to all reviewers. We will also describe this in our final manuscript as well.

Reviewer comment (Theory):
The statement that the main goal of function-space BNNs is "directly assigning prior distributions to the outputs of neural networks" is questionable. This can be done without function-space formulation, and it is arguable that the main interest lies directly on the properties of the function space itself. Furthermore, the approach does not necessarily become more "user-friendly" due to the difficulties intrinsic to function space. More insight on these points would be appreciated.

Author response:
When considering one of the Bayesian perspectives, which involves incorporating prior knowledge into the model, we believe that providing a more directly interpretable function-space prior for BNNs is a desirable direction for practitioners. We acknowledge that this perspective can be subjective and thus would like to moderate our claim regarding the main goal. Additionally, we agree with your assertion that addressing the characteristics of the function space and overcoming the challenges of function-space inference are indeed meaningful directions for improving function-space inference.

However, if we focus solely on the characteristics of the function space and technical research, we might overlook the advantages of the Bayesian approach—especially its ability to enhance performance with limited data by leveraging prior knowledge. Therefore, we believe that research on BNNs that can incorporate intended prior knowledge is crucial. Exploring ways to make the function-space prior more user-friendly is the key direction in this context.

Reviewer comment (Novelty):
Given the BMA approach and the nature of the contribution in Rodríguez-Santana et al. (2022), further discussion is required to clarify how these relate to each other and to provide a more comprehensive overall picture.

Author response:
Based on our understanding, Sparse Implicit Process (SIP) [1] primarily contributes to improving the scalability and flexibility of the Implicit Process (IP), known for constructing an implicit function-space distribution to model data. SIP achieves this improvement by incorporating the inducing input idea and adversarial Variational Inference. Our research, on the other hand, focuses on BNNs, considered as an example of IP. We aim to enhance function-space Variational Inference for BNNs, inspired by the ideal role of function-space predictive sample functions.

Regarding the connection to these works, we believe that these studies offer interesting insights into the construction of a function-space prior. SIP forwards Gaussian noise and data through a neural network to obtain random functions. It then uses the empirical distribution (averaged mean and covariance of the random functions) to form and learn the implicit prior distribution on function space. In contrast, our study utilizes the paths of weight parameters during training to construct an input-dependent function-space prior with different predictive entropy for in-distribution and out-of-distribution datasets. This difference seems to originate from the fact that, while SIP focuses more on obtaining the implicit function-space prior for flexible modeling, our work concentrates more on the adaptive function-space prior for Bayesian Model Averaging (BMA) prediction. This represents our perspective on the connection between our work and SIP. If there are different opinions, feel free to share them. It could contribute to a more comprehensive discussion of the papers.

[1] Function-space Inference with Sparse Implicit Processes - ICML 22