Improving Generalization with Flat Hilbert Bayesian Inference

We appreciate the reviewer’s valuable feedback and will address the concerns as follows:

• Regarding the originality: While the reviewer correctly noted that the RKHS framework has been used in previous works on variational families, this is not the main contribution of our paper. The first novelty of our work is that to improve generalization, we propose in Section 4 a novel Bayesian framework that aims to approximate the general posterior $p(\theta|D)$ instead of the empirical posterior $p(\theta|S)$ like prior works such as SVGD [1]. In doing so, the model sampled from the final posterior will fit the whole data distribution $D$ instead of overfitting to the specific dataset $S$. Notice that we do not have access to this general posterior because $D$ is unknown. To address this general posterior, our second contribution is that we develop the concept of functional sharpness on the RKHS from Theorem 1, and propose to apply this concept to the context of Bayesian Inference. In short, our main contribution is not the usage of the RKHS in Bayesian Inference, but the introduction of the concept of sharpness on the RKHS, and the usage of this concept in Bayesian Inference.
• Clarification about the empirical and the general posteriors: Recall that prior works in Bayesian typically define the likelihood as $p(y_i|x_i,\theta) = \exp(-\ell(f_\theta(x_i),y_i)),$ which gives the posterior as the reviewer mentioned $p(\theta|S)\propto\exp(-nL_S(\theta))p(\theta)$ Our work slightly modifies the likelihood by scaling with the dataset size: $p(y_i|x_i,\theta)=\exp(-\frac{1}{n}\ell(f_\theta(x_i),y_i)).$ We emphasize that this formulation has been employed by prior works such as SA-BNN [3]. Moreover, the modification does not change the intuition of the original formulation that a smaller loss yields a higher likelihood and, therefore retains its validity. With this formulation, the empirical posterior becomes: $p(\theta|S)\propto\exp(-L_S(\theta))p(\theta)$ Then, we can define its general posterior counterpart $p(\theta|D)\propto\exp(-L_D(\theta))p(\theta)=\exp(-\mathbb{E}_{(x, y)\sim D}[\log p(y|x,\theta)])p(\theta),$ which is valid because this definition does not contain $n$.
• Regarding the benefits of addressing the general posterior: To understand the benefits of addressing the general posterior, let us first consider the method proposed by SAM [1]. SAM seeks to minimize $L_D(\theta)$ instead of $L_S(\theta)$, thereby reducing overfitting by fitting to the data distribution rather than a specific dataset. Since the $D$ is unknown, $L_D(\theta)$ cannot be computed. Then, SAM has to implicitly minimize $L_D(\theta)$ with a sequence of approximations to the worst-case empirical loss within a radius. Our approach extends this idea of SAM to the regime of distributions. As discussed in Section 4: rather than approximating $p(\theta|S)$ as done in conventional Bayesian methods, we aim to approximate $p(\theta|D)$. By doing so, the posterior fits the whole data distribution $D$ rather than overfitting to a specific dataset $S$, so the model particles sampled from this posterior will not overfit to the specific dataset $S$ and hence improve generalization upon prior Bayesian techniques. Similar to SAM, our approach requires approximations to address the general loss/posterior. Nevertheless, this approximation error can be mitigated with the proper choice of the hyperparameter $\rho$, which is typically a small constant.
• Regarding the claim of improving generalization: As discussed above, the feature showing that FHBI improves generalization is that it approximates $p(\theta|D)$ instead of $p(\theta|S)$. Then, the model particles sampled from the final posterior follow the true distribution instead of overfitting to the dataset $S$. This improvement is also evident in our Ablation studies. As noted by SAM, lower sharpness correlates with better generalization. Section 6.2 demonstrated that FHBI reduces the sharpness of each particle compared to SVGD and also prevents particles from collapsing into a single mode. Together, these properties support our claim that FHBI leads to better generalization.
• Regarding other minor issues: We appreciate the suggestions of the reviewer about the minor issues and have incorporated them in the revision.

References:

[1] Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A general purpose Bayesian inference algorithm. Advances in neural information processing systems, 29, 2016a.

[2] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.

[3] Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat-seeking Bayesian neural networks. Advances in Neural Information Processing Systems, 2023.

We thank the reviewer for their thoughtful and detailed response. We would like to address the concerns as follows:

1. Clarification on "overfitting posterior": We wish to clarify that, throughout the main text and this discussion, when we state that the method addresses overfitting by approximating the population posterior $p(\theta|D)$, we are referring to alleviating the overfitting issue associated with the ensemble of particles rather than the posterior itself. To further address the reviewer’s concern theoretically and verify the theoretical advantages of targeting $p(\theta|D)$, we have updated the rebuttal revision to include Proposition 1 at the beginning of Section 4. The statement of the proposition is as follows:

Consider the problem of finding the distribution $\mathbb{Q}$ that solves:

$\mathbb{Q}^*=\min_{\mathbb{Q} \ll \mathbb{P}}E_{\theta \sim \mathbb{Q}}[L_D(\theta)] + KL(\mathbb{Q} || \mathbb{P})$

where we search over $\mathbb{Q}$ absolutely continuous w.r.t the prior $\mathbb{P}$, and the second term is the regularization term. The closed-form solution to this problem is exactly the population posterior $p(\theta|D)$.

In this proposition, we aim to identify the posterior distribution that minimizes the expected population loss, where the expectation is taken over the entire parameter space with $\theta\sim\mathbb{Q}^*$, while staying sufficiently close to the prior distribution to ensure simplicity. With access to $\mathbb{Q}^*$, we can sample a set of particles whose average performance optimally minimizes the population loss. Since the solution to this optimization problem corresponds exactly to the population posterior, the ensemble of the particles sampled from $\mathbb{Q}^*\equiv p(\theta|D)$ effectively minimizes the average value of the population loss. This is because $\mathbb{Q}^*$ is explicitly chosen to minimize the expected value of the population loss $L_D$, which means the ensemble fits the whole data distribution instead of overfitting to the specific dataset $S$, therefore establishes improved generalizability. Consequently, this proposition theoretically asserts that sampling from $p(\theta|D)$ improves the generalizability of the ensemble.

2. Regarding the practical approximations: Our theoretical derivation relies on two approximations, both of which are widely established practices in prior works:

• First approximation: We implicitly minimize the KL divergence to $p(\theta|D)$ by minimizing the worst-case loss within a neighborhood $\max_{|f' - f|\leq\rho}\text{KL}(q_{[I+f']}\|p(\theta|S))$. This step is necessary and inevitable because the dataset distribution $D$ is unknown. Importantly, we emphasize that the Euclidean analog of this approach is widely used in almost all sharpness-aware optimization techniques, including SAM or Fisher-SAM.

• Second approximation: This occurs in Eq. (8), where we apply a first-order Taylor expansion. Such approximations are a standard practice in first-order optimization methods, including SGD and particularly SAM. A potential concern might arise because the perturbation and approximation are conducted in an RKHS, which may behave differently compared to perturbations in Euclidean space (as in SAM). However, it is noteworthy that the RKHS in question represents the space of functions $f$, which encode the movements of particles residing in Euclidean space. Therefore, perturbations in the RKHS are directly tied to the perturbations of individual particles. Consequently, higher-order terms like $O(\rho^2)$ in the RKHS can be safely omitted, just as they are in Euclidean space with SAM.

In summary, the derivation relies on only two approximations, both of which are well-established in the literature on first-order methods and sharpness-aware optimization. These approximations are widely used practices and can be safely applied without introducing significant deviations.

3. Regarding the isolation of the contribution of approximating $p(\theta|D)$: We would like to emphasize that the comparison between FHBI and SVGD already isolates the contribution of approximating $p(\theta|D)$, both theoretically and empirically. Specifically, if we modify the problem formulation in Eq. (6) by replacing $p(\theta|D)$ with $p(\theta|S)$, the ascending step (Eq. (10)) is no longer present, effectively reducing FHBI to SVGD. Thus, this comparison inherently represents the isolation of the contribution of approximating $p(\theta|D)$.

Finally, we note that while addressing the reviewers' concerns, the theoretical development of this paper has been further strengthened and refined. We sincerely thank the reviewer for their insightful and constructive comments, which have greatly enhanced the final revision. Please let us know if there are any remaining concerns or points that need further clarification.

We would like to thank the reviewer for the useful feedback. We would like to address the concerns as follows:

• Regarding the connection with SAM [1]: While our method is not directly derived from SAM, we frequently reference SAM due to the meaningful conceptual connections between the two approaches. As shown in Equations (6)–(8), our method performs functional sharpness-aware minimization (SAM) in the RKHS. As we have mentioned, FHBI reduces to SAM when $m=1$. The key insight behind this fact is that FHBI is performing sharpness-aware minimization on the RKHS with $f$ being the parameter, and since $f$ captures the transportation of the particles, it indicates that FHBI is simultaneously performing SAM on all particles. Moreover, as demonstrated in Section 4, FHBI not only controls the sharpness of all particles via the function $f$, but also controls the interaction between the particles to further reduce sharpness and promote particle diversity to further improve generalization. Consequently, in the single-particle case $m=1$, FHBI naturally reduces to SAM, aligning with the intuitive understanding of its behavior in this special case.

• Regarding the connection with SVGD [2]: This connection arises from the interpretation of FHBI as performing SAM in the RKHS, while SVGD is analogous to SGD in the RKHS. When $\rho = 0$, SAM reduces to SGD due to the absence of an ascent step. As a result, FHBI simplifies to SVGD in this case. However, we emphasize that this connection is based on the reduction of SAM to SGD when $\rho = 0$, but cannot be inferred from Theorem 1, as $h(1/\rho^2)$ becomes undefined when $\rho = 0$.

• Regarding the choice of the kernel: In our experiments, we utilized the RBF kernel, renowned for its strong representational capabilities and its ability to balance underfitting and overfitting through the kernel width $\sigma$. However, based on the reviewer’s suggestion, we conducted an additional ablation study to evaluate performance with alternative kernels. We tested our method with a polynomial kernel of degree $10$ on the Specialized datasets. The results, presented below, show that while the polynomial kernel slightly underperforms compared to the RBF kernel, the performance gap is negligible. We have addressed the reviewer's suggestion and included this experiment in Appendix B of the revised rebuttal.

• Regarding the insight on $\rho$ and $\epsilon$: While the analogy on the Euler-Maruyama method for solving ODE provided by the reviewer is insightful and thought-provoking, it is noteworthy that smaller $\rho$ and $\epsilon$ do not necessarily lead to better performance. From the perspective of solving a minimax optimization problem, $\rho$ and $\epsilon$ serve as the step sizes/learning rates of the inner maximization and outer minimization problems, respectively. While smaller values of the learning rate $\epsilon$ ensure convergence, they do not guarantee better performance due to the potential for getting stuck in local minima. Similarly, the ascent step size $\rho$ requires careful tuning: excessively large values can destabilize training, while very small values may render the ascent step ineffective, thus failing to enhance generalization. We also emphasize that these principles also hold for SAM. In summary, smaller step sizes with more iterations may promote convergence but do not ensure better performance due to the non-convex nature of the loss function and the risk of local minima in practice.

• Regarding the choices of $\rho$ and $\epsilon$: Since $\rho$ and $\epsilon$ govern distinct optimization problems—$\rho$ for the inner maximization and $\epsilon$ for the outer minimization—they are tuned separately. Based on the above interpretation of $\rho$ and $\epsilon$ as learning rates, we tuned these hyperparameters using a grid search within a reasonable range. Details of this process are provided in the Experimental section and Appendix C in the rebuttal revision.

• Regarding Section 6.2: The goal of Section 6.2 is to empirically study the interaction between particles, supporting the claims made at the end of Section 4. This setting is specifically designed for multi-particle methods, so comparisons with SAM, which inherently operates in a single-particle framework, are not applicable in this context.

References:

[1] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.

[2] Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A general purpose Bayesian inference algorithm. Advances in neural information processing systems, 29, 2016a.

Thank you for responding to my questions. Although I was a little disappointed that the proposed choice of $\rho$ and $\epsilon$ could not further promote FBHI, the explanation regarding $\rho$ and $\epsilon$ was very convincing as well. Thank you also for testing a different kernel.

With regard to the connection with SAM, while I understand the authors' point and intention, I still feel that there might be a better way to reorder the explanations--at the same time, I feel this might also be a matter of this reviewer's personal preference. I am thankful for the authors' additional efforts, and I would like to keep my score as is to show my support and I am willing to defend the acceptance of this work within my capability.

Dear reviewer yQt2,

We thank you for your response and for maintaining a rating of 8. Your comments helped significantly improve the final revision. Please feel free to let us know if you have any further concerns.

Best,

The Authors

Here is the comparison between our method, SVGD+SAM (SA-BNN), and SGLD+SAM (SADA-JEM). These results have also been included in the rebuttal revision.

Accuracy:

ECE:

• Regarding other minor issues: We appreciate the reviewer pointing out the typo and the notation for the distributions and we have also included those suggestions in the revision.

Once again, we would like to thank the reviewer for the insightful feedback for an improved revision.

References:

[1] Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat-seeking Bayesian neural networks. Advances in Neural Information Processing Systems, 2023.

[2] Germain, P., Lacasse, A., Laviolette, F., & Marchand, M. (2006). PAC-Bayesian theorems for domain adaptation. In Advances in Neural Information Processing Systems (NeurIPS).

[3] Bishop, C. M. (2006). Pattern Recognition and Machine Learning. New York: Springer.

[4] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https://openreview.net/forum?id=6Tm1mposlrM.

[5] Feng, Y., & Li, C. (2022). Fast Stein Variational Gradient Descent for Neural Network Sampling. IEEE Transactions on Neural Networks and Learning Systems.

[6] Kim, Minyoung, Da Li, Xu Hu, and Timothy Hospedales. "Fisher SAM: Information Geometry and Sharpness Aware Minimisation." Presented at ICML 2022.

[7] "AdaptiveSAM: Towards Efficient Tuning of SAM for Surgical Scene Segmentation." Preprint available on arXiv, 2023.

[8] Yang, Xiulong, Qing Su, and Shihao Ji. "Towards Bridging the Performance Gaps of Joint Energy-based Models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

We sincerely thank the reviewer for their valuable and constructive feedback. Below, we address the concerns raised in detail:

• Regarding the notations of the distributions: Even though these notations do not affect the theoretical development since they are not used throughout the derivation, we thank the reviewer for the comment and have incorporated this suggestion in the revision.

• Regarding the derivation of the empirical posterior: First, by applying Bayes' rule and assuming that the samples are i.i.d., we obtain

$p(\theta|\cal{S}) = p(\theta)p(\cal{S}|\theta)/p(\cal{S}) \propto p(\theta)p(\cal{S}|\theta) = p(\theta)\prod_{i=1}^n p(y_i | x_i, \cal{S}, \theta).$

We define that the likelihood is proportional to $\exp\left(-\frac{1}{|\cal{S}|} \ell(f_\theta(x), y)\right)$, where $\ell$ is a loss function such as the Cross-Entropy loss, and $f_\theta$ represents some sufficiently expressive model. Then, the posterior has the form:

p(\theta|\cal{S}) \propto p(\theta)\prod_{i=1}^n \exp(\frac{-1}{n}\ell(f_\theta(x_i), y_i)) = p(\theta) \exp(-\cal{L}_\cal{S}(\theta)) .

We emphasize that the formulation that expresses likelihood as the exponential of the negative loss value that we used above is a popular formulation that has been widely employed in prior works in learning theory and Bayesian Inference such as [1] and [3]. The intuition behind this formulation is that when $f_\theta$ is sufficiently representative to capture the likelihood, a smaller loss value results in a larger exponential term, indicating a higher likelihood, and vice versa.

• Regarding the usage of the terminology "general loss": The term \cal{L}_\cal{D} is referred to by various names in the optimization and statistical learning theory literature, including population loss and generalization loss, as the reviewer mentioned, and general loss that we used which has been employed in several previous works on learning theory and PAC-Bayes theory, including those cited by the reviewer [1, 2]. In response to this suggestion, acknowledging that the term population loss is more commonly used in learning theory, we have replaced general loss with population loss in the final revision to enhance readability. We appreciate this insightful feedback.

• Regarding the scale of $\rho$: In the theory of sharpness-aware minimization, $\rho$ represents the size of the neighborhood, which is also the size of the ascending step. We emphasize that similar to SAM, $\rho$ is a hyperparameter that we will tune in practice, and is typically a small constant. However, if $\rho$ is excessively small, the term $\max_{|\theta' - \theta| \leq \rho} L_S(\theta')$ also gets smaller and becomes a closer approximation of $L_S(\theta)$; consequently, a larger error term (which is $h(1/\rho^2)$) is required to ensure that inequality is maintained. Conversely, as $\rho$ increases, the term $\max_{|\theta' - \theta| \leq \rho} L_S(\theta')$ grows due to the larger neighborhood, necessitating a smaller error term to maintain the validity of the inequality. In short, $\rho$ is a small constant but is not excessively small. This intuition is also reflected in our practical algorithm as well as other sharpness-aware algorithms: when $\rho$ is too small, the ascent step becomes negligible and ineffective; conversely, if $\rho$ is too large, the ascent step becomes excessively big, leading to unstable training.

• The role of the sample complexity $n$ in Theorem 2: As mentioned above, one can find in the formal statement of Theorem 2 in the Appendix that the error term has the form $O\left(\sqrt{\frac{\log(1+\frac{1}{\rho^2})+4\log\left(\frac{n}{\delta}\right)}{n-1}}\right)$, which tends to $0$ as the sample complexity grows. Acknowledging the importance of the information about the sample complexity, we have incorporated this suggestion in the revision for better clarity.

• Regarding the experiments:

  • Data augmentation: We adopted the data augmentations used by the V-PETL benchmark for model-finetuning. To ensure fairness, on each dataset, the same augmentation was applied across all methods.
  • Runtime and memory limitations: As discussed in the Limitations section, our method shares the same limitation with SVGD and other particle-based approaches that runtime and memory requirements grow linearly with the number of particles due to the sequential implementation. For future research, we suggest addressing this limitation by employing parallel and efficient implementations of SVGD, such as the one proposed in [5]. However, since this issue is not the primary contribution of our work, we opted not to discuss it in great detail in the main text.

We appreciate the reviewer’s useful feedback. The motivation for testing our method on fine-tuning stems from its growing significance in various applications, such as model editing and personalized AI, where tasks often rely on learning from personalized, domain-specific datasets. Additionally, particle-based Bayesian methods, which involve storing multiple model instances, are suitable for fine-tuning methods like LoRA and prompt tuning, since these methods typically freeze most of the model’s parameters and train only a lightweight module, making them computationally efficient and suitable for Bayesian techniques. As for not evaluating the proposed method on training models from scratch, this is discussed in the limitations section. Our method, like other particle-based Bayesian inference approaches, requires storing multiple models, which can lead to memory bottlenecks when dealing with large-scale models.