Improved Function Space Variational Inference with Informative Priors

Q1. Why was the last layer set as BNN layer? Why not set the whole network as BNN?

We request you to confirm our response regarding the Reasons for taking the Last-layer BNNs in the general review.

Q2. The experimental results are only marginal. Is that because the only last layer is BNN? Why not try more BNN layers that can demonstrate more difference between new prior with previous informative prior?

We request you to confirm our response regarding the Marginal Performance improvement, possibly due to the Last-layer BNN in the general review.

Q3. Since the prior is changing during the training, how to ensure the convergence of the procedure?

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.

Q4. Since the method is specially designed for the classification task, I suggest the author revise the title and introduction accordingly to highlight the classification task.

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.

Q5. What is the role of Eq. (12)?

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)$.

Q6. Some symbols are not defined, like N^q

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.

Q1. 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.

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

Q2. “Denote auxiliary input, 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 features …” What is an empirical parameter?

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.

Q3. When applying the last-layer, the corresponding model would be Bayesian linear regression with interpretable parameters. 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?

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.

Q4. The results reported in Figure 2 are mixed. Figure 2.(a) and (C) are favorable for the central message of the paper while Figure 2.(b) and (d) are just the opposite.

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.

Q5. 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 even needs an experiment for that.

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).

Q6. Only T-FVI or all the three baselines? If first, why not others?

Since we aim at tackling the issue of input-independent function-space prior for T-FVI in this work, we set the T-FVI as our main baseline.

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.

Q2. Our method is somehow ad-hoc without well-elaboration for Eq (9) and Eq (13).

We request you to confirm our response regarding the Elaboration of the proposed algorithm in the general review.

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

Q4. Regarding the marginal performance gain in Section 5.2.

We request you to confirm our response in the Marginal Performance improvement, possibly due to the Last-layer BNN

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.

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.

Q1. Authors do not make the efforts necessary to justify this last-layer approximation.

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.

Q2. Does the last-layer approximation play a role in the final performance metrics? What results are achieved if this restriction is not applied and instead a full Bayesian NN is used ?

We request you to confirm our response in the Marginal Performance improvement, possibly due to the Last-layer BNN in the responses to all reviewers above.

Q3. Comment about insufficient extensive evaluation results (evaluation on regression task) in Weakness

We request you to confirm our response in the Response to all reviewers for ""Extension for regression tasks"" above.

Q4. Why would you argue that the main goal of function-space BNNs is "directly assigning prior distributions to the outputs of neural networks?

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.

Q5. Connection to generalistic approach; Given the BMA approach and the nature of the contribution in Rodríguez-Santana et al. (2022), how do these relate to each other?

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

Reasons for taking the Last-layer BNN.

Memory Issue:

For widely-used deep neural networks (DNNs) like ResNet, computing the Jacobian matrix for the entire weight parameters demands a substantial amount of GPU memory due to the large size of the Jacobian matrix, which is #batch size x # weight parameters x #class; the size of Jacobian matrix for the ResNet 18 on CIFAR 10 using 128 batch is 128 x 11.8M x 10. This can potentially lead to out-of-memory issues, making it impractical to consider linearization for all weight parameters.

Indeed, due to this potential issue, the function-space prior in [4] employs the Jacobian matrix of the last layer for the covariance of the function-space prior. This approach motivates us; if the last layer is modeled as the Bayesian layer allowing the closed form of the function-space distribution, the function-space prior can be constructed even with limited GPU resources. The sample function of the variational distribution can also be efficiently constructed, as described in A.2 in the Appendix (page 12).

Effectiveness of the Last-layer BNN:

In large-scale BNNs, considering partially stochastic layers for BNNs can be more effective than using fully stochastic layers [5]. Based on this observation, we designate the last layer as the Bayesian layer. To address the necessary uncertainty modeled by the $L-1$ stochastic layers $\\{f^{l} \\}_{l=1}^{L-1}$, we directly model this uncertainty through the latent variable $z(x)$, by leveraging the $Q$-dimensional mahalanobis distance $|| \Delta \ h(x) ||$ for the penultimate feature $h(x)$ from using the averaged penultimate features in Eq. (9).

Marginal Performance improvement, possibly due to the Last-layer BNN.

To investigate the effectiveness of the last-layer approximation, we conducted additional experiments on CIFAR-10 using ResNet-18 as the baseline model. Due to the memory issues explained above, we extended the Bayesian last layer for the baseline (T-FVI) by replacing the last linear layer ([512, 10] with 512 layer features and 10 classes) with a Bayesian 2-hidden MLP layer ([512, 128] -> ReLU -> [128, 10]). Note that this Bayesian 2-hidden MLP layer uses an increasing number of weight parameters for the mean and variance parameters of the variational and prior distributions.

Additionally, we considered varying numbers of training sets, specifically $\\{10,000, 20,000, 30,000, 45,000\\}$, to demonstrate the effectiveness of the proposed inference. For the hyperparameters, we tuned the regularization hyperparameter $\lambda$ in $\\{10^{-3}, 10^{-4}, 10^{-5}\\}$ for KL divergence and $r$ in $\\{0.05, 0.10\\}$ for the adversarial hidden feature with the temperature parameter $\tau=1/0.9$ for the variational distribution.

We report ACC, NLL, AUROC, and ECE, respectively, following the protocol in the manuscript. We observe that R-FVI consistently outperforms the baselines in terms of NLL and AUROC, except when using a $10,000$ training set. We would like to emphasize that the improvement of R-FVI is not marginal when comparing its performance with that of T-FVI (2 MLP Layers).

Modification for Regression

For regression task, the proposed method can be also used with a minor modification as the following:

For a continuous-valued output $Y=\\{y_{n}\\}_{n=1}^{N}$ in the training set, we consider the range of $[-\infty, \min{Y}] \cup [\min{Y}, \max{Y}] \cup [\max{Y}, \infty]$, partitioning this range into $Q$ intervals, including $(-\infty, \min{Y})$ and $(\max{Y}, \infty)$. We assign the order of each interval as the discrete label. For instance, if $y_n$ belongs to the $q$-th interval out of $Q$ intervals ($1 \leq q \leq Q$), we define the discrete label of $y_n$ as $q$. Then, we compute Eq. (8) using the assigned discrete label for each $y_n$, building the function-space prior in Eq. (10) as done in the classification task.

Regarding the variational distribution, we consider the latent variable $z(x) \sim Cat( \hat{p}(x) )$ using mahalanobis distance, i.e., $\hat{p}(x) = || \Delta h(x) ||$ for an input $x$. Then, we apply the parameter $\tau$ to the only covariance parameter $\Sigma_{f|z}$ in Eq. (15), enabling the sample function $f(x)$ to exhibit varying levels of predictive uncertainty depending on the interval corresponding $x$. However, we believe that covering this approach in this work is too extensive. Therefore, for the regression, we recommend using the proposed function-space prior without using the latent variable $z$. Indeed, we confirm that using the regularization by the proposed function-space prior yields good performance, as below:

Experiment Result

We validate the proposed method (R-FVI) on the UCI regression task and compare R-FVI with FVI [6] and T-FVI [4] using the GP prior with the RBF kernel. We follow the same setting as the UCI regression experiment in the appendix of [4]. For the proposed method, we set $Q=20$ intervals and find the best regularization hyperparameter $\lambda \in \\{0.001, 0.01, 0.1\\}$ for KL divergence and the radius $r \in \\{0.01, 0.1 \\}$ in Eq. (12). We apply map inference for the first 50 percent of the total training iterations and then apply function-space variational inference for the remaining iterations.

The first table reports RMSE ($\downarrow$: better), and the second table reports log likelihood, ($\uparrow$: better). We observe that the proposed method tends to outperform the baselines FVI and T-FVI.

Sensitivity analysis for $Q$ intervals

Furthermore, we conducted additional experiments to assess the robustness of the proposed method across different $Q$ interval on the ''Concrete,' 'Energy,' and 'Protein' datasets. The results, presented below, illustrate the consistent performance of the method under varying number of interval $Q \in \\{5, 10, 15, 20 \\}$.

[6] Functional Variational Bayesian Neural Networks - ICLR 19

[4] Tractable Function-Space Variational Inference in Bayesian Neural Networks - NeurIPS 22

[5] Do Bayesian Neural Networks Need To Be Fully Stochastic? - AISTATS 23

How the proposed methods achieve the BMA property using the inductive bias.

Role of the function-space prior:

The proposed function-space prior is designed to induce a more peaked probability distribution (low-entropy) when the auxiliary input is more similar to the in-distribution dataset. Specifically, the proposed function-space prior $f_{BMA}(\cdot) = \Theta^{L} \ proj(h(\cdot))$ in Eq. (10) is constructed using the averaged weight parameters and penultimate hidden features in Eq. (9), which are obtained from pre-defined iterations ($\mathcal{T}=\\{145, 150, 155, 160\\}$ in the above example). We would like to emphasize that these averaged weight parameters and hidden features can form the approximate posterior based on the fact that the averaged weight parameters using the finite weight samples $\\{w(t)\\}_{t=1}^{T}$, ($T=200$) during training iterations can serve as an approximate posterior of the BNN in SWAG [2].

In this context, when the auxiliary input $x$ is more similar to the in-distribution dataset, the corresponding $proj(h(x))$ is likely to be one of the penultimate hidden features in Eq. (9), such as ${\hat{m}^{q}_{h}}$ representing the $q$-the label. Then, the sample function of the prior distribution would have a peaked predictive probability for the $q$-th label using the random weight $\Theta^{L}$ modeled by the averaged weight parameters in Eq. (9) .

Role of the function-space prior:

The proposed function-space variational distribution $q(f)$ is designed to allow the sample functions $\\{ f_{j} (x) \\}$ of the variational distribution to exhibit varying levels of disagreement between predicted classes on $x$. Specifically, for an input $x$, we model the level of disagreement by introducing the categorical latent variable $z(x) \sim Cat(\hat{p}(x))$ parameterized by $\hat{p}(x) = \frac{ \mu^{L}_{w_q} \ h(x) }{ \sqrt{ 1 + \pi/8 || \Delta \ h(x) ||} } \in \Delta^{Q-1}$ in Eq. (13).

We would like to emphasize that $|| \Delta \ h(x) || \in R^{Q}$ in the denominator denotes the $Q$-dimensional mahalanobis distance of $h(x)$ from the averaged penultimate hidden features $\hat{m}^{q}_{h}$ for $q=1,..,Q$ in Eq. (9). The $|| \Delta \ h(x) ||$ has known to be effective for OOD detection in MAHALANOBIS [3] because its has the small distance for the specific $q$-th dimension or has the large distance for all $q=1,..,Q$, depending on how close a input $x$ is to the in-distribution.

In this context, when the $x$ seems to be a out-of-distribution dataset, the $\hat{p}(x)$ would become a flat probability due to $|| \Delta \ h(x) ||$ . This induces the latent variable $z(x)$ to have a different level of disagreement depending on $x$, and thus affects the sample functions $f (x)$ to exhibit varying levels of disagreement by applying the temperature parameter $\tau$ on the specific dimension modeled by $z(x)$.

The distinct property of the function-Space prior as input-dependent prior

A BNN can be understood as consisting of two stochastic Neural Networks (NNs): one utilizes the prior distribution of weight parameters, and the other employs the variational distribution of weight parameters. In this context, the function-space prior represents the distribution of the (linearized) stochastic NN's output using the prior distribution of weight parameters. The function-space prior can be used to regularize the stochastic NN (variational distribution) by minimizing the KL divergence between these stochastic NN output distributions. To this end, the function-space prior should specify the auxiliary inputs to be used for constructing the function-space prior (input specification).

This is a distinctive property of the function-space prior compared to the weight-space prior because the function-space prior can regularize the stochastic NN (variational distribution) by employing "input-dependent” prior distributions on the function-space. In contrast, the weight-space prior for the conventional BNN regularizes the stochastic NN (variational distribution) by using “input-independent” prior distributions on weight-space.

Issue of the existing function-Space prior

The function-space BNN constructs a function-space prior by assigning a zero-mean Gaussian distribution as the weight prior distribution because it was not well understood how the choice of auxiliary input and weight prior distribution affects the corresponding function-space prior distribution. However, this approach results in the input-independent zero-mean function-space prior. This can introduce unnecessary uncertainty if the auxiliary input is similar to the in-distribution dataset, as illustrated in Figure 1 (a). To understand this, let us consider the training set as CIFAR 10 and the auxiliary inputs as CIFAR 100 where some inputs in CIFAR 100 share the same labels with CIFAR 10. Regularizing the zero-mean function-space priors of these auxiliary inputs (specific subset of CIFAR 100) into the BNN would introduce unnecessary uncertainty.

Direction of the proposed approach

To establish a reasonable function-space prior, we are exploring how the predictive probability, sampled from the posterior distribution, can exhibit appropriate uncertainty based on the similarity of auxiliary inputs to the training set (for example, auxiliary inputs from CIFAR 100 sharing labels with CIFAR 10). Consequently, we have observed that even if each predictive sample function has a peaked probability, the averaged prediction can be uncertain if the predicted classes exhibit varying levels of disagreement. This principle can elucidate why Deep ensemble [1] achieves good performance. Our goal is to integrate this insight into both the function-space prior and the corresponding function-space posterior distribution.

Elaboration of the proposed algorithm

To achieve this, we propose a learning algorithm based on two inductive biases commonly used in deep learning methods:

SWAG [2]: The empirical mean and covariance of weight parameters has been known to play a similar role with the approximate posterior distribution of BNNs weight parameters. We utilize the empirical mean and covariance, obtained during pre-defined iterations, to construct the function-space prior.

MAHALANOBIS [3]: The Mahalanobis distance of the penultimate hidden feature $h(x)$ has known to be effective in discriminating whether a input $x$ belongs to the in-distribution dataset. We use this inductive bias to enable the predictive sample functions of $x$ to have varying degree of disagreement in their predicted class.

Procedure of the proposed algorithm

The proposed learning algorithm consists of two phases. Let $T=200$ be a total of training iterations, and $\mathcal{T}=\\{145, 150, 155, 160\\}$ be the pre-defined iterations as following the notation in our manuscript.

Initial Phase (0-160): Following the methodology similar to [1], we train the mean parameters of the variational distribution of the BNN weight parameters using MAP inference. During this phase, we compute the empirical parameters in Eq. (9) at the pre-defined $\mathcal{T}=\\{145, 150, 155, 160\\}$ to construct the function-space prior of Eq. (11).

Remaining Phase (161-200): We train the mean and variance parameters of the BNNs' variational distribution using function-space variational inference with the proposed function-space prior of Eq. (11) and variational distribution of Eq. (15).

[1] Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles - NeurIPS 17

[2] A simple baseline for bayesian uncertainty in deep learning - NeurIPS 19

[3] A simple unified framework for detecting out-of-distribution samples and adversarial attacks - NeurIPS 18