Flat Seeking Bayesian Neural Networks

Summary

This paper proposes modifying the loss used for Bayesian neural networks (BNNs) to take into account the sharpness/flatness of the loss with respect to the model parameters. Theory, based on that of sharpness-aware minimization (SAM), is developed to propose this loss modification. Making BNNs sharpness-aware led to increased accuracy, decreased negative log-likelihood (NLL), and decreased expected calibration error (ECE) for large neural network models across CIFAR-10, CIFAR-100, and ImageNet.

Strengths

This paper does an excellent job of combining sharpness-aware and BNN methods. The theory for this is given and the theory is applied to several BNN method, namely stochastic weight averaging Gaussian (SWAG), stochastic gradient Langevin dynamics (SGLD), Monte Carlo (MC) dropout, and deep ensembles, for CIFAR-10 and CIFAR-100. These experiments show reasonable improvements.

Weaknesses

1. The difference between the optimization problems used in the different Bayesian methods was not clear enough. The main optimization problem discussed was for variational inference. To my knowledge, however, SWAG and SGLD typically use a procedure that is based on the gradient of the log of the posterior. 2. The incomplete ImageNet experiment. Only SWAG and FSWAG are evaluated for ImageNet, versus all of the methods that were applied to CIFAR-10 and CIFAR-100.

Questions

1. Are different optimization problems used for some of the studied methods? If so, please add those details to Section 3. 2. Is it possible to have results on ImageNet for more of the studied methods? 3. I'd recommend putting the number of repetitions used to create the confidence intervals in the captions of the tables. I see that, in the supplementary material, it was three, but I think that this information, along with the number of samples used for the methods, should be in the main text. 4. How do you think that sharpness-aware methods would interact with Laplace approximation methods, another popular BNN method? (This is not a necessary addition to the paper, in my opinion. However, I am interested in your answer to this.)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Limitations and societal impact are not discussed. One limitation that could be added to the conclusion is that the proposed Gaussian variational approach only uses a diagonal covariance for the posterior approximation. While this is standard, it is a limitation that, by mentioning it, could show areas of potential improvement for the proposed variational method.

Summary

This paper introduces a method called Sharpness-Aware Bayesian Neural Networks (SABNN) that aims to improve generalization performance on test datasets. The key idea is to replace the negative empirical loss function with the negative SAM (Sharpness-Aware Minimization) [1] loss function, enabling consideration of the flatness of the loss surface. The authors provide both theoretical analysis and empirical evidence to demonstrate that SABNN achieves better generalization compared to existing methods when evaluated on test datasets. References [1] Foret, P., A. Kleiner, H. Mobahi, et al. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations. 2021.

Strengths

Originality - The paper presents a theoretical demonstration of the robustness of sharpness-aware posterior inference for the entire distribution of the true dataset. - The authors propose a straightforward variational approach that utilizes a Gaussian approximate posterior for efficient and straightforward inference. Clarity - The paper is well-written and provides clear explanations, making it easy to understand and follow.

Weaknesses

Method - The method described in the paper involves a straightforward and nearly direct modification of the likelihood function by replacing the negative empirical loss with the negative SAM loss function. Experiment - In section 4.3 of the paper, the authors mention that the sharpness scores are provided in the supplementary material. However, I was unable to locate the results in the supplementary material. - The paper does not discuss any potential additional computational costs associated with the proposed method in comparison to the baselines.

Questions

Experiment - It is requested to provide the sharpness scores and the largest Hessian eigenvalues for the entire dataset, as well as for the SWAG and F-SWAG models. I think the eigenvalues can be calculated for the mean $\mu$ of each method. - There is a desire for a discussion regarding potential additional computational costs associated with the proposed method in comparison to the baselines. - The results of the Deep Ensemble, which is an ensemble of independently trained models without KL in Bayesian Deep Ensemble and with L2 regularization loss, trained with SAM loss, are also requested. The aim is to assess whether the proposed SABNN method is practically beneficial or not. I am willing to increase my score after considering the experiment results provided in the rebuttal.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Limitation - This paper did not address the limitation of the proposed method. - One possible limitation can be additional computational costs.

Summary

This paper extends SAM -- sharpness aware minimization framework of Foret et al, who seeks parameters of Neural Networks in the regimes of flat loss landscape. The current understanding is that a model that exhibits a flat loss landscape exhibits better generalization performance. Particularly in this paper, such methods are extended to variational inference for neural networks, where the posterior (and hence neural network training) is obtained through an objective function that seeks flat local minima. Experiments are conducted, which were designed to incorporate their objective function into the current Bayesian Neural Networks. The results show consistent improvements over the vanilla counterparts.

Strengths

The contribution of the paper seems interesting. The framework is utilizing PAC-Bayes like theory, where an upper bound to the true loss function is obtained through a careful balance between empirical loss and certain KL divergence. The later regularizes the Bayesian neural network overfitting. It is interesting to see, how variational inference fits the true posterior, which is in turn modified according to the flat local minima seeking objectives. As SAM is one of the popular methods to understand generalization (while being simple), such extensions to Bayesian Neural Networks might be meaningful to the community.

Weaknesses

I am not an expert in this specific area of combining geometric insights into variational inference. I lean to accept the paper due to my inability to fully grasp this concept. I hope that other reviewers can provide more meaningful feedback for improvements here. For the weakness: - presentation could perhaps improve a bit more. 1. in the abstract, I do not see the logical connection between generic explanations of Bayesian Neural Networks, and generalization. The switch was to me, a bit disconnected and sounded a bit sudden. In the introduction, I also did not see much the role of the 2nd paragraph. Indeed, there have been many variation inference approaches, but I could not see the connection between the expressivity of the posterior, and the approaches of seeking flatness in variational inference. 2. in the theoretic part, I see several equations that repeat also in other papers, e.g., the upper bound of true loss with empirical and KL divergence, how the true posterior provides tight generalization bound, etc. While I appreciate the author's efforts, especially in the appendix for many derivations, it may also make sense to provide citations more in the corresponding texts, so that the readers may be able to learn also through other papers. 3. lines 170-180 should be extended a bit more. I could not understand the connection between each question on page 5 and also missed insights between each derivation. 4. It might be helpful to include an algorithm part, so that the readers may better comprehend how the method works (and hopefully also why). - empirical improvements seem rather minor, and it was not clear to me, how these experiments demonstrate generalization. While I appreciate the ample empirical evidence provided by the paper, most of the improvements in terms of accuracy are less than 1 percent. This may raise a question if the experimental design here is adequate to show the generalization performance. Table 3, 4, and 5 need error bars, which is requirements in Bayesian machine learning works. On how to test generalization, I think the paper could adapt few-shot settings or experiments from meta-learning. Therein, little training data is provided, and one may be able to test better how the model still generalizes. I am not sure, how testing on vanilla cifar10, cifar100 or imagenet classification score is an indication of generalization performance.

Questions

--

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

3 good

Limitations

I could not find the limitation section in the main body of the paper.

Summary

The paper proposes a new approach to posterior inference for Bayesian neural networks that takes into account the sharpness/flatness of deep learning models, leading to better generalisation ability. The authors introduce the Sharpness-Aware Posterior (SA-Posterior), which allows the sampling of a set of flat models that improve model generalisation. The paper presents a theoretical framework for the SA-Posterior, including a Bayesian framework and a variational inference approach. The authors demonstrate the effectiveness of their approach through experiments on various datasets, showing that it outperforms existing methods on all metrics of interest.

Strengths

- The paper presents a novel approach to posterior inference for Bayesian neural networks. - The theoretical framework for the SA-Posterior is developed, including a Bayesian setting and a variational inference approach. - The effectiveness of the SA-Posterior is demonstrated through experiments on various datasets. - The paper provides insights into the importance of sharpness/flatness in deep learning models. - The paper contributes to the growing body of research on improving the generalisability of deep learning models.

Weaknesses

- The experiments could be extended to include more datasets and models. - The paper does not discuss the computational complexity of the SA-Posterior approach. - The paper does not compare the SA-Posterior approach with other recent approaches to improve the generalisation ability of deep learning models.

Questions

- How does the SA-Posterior approach compare to other recent approaches to improve generalisation in deep learning models? - How does the computational complexity of the SA-Posterior approach compare with other approaches to posterior inference for Bayesian neural networks?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

- Investigate the potential of the SA posterior approach for other types of neural networks beyond Bayesian neural networks. - Investigate the computational complexity of the SA-Posterior approach and develop methods to reduce it. - Investigate the potential of the SA-Posterior approach for transfer learning and domain adaptation.

Summary

This paper introduce theories in Bayesian settings and propose variational inference for the sharpness aware posterior in the context of Bayesian Neural Network. The proposed approach is incorporated with existing state of the art Bayesian Neural Networks and experiments were conducted to show the effectiveness of sharpness aware posterior. The results indicate that incorporating sharpness aware posterior methodolody outperforms baseline in terms of ensemble accuracy , expected callibration error and negative likelihood. Also, experimensts show that models from poroposed approach are less sensitive to noise and have improved generalization ability

Strengths

1) The paper proposes to approximate posterior of Bayesian Neural Networks with sharpness aware posterior. 2) The proposed approach can easily be incorporated to already existing Bayesian Neural Networks 3) The experiments use state-of-the-art settings 4) The paper outperforms the baseline methods

Weaknesses

1) The experments show that proposed approach most of teh times outperforms baseline methods , however, improvement is not that significant. 2) The experimenst are only for image-classification task

Questions

Q1) What is the computational head of the proposed approach?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

It would be great to add compared to baseline approaches , what is the computational overhead of the proposed approach.

Please kindly let us know if our additional experiments address your concerns

Dear Reviewer, Could you please look at our rebuttal and kindly let us know if the additional experiments address your concerns? Really appreciate your time and effort to review our paper. Regards, Authors of the paper 6164

Response to the authors

I would like to thank the authors for the efforts. I have read the rebuttal and comments from other reviewers. I decided to stand with the score of borderline. More thoughtful experiments, which better showcase the strength of the methods could make the paper more relevant to the community.