Paper: Jensen-Shannon Divergence Based Novel Loss Functions for Bayesian Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The paper states that the problem with the KL-divergence in VI-BNNs is due to its asymmetry and unboundedness, which leads to instability in optimization. While instability is indeed a key issue with VI-BNNs, the connection between asymmetry/unboundedness and instability is not clearly shown. The link should be either theoretically or empirically demonstrated so that the proposed JS-based loss can be better motivated.

Author response:
Motivation: Boundedness affects the optimization as seen in the figure in Appendix A. It shows that the KL and the JS-G divergence terms can go to infinity but the JS-A divergence is always bounded. Since the KL term can go to infinity during optimization, the loss function encounters singularity, and as a consequence, the optimization stops. Our implementation shows that the optimization fails at the first epoch when the KL divergence is evaluated numerically through MC sampling even for Gaussian priors. 
On the other hand, non-symmetry does not affect the stability of optimization. We seek symmetry because a symmetric divergence resembles a metric.

Reviewer comment (Writing):
The paper generally lacks polish, with many typos and inconsistent naming (e.g., CIFAR vs Cifar).

Author response:
We have addressed the inconsistencies in the revised version of the manuscript.

Reviewer comment (Theory):
The derivation of the JS-based loss from the constrained optimization perspective—simply replacing `KL(variational_posterior, prior)` with `JS(variational_posterior, prior)`—raises concerns about correctness. Since the form `expected_nll + KL(variational_posterior, prior)` is derived from `KL(variational_posterior, true_posterior)` and relies on properties of KL divergence, it is unclear if the proposed objective is still a valid ELBO. This needs discussion.

Author response:
Presentation: \
a)  The proposed objective maximizes the log-likelihood of observing the data with a constraint that the distribution learned is similar to the assumed prior distributions of the parameters. If we measure the similarity between the learned distribution q and the prior p using KL divergence then we obtain the ELBO. For the JS divergences, the proposed objective is not an ELBO in general. However, we can get back the ELBO from the proposed JS divergence-based losses with alpha = 0 and lambda = 1. \

Reviewer comment (Experiments):
The paper focuses only on generalization performance, but uncertainty quantification (UQ) is equally important in the context of BNNs. Extensive UQ experiments should be performed to make the paper complete.

Author response:
Experiments: \
b) To implement the reviewer’s suggestion, we performed UQ on the histopathology images by the standard BNN and JS divergence-based BNNs. The mean uncertainty quantified by the JS-G and JS-A divergence-based loss is 74% and 29% higher as compared to that of KL. However, these results need to be further analyzed to study the difference between the two. These additional results are reported in Sec.J.5 of the revised manuscript. A detailed study on the UQ aspects of the loss functions is an interesting question that we plan to address in future work.

Reviewer comment (Experiments):
The model's performance on CIFAR10 and CIFAR100 datasets seems suboptimal, potentially due to the absence of batch normalization. It would be worthwhile to explore the impact of incorporating batch normalization and also consider using the Filter Response Normalization (FRN) layer if issues persist.

Author response:
Suboptimal performance: The model’s performance is suboptimal due to the presence of a significant degree of noise (0.1<=sigma<=0.9). Similar performances are obtained for datasets with noise in the literature [Hendrycks et al., ICLR 2019 ]. With lesser degree of noise, the accuracy significantly increases.

Reviewer comment (Experiments):
It is recommended to provide ablation results concerning the parameters $\lambda$ and $\alpha$ in the experimental analysis.

Author response:
Ablation results: Removing the parameters lambda and alpha is equivalent to setting alpha =0 and lambda=1. With this setting, KL divergence-based loss is obtained for the proposed loss functions. Therefore a separate ablation study is not performed in this work.

Reviewer comment (Experiments):
The next concern is the lack of the comparison between JS-A and JS-G. Although the theoretical properties are summarized in Table 1, I could not understand which is better when used in VI. From the theoretical properties, JS-A seems better than JS-G. However, when looking at numerical experiments, such as Fig. 12, JS-G seems better, especially for large-scale experiments. Could you add a comparison for JS-A and JS-G and discuss their differences as VI?

Author response:
Both these divergences have advantages and limitations. The major difference is that JS-A is bounded but doesn't provide a closed-form expression, unlike JS-G which is unbounded but has closed form expression when p and q are Gaussians. Because of the presence of analytical expression, the JS-G-based loss is faster to train as seen in the time comparison results. On the other hand, JS-A eliminates the instability in optimization due to its boundedness property. Both these loss functions can generalize better as shown in the results.

Reviewer comment (Evaluation):
In Fig. 1, although the results of both JS-A and JS-G are presented, they are never discussed. What difference between JS-A and JS-G can we interpret from the figure? Also, the author discussed only the mean. Please discuss the results of the variance in Fig. 1.

Author response:
Fig 1 shows the higher penalization of the JS divergence-based loss functions when the posterior moves away from the prior distribution. The value of JS-G>KL when sigma_p > sigma_q. The variance figure of JS-G is a depiction of theorem 3. It shows the change of JS-G with respect to KL as sigma is changed.  In the case of JS-A divergence, the value of KL>JS-A when lambda = 1. However, by increasing lambda, we obtain a higher penalization of the loss which is depicted in Fig 1. Sec.3.4.1 and Sec 3.4.2 of the revised manuscript are modified to reflect these comments.

Reviewer comment (Evaluation):
In Table 4, the proposed JS-A seems slower than the other methods in terms of wall clock time. Why?

Author response:
The JS-A divergence is slower because it does not have a closed-form expression and it has to be implemented through MC sampling as mentioned in Appendix K of the manuscript.

Reviewer comment (Experiments):
According to the explanation of Eq(18), the proposed method has both mode- and mean-seeking properties. I think this resembles alpha divergence VI. Could you show the comparison of alpha-divergence VI numerically?

Author response:
To implement the reviewer's comments we performed a comparison study on multiple regression datasets for various state-of-the-art VI techniques for BNNs including alpha divergence VI. The results are provided in Appendix N of the revised manuscript. We found that the JS divergence-based losses perform as good as the other techniques or in some cases better.

Reviewer comment (Theory):
In this paper, the authors adopted the problem that the KL-based VI objective used in BNN, such as [Blundell et al., 2015], is unbounded and lacks symmetry, which negatively impacts optimization stability and generalization performance. As the basis for this assertion, they reference [Hensman et al. (2014); Dieng et al. (2017); Deasy et al. (2020)]. However, in my opinion, these references do not explicitly suggest that the unbounded and non-symmetric nature of the objective affects generalization performance. For instance, [Dieng et al. (2017)] identify the decrease in generalization performance due to **the underfitting of posterior variance** by KL-based VI and demonstrate that minimizing the upper bound based on the chi-squared distribution (called CUBO) can resolve this issue, and it is not stated that this underfit of posterior variance is caused by the unbounded and asymmetric properties of KL objective. This raises significant concerns regarding the validity of the problem setting adopted in this paper unless it is theoretically and/or empirically demonstrated that the unbounded and non-symmetric objectives affect the stability of optimization and generalization performance of VI. There still be room for reconsidering whether the issues of divergence boundedness and asymmetry should be approached in the context of BNN+VI rather than in the context of divergence measures for distributions.

Author response:
Boundedness affects the optimization as seen in the figure in Appendix A. It shows that the KL and the JS-G divergence terms can go to infinity but the JS-A divergence is always bounded. Since the KL term can go to infinity during optimization, the loss function encounters singularity, and as a consequence, the optimization stops. Our implementation shows that the optimization fails at the first epoch when the KL divergence is evaluated numerically through MC sampling even for Gaussian priors. We have not observed any direct effect of lack of boundedness on the generalization performance.
Non-symmetry does not affect the stability of optimization. We seek symmetry because a symmetric divergence resembles a metric. However, we agree with the reviewer that lack of symmetry does not affect the stability of optimization and generalization performance. 
The Introduction section of the manuscript is changed to reflect these facts.

Reviewer comment (Novelty):
If we utilize several practical approaches proposed in the context of PAC-Bayes, such as those suggested by [DR18] and [P21] (bounded cross entropy) or appropriately initializing model weights and prior, it becomes possible to bound the KL-based objective in the case of the Gaussian distribution that is the focus of this paper. Therefore, it is necessary to carefully consider whether the approach to resolving the issue of unbounded KL-based objectives by altering the upper-bounded divergence is valid. However, the merits of the proposed method using JS divergence have not been discussed in comparison to these approaches, thus the validity of the approach remains unconfirmed.

Author response:
Yes, it is possible to keep the KL divergence bounded by selecting appropriate priors. However, that imparts restrictions on the priors that can be used. Techniques found in [DR18; P21] provide bounds on the KL divergence only for specific distributions (exponential distributions) whereas, the boundedness of JS-A divergence shown in the present work is true for any arbitrary distributions. In addition, bounding the cross-entropy loss ( [DR18; P21]) does not resolve the unboundedness of the divergence part of the loss function, which is the goal of this work.

Reviewer comment (Experiments):
If my understanding is correct, Theorems 2 and the corollary demonstrate that the proposed objective function becomes tight. However, as numerical experiments to verify this have not been conducted, the validity of the theoretical analysis remains unconfirmed. In this regard, plotting the progression of ELBO and the proposed objective values for each iteration would quickly provide confirmation.

Author response:
To implement the reviewer's suggestion, the evolution of the divergence and NLL part of the loss function during optimization is shown in Fig 13 in Appendix L of the revised manuscript for the histopathology dataset. It is aligned with the theoretical results that the JS-G divergence penalizes higher than the KL divergence when the distribution q is farther from p. It is also seen that the negative log-likelihood is higher for the JS-G divergence as compared to the KL divergence. However, for the test dataset, the negative log-likelihood is lower for the JS divergences than the KL divergence which is desirable.

Reviewer comment (Evaluation):
The authors solely compare the performance of their proposed method and existing methods based on predictive accuracy. However, since BNNs are probabilistic models, it is crucial to consider how well they capture the uncertainty in predictions, which is an important performance metric in many BNN studies, as confirmed experimentally by researchers such as [I21, P21]. To claim state-of-the-art performance, a more comprehensive performance evaluation on the test data, including metrics like test-ELBO, test-NLL, and Expected Calibration Error (ECE), is necessary.

Author response:
To implement the reviewer’s suggestion we now compare the Test NLL, Test loss, and ECE for the three losses for the histopathology dataset. This is provided in Table 5 in Appendix M. It is observed that the proposed JS divergences perform better in terms of both ECE and negative log-likelihood than the KL divergence for the histopathology dataset.

Reviewer comment (Evaluation):
There is no information about how predictions are being made. Is this an evaluation based on deterministic predictions using the mean parameters of the variational distribution? Or is it a probabilistic prediction through Bayesian model averaging (BMA) using samples from the variational distribution? Since BNNs are probabilistic models, performance evaluation through BMA is essentially necessary.

Author response:
Samples from the posterior distribution are considered while training the Bayesian NN. However, we are predicting the labels of the test set using the mean of the parameters. To implement the reviewer’s suggestion we have compared the predictions obtained with the mean and the BMA. The result shows that the predictions by both approaches differ by about 0.006 % in accuracy for the histopathology dataset.

Reviewer comment (Experiments):
The initial values for network parameters are taken from pre-trained neural networks on ImageNet. However, in actual experiments, the authors use datasets such as CIFAR-10/100 and they do not employ ImageNet. This might seem somewhat unusual, as the reasons for this decision are not clearly articulated.

Author response:
This was done to get a better initialization of the weights of the network for faster convergence. However, the network parameters are not frozen and hence are allowed to be learned during optimization. We found that initializations with pre-trained network weights performed better than random initialization.

Reviewer comment (Writing):
Sec.1, 2th paragraph: What does it mean "the two most commonly used techniques to approximate them (i.e., Bayesian inference) are the Variational Inference (VI)?"

Author response:
RMisc: All the miscellaneous comments are addressed in the revised manuscript.

Reviewer comment (Novelty):
Sec.1, 2th paragraph: Citation for VI and MCMC could be required, e.g., [JGJS99]

Author response:
RMisc: All the miscellaneous comments are addressed in the revised manuscript.

Reviewer comment (Writing):
Sec.3.1, the first sentence: The Generalized JS... --> The generalized JS...

Author response:
RMisc: All the miscellaneous comments are addressed in the revised manuscript.

Reviewer comment (Writing):
Sec.4.1, the first sentence: ...on two Data sets... --> ... on two data sets...

Author response:
RMisc: All the miscellaneous comments are addressed in the revised manuscript.

Reviewer comment (Experiments):
The loss relies on 2 hyper-parameters (lambda, alpha) which might make the practical use of the method problematic (e.g., in cases where hyper-parameters search is too expensive).

Author response:
This fact is mentioned in the limitation section of the manuscript. When alpha = 0 and lambda =1, we get back the KL divergence-based loss. This can be utilized in such cases where hyper-parameters are expensive. However, our empirical experiments show that finding the optimal values of the two additional parameters is not significantly more expensive. The addition of these two new hyperparameters to the existing set of hyperparameters (such as learning rate and batch size) does not add significant computational cost.

Reviewer comment (Experiments):
It is unclear if simpler regularization methods (e.g., dropouts) might yield similar improvement.

Author response:
The key contribution of this work is to overcome the unboundedness due to KL divergence in the state-of-the-art BNNs. In doing this, we observe better regularization through theoretical and empirical analysis of the loss functions. Since the motivation for this work is not to improve the regularization performance, a comparison with other regularization methods is not performed.  In addition, BNNs automatically regularize through model averaging and the prior distribution. For example, when Gaussian prior is used, it yields L2 regularization.

Reviewer comment (Experiments):
Experimentation is limited to two simple datasets.

Author response:
To implement the reviewer's comments we performed a comparison study on multiple regression datasets for various state-of-the-art VI techniques for BNNs. The results are provided in Appendix N of the revised manuscript. We found that the JS divergence-based losses perform as good as the other techniques or in some cases better in terms of both accuracy and log-likelihood for these datasets.

Reviewer comment (Experiments):
Comparison to other simpler regularization methods (e.g., dropout) is missing.

Author response:
The key contribution of this work is to overcome the unboundedness due to KL divergence in the state-of-the-art BNNs. In doing this, we observe better regularization through theoretical and empirical analysis of the loss functions. Since the motivation for this work is not to improve the regularization performance, a comparison with other regularization methods is not performed.  In addition, BNNs automatically regularize through model averaging and the prior distribution. For example, when Gaussian prior is used, it yields L2 regularization.

Reviewer comment (Theory):
In the introduction: The number of parameters has nothing to do with the ability to provide a robust measure of uncertainty. The model choice allows such ability.

Author response:
Changed to: due to deterministic parameters, CNNs cannot provide a robust measure of uncertainty.

Reviewer comment (Theory):
In the introduction: The variational distribution is learned by minimizing its dissimilarity with respect to the true posterior - VI does not minimize this quantity since the true posterior is unknown. It can be derived starting from the KL divergence between the approximate and true posterior though.

Author response:
Changed to: The variational distribution is learned by minimizing an objective function derived from its dissimilarity with respect to the true posterior

Reviewer comment (Theory):
Derivation of Eq. 5 from Eq. 4 is awkward and the justification for dropping p(D) is not clear (being a constant that multiplies the approximate posterior). Typically derivation entails the use of Jensen’s inequality which also facilitate the relation between Eq. 5 and Eq. 4.

Author response:
Modified to explain it better

Reviewer comment (Theory):
Near Eq. 2: Gα(x, y) = x1−αyα - G_alpha is an unnormalized distribution in the general case, is that true?

Author response:
Yes, the geometric JS divergence is unnormalized in the general case.

Reviewer comment (Theory):
Near Eq.3: JS-G(p||q)|α = JS-G(p||q)|1−α type, should it be non equality?

Author response:
This is equality. Symmetry in this context means JS-G(p||q)|α = JS-G(q||p)|α

Reviewer comment (Experiments):
Figure 2: it is unclear whether other regularization techniques (like dropouts) could narrow down the gap for VI with KL divergence, without the need for the forward KL regularization term in the loss. Also, 2b looks like overfitting for the VI with KL divergence. Did you try other regularization methods?

Author response:
BNNs automatically regularize through model averaging and the prior distribution as mentioned earlier. Therefore, we did not try additional regularization methods.