Paper: Functional Wasserstein Bridge Inference for Bayesian Deep Learning (ICLR-2024)

Reviewer comment (Experiments):
While the experiment setup is good (it encompasses different applications such as regression, bandits, and classification), it ignores non-variational baselines, such as the deep kernel learning (DKL, Wilson et al., 2016) and the Laplace approximation (LA, Daxberger et al., 2021). They are both function-space BNNs—for the latter, it is given by the linearized (Immer et al., 2021) and the last-layer Laplace (Riquelme et al., 2018) formulations. They have been shown to be very scalable (for the last-layer version, even applicable to transformers) and yield good uncertainty quantification performance. Moreover, they have also been applied in similar experimental setups as in the present paper (e.g., bandits (see Sec. 3.2 here: https://arxiv.org/pdf/2310.00137.pdf) or Bayesian optimization (https://arxiv.org/abs/2304.08309)). Comparing the FWBI with a Laplace baseline is seen as a must, both in terms of performance and costs, the latter being very important since it dictates whether a method is practical or not.

Author response:
$\textbf{Reply:}$ For deep kernel learning, unlike our FWBI, we think its main application is not to BNNs, but to some tasks about GPs, so we did not compare FWBI to it for the time being. For Laplace approximation (LA, Daxberger et al., 2021), we have added this baseline results on the toy example in Figure 6 in Appendix E.2.  Compared to it, our model shows stronger fitting ability and competitive uncertainty estimation.

Reviewer comment (Experiments):
Regarding Table 2, the classification performance presented is very bad. The reviewer suggests the authors use more standard networks for those problems, e.g., LeNet-5 for MNIST and FMNIST, and ResNet for CIFAR-10. Otherwise, Table 2 will only cast FWBI in a bad light.

Author response:
$\textbf{Reply:}$ Yes, we may obtain better performance using the other more sophisticated network architectures, and so will other methods. Other possible ways to further improve the performance may include batch normalization, data augmentation, adaptive learning rate, hyperparameter tuning using black-box optimizer, etc. However, the classification experiment here is used to demonstrate the effectiveness of the proposed functional inference compared with other methods. We think a fair setting for all comparative methods is enough.

Reviewer comment (Novelty):
The authors mention that "prior distillation"—the first step in FWBI—and "Wasserstein variational inference"—the second step of FWBI—have been done separately before. This makes FWBI seem to simply combine two prior things together, raising uncertainty over its novelty. The reviewer asks the authors to comment on this.

Author response:
$\textbf{Reply:}$ We revised the ‘Related Work’ section to emphasize our innovations and contributions compared to prior works. We did use the similar idea with “prior distillation”, but we want to highlight that “prior distillation” is not for functional variational inference. The task of “prior distillation” is the prior matching. The other difference is that the 1-Wasserstein distance to measure the distance between two stochastic processes is only one terms of our algorithm, and we have another two terms: one is log-likelihood and the other is the distance between variational posterior and bridging distribution. The three terms are jointly optimized in a functional variational inference framework which is our main contribution in this study. “Wasserstein variational inference” is firstly a parameter-space method but ours is functional space method. Secondly, “Wasserstein variational inference” defines the loss function using the Wasserstein distance between variational distribution and prior, but ours is a log-likelihood term plus the Wasserstein distance between variational distribution with bridging distribution. The two forms are not equal with each other because the distance is Wasserstein rather than KL divergence.

Reviewer comment (Novelty):
The "Related Work" section is supposed to address the confusion about novelty. However, the reviewer finds that the section merely describes prior works without discussing the differences between them and the authors’ present work. The reviewer suggests reworking that section for better clarity.

Author response:
$\textbf{Reply:}$ We revised the ‘Related Work’ section to emphasize our innovations and contributions compared to prior works. We did use the similar idea with “prior distillation”, but we want to highlight that “prior distillation” is not for functional variational inference. The task of “prior distillation” is the prior matching. The other difference is that the 1-Wasserstein distance to measure the distance between two stochastic processes is only one terms of our algorithm, and we have another two terms: one is log-likelihood and the other is the distance between variational posterior and bridging distribution. The three terms are jointly optimized in a functional variational inference framework which is our main contribution in this study. “Wasserstein variational inference” is firstly a parameter-space method but ours is functional space method. Secondly, “Wasserstein variational inference” defines the loss function using the Wasserstein distance between variational distribution and prior, but ours is a log-likelihood term plus the Wasserstein distance between variational distribution with bridging distribution. The two forms are not equal with each other because the distance is Wasserstein rather than KL divergence.

Reviewer comment (Theory):
Proposition 1 tells us that the FWBI objective is a lower bound to the standard ELBO. While it is good that FWBI defines a lower bound to the marginal likelihood, the reviewer thinks the result is incomplete. For example, a sufficiently small constant function would also be a valid lower bound, but it would obviously be a bad one. The reviewer asks the authors to comment on how good the FWBI objective actually is.

Author response:
$\textbf{Reply:}$ We want to firstly explain the role of ELBO here. The reason why we want to have a valid ELBO of the data likelihood is because we can achieve the marginal likelihood maximization via maximizing the valid ELBO. Yes, a sufficiently small constant function could also be a lower bound for marginal likelihood, but we cannot achieve the marginal likelihood maximization via maximizing this “sufficiently small constant function”. Maximizing a good lower bound could improve the data marginal likelihood and then improve the predictive performance, so we can use our experimental results to evaluate our FWBI objective.

Reviewer comment (Evaluation):
Still in Proposition 1, and ignoring the obvious flaws presented in Section 2, the reviewer asks the authors to comment on why FWBI yields better performance than other (KL-based) FVBNNs even though the KL-ELBO is a tighter lower bound.

Author response:
$\textbf{Reply:}$ Firstly, the tighter bower is preferred due to the larger possibility of obtaining maximized data likelihood, but it does not guarantee to obtain better performance in practice. One example is Trust Region Policy Optimization (TRPO) [1] which uses KL divergence to replace the Total Variation divergence in practice but, in fact, Total Variation divergence has a tighter bound than KL.  

Secondly, apart from the obvious flaws presented in Section 2, one possible reason is that KL divergence is easily caught in a local minimum, but Wasserstein distance has the better ability to jump out of local minimum. To show that, we add a Figure 17 in Appendix E.10, where we use KL and Wasserstein distances as the loss function to approximate a Gaussian mixture model (with three components) using a (single-mode) Gaussian distribution. We found that 1) KL divergence is sensitive to the initialization. For different initializations, there are two different results (Figures 17(c) and 17(f)) from KL divergence. In contrast, the results from Wasserstein distance (Figures 17(b) and 17(e)) are the same under different initializations; and 2) Wasserstein could jump out of the local optimum and move close to the global optimal mode (which is the up-left corner one with the darkest colour in Figure 17). 

[1], https://arxiv.org/pdf/1502.05477.pdf

Reviewer comment (Evaluation):
**Evaluation Index**: My expertise primarily lies in regression tasks, and I wonder whether the authors have considered employing alternative evaluation metrics such as calibration curves, as suggested in Reference [1]. The inclusion of such metrics could potentially provide a more comprehensive assessment of the model's performance.

Author response:
$\textbf{Reply:}$ As suggested, we have added the calibration curves of all parametric and functional methods for toy regression examples shown in Figure 8 in Appendix E.4, where our FWBI shows the most superior calibration results.

Reviewer comment (Theory):
**Definition of Convergence**: Upon reviewing Appendix D.3, I noted that the convergence line appears to be based on the learning objective. Could you please clarify your definition of convergence within this context? Furthermore, is it possible to theoretically prove the algorithm's convergence, given your chosen definition?

Author response:
$\textbf{Reply:}$ The convergence figures are the convergence processes of the 1-Wasserstein distance between the bridging distribution over functions and the functional prior (the top row), and the 2-Wasserstein distance between the variational posterior and the bridging distribution (the bottom row) during the training. Additionally, we added convergence plots for the overall variational objective function of FWBI in Figures 9 and 10 of Appendix for two toy examples respectively. The convergence here means the objective function to be converge to some stable values during the optimization. The comparisons of methods are valid only when they are converged well. That is the reason why we plot them here. Note that our algorithm is to design a loss function rather than optimize the loss function. Specially, we have designed a loss function in Eq. (11) for functional BNNs problem, and then we use an optimizer (e.g., SGD or Adam) to optimize this loss function. Given a loss function, the convergence is guaranteed by the optimizer rather than our algorithm.

Reviewer comment (Experiments):
While the paper already evaluates several priors (RBF, Matérn, linear kernels, and a BNN), it lacks some more interesting ones, e.g., how would a periodic kernel-based prior (as mentioned in the introduction) perform in Figure 2?

Author response:
$\textbf{Reply:}$ We have added posterior results for periodic-kernel based GP prior in Figure2, and FWBI still performs well with periodic kernel compared to other methods.

Reviewer comment (Evaluation):
The contribution section promises "reliable uncertainty estimation", while barely providing such an experiment in the form of OOD detection. A more detailed evaluation would involve reporting predictive log-likelihoods for all regressions and classification experiments (in addition to the current RMSE/Accuracy), as well as, e.g., expected classification errors for the classification setup.

Author response:
$\textbf{Reply:}$ We have added the NLL results for UCI regression in Table3 in Appendix E.7. For OOD detection in classification task, there are AUC results based on Entropy of the predictions shown in Table 2 in the main body, and Figure 16 in Appendix E.9 also shows the ROC curve. Our FWBI shows competitive uncertainty estimation performance compared to other methods. We also added the calibration curves of all parametric and functional methods for toy regression examples shown in Figure 8 in Appendix E.4, where our FWBI shows the most superior calibration results.

Reviewer comment (Evaluation):
For Figures 1/2 the mean fit looks a lot better than for the baselines, but still lacks a lot with respect to predictive uncertainty. Compare this to Figure 1 in Wild et al. (2022) who use the same function.

Author response:
$\textbf{Reply:}$ It is possible that the experimental results will not be the same under different settings. We have added the results of Wild et al. (2022) under our settings on the toy example shown in Figure 6 of Appendix E.2, where we can see that our FWBI still shows stronger predictive ability to recover the target function and competitive uncertainty estimation in the unseen regions.

Reviewer comment (Theory):
Sec 2, second paragraph claims "if prior and likelihood distributions are both assumed to be Gaussian, the posterior $p(\mathbf{w}|\mathcal{D})$ can be solved analytically as a Gaussian. As the posterior for neural-network-based mappings within the likelihood is highly multimodal, this is a false statement.

Author response:
$\textbf{Reply:}$ Sorry for the misunderstanding. In fact, our statement was intended for the general Bayesian inference not for BNNs, and we have removed this sentence in the revision to avoid the misunderstanding.

Reviewer comment (Theory):
Similarly, on page 3 (top) it is stated that parameter-space variational objectives can be optimized if the variational posterior "is assumed to be a fully-factorized Gaussian distribution". While true, the sentence reads as if that were the only case, but $q$ could also be a variety of other distributions and still remain tractable to stochastic gradient descent.

Author response:
$\textbf{Reply:}$ Sorry for the ambiguity, and we have removed this sentence in the revision.

Reviewer comment (Presentation):
Figure 1/2 only visualize an FBNN baseline. What would a corresponding figure look like for your main competitor, i.e., GWI?

Author response:
$\textbf{Reply:}$ We have tested GWI using Figure 1’s setting and the results are given in Figure 6 (a), (b), and (c) of Appendix E.2. We can see that our FWBI still shows stronger predictive ability to recover the target function and competitive uncertainty estimation in the unseen regions. Since Figure 2 is used to show the ability of our method to use non-GP prior, we did not include them because GWI cannot use BNN non-GP priors.

Reviewer comment (Presentation):
How does the posterior for a GP with each of the three kernels look like for Figure 1/2?

Author response:
$\textbf{Reply:}$ Since we want to highlight that our method (and other functional BNNs) could easily incorporate a meaningful prior via GP contrast to parameter-space BNNs, we give a presumed meaningful GP prior to all functional BNNs in the Figures 1 and 2. To obtain a “meaningful” GP prior, we used all (both training and test) data points to train a GP. Such setting is unfair to GP. Hence, in Figure 7 of Appendix E.3, we change the prior to a GP trained only on training data points and then the results from GP (with three kernels) and our method are compared, where our method has similar large variances in the regions without data but better (smaller) variance on the regions with observed training data.

Reviewer comment (Evaluation):
I am mainly missing some empirical analysis of the proposed objective compared to competitors. The running time is faster, but how does, say, the stability of the training or the convergence speed compare to other methods?

Author response:
$\textbf{Reply:}$ We have added a plot of convergence processes of training loss for our FWBI in Figures 9 and 10 in Appendix E.5 on toy examples and convergence processes for all methods in Figure 15 in Appendix E.8 on UCI regressions. Our FWBI demonstrates better convergence speed and stability compared to other methods.

Reviewer comment (Experiments):
The proposed objective function contains two hyperparameters weighing the contributions from each of the Wasserstein distances, which makes sense practically, although it is a little unsatisfying theoretically. The authors do not discuss how to choose or tune these hyperparameters.

Author response:
$\textbf{Reply:}$ We did not use any hyperparameter tuning method to obtain the optimal ones but only chose the hyperparameters to ensure different components are within the same scale. We believe our method could achieve further better performance after thorough hyperparameter tuning using some black-box optimization methods, like Optuna and scikit-optimize.

Reviewer comment (Presentation):
It is a shame that quite many important results, such as the proof that the objective is an ELBO and some of the empirical analysis, have been put in the appendix. There is, of course, a limit to how much one can fit on 9 pages, but for instance the introduction and preliminaries could perhaps have been shortened a bit to at least make room for a proof sketch. Given that this proof is highlighted as a contribution, at least something should be said about it in the main paper.

Author response:
$\textbf{Reply:}$ Thanks for the suggestion! We did move a lot of content to the Appendix due to the page limitation. As recommended by the reviewer, we move the statement of the theorem to the main body but have to still leave the proof details in Appendix.

Reviewer comment (Theory):
In section 3.2, you call $\boldsymbol{\theta_q}$ and $\boldsymbol{\theta_q}$ "stochastic parameters". It also seems like they are sampled in Eq. (12), but in Eq. (11), they are minimised, which confuses me. Can you elaborate on this?

Author response:
$\textbf{Reply:}$ Sorry for the confusion! $\theta_q$ is the stochastic variational parameter of the optimization problem in Eq. (11). Once the optimal $\theta_q^*: =(\mathbf{\mu}_q^*, \mathbf{\Sigma}_q^*)$ is obtained, we can obtain the approximation posterior predictive function draws by sampling weights $\mathbf{w}^{(j)}$ from $\mathcal{N}(\mathbf{\mu}_q^*, \mathbf{\Sigma}_q^*)$. We have revised the Eq. (12) as follow: 

$q(\mathbf{y}^*|\mathbf{x}^*) \approx \frac{1}{S} \sum_{j=1}^{S} p(\mathbf{y}^*|f(\mathbf{x}^*; \mathbf{w}^{(j)})), \mathbf{w}^{(j)} \sim \mathcal{N}(\mathbf{\mu}_q^*, \mathbf{\Sigma}_q^*), \boldsymbol{\theta_q}^* := (\mathbf{\mu}_q^*, \mathbf{\Sigma}_q^*).$

Reviewer comment (Experiments):
Do you have advice on how to select or tune the hyperparameters in the objective?

Author response:
$\textbf{Reply:}$ We did not use any hyperparameter tuning method to obtain the optimal ones but only chose the hyperparameters to ensure different components are within the same scale. We believe our method could achieve further better performance after thorough hyperparameter tuning using some black-box optimization methods, like Optuna and scikit-optimize.

Reviewer comment (Evaluation):
What does the running time in E.2 measure? Is it training time or prediction time? If training time, is it per epoch or until convergence?

Author response:
$\textbf{Reply:}$ The running time in E.2 is the total training time for all 2000 training epochs.