Improving Generalization with Flat Hilbert Bayesian Inference

We appreciate the reviewer's feedback. We would like to address the concerns as follows:

• Regarding the relationship between FHBI and SAM: As discussed at the end of Section 4, FHBI is a generalization of SAM with multiple model particles. This property is reflected more clearly in the update rules specified in the pseudocode of Algorithm 1 rather than the objective function. In particular, for the case of $m=1$ particles, the kernel terms become constant, and hence, the first step becomes the ascend step in SAM, while the second step in Algorithm 1 becomes the descend step of SAM.

• Regarding the notations: Both terminologies in question refer to population loss. We thank the reviewer for highlighting this inconsistency and will fix this to "population loss" in the final revision to ensure consistency.

• Regarding the runtime comparison with other methods in Figure 5: Firstly, it appears that the reviewer is referring to Figure 5 instead of Figure 3, since Figure 3 does not contain the comparisons with SVGD. Nevertheless, even though we compared FHBI with many baselines in the main experiments, we decided to ablate the runtime of FHBI by comparing it with SVGD since SVGD is most directly related to FHBI. Moreover, in Figure 5, we cannot compare with SAM because the figure presents the runtime based on the number of particles, which is not suitable for deterministic, single-particle methods like SAM.

We thank the reviewer for the constructive feedback and would like to address the concerns as follows:

• Memory consumption and scalability with larger models: All experiments were conducted using a single Tesla V100 GPU with 4 model particles. The memory usage and training time for the CIFAR-100, Caltech-101, and PatchCamelyon datasets are reported in the tables below. In line with the reviewer’s suggestion, we will incorporate this ablation study into the appendix of the final revision.

VRAM consumption (MB):

Training time (s/iter):

• Sensitivity to the RBF kernel length scale: As noted in both the main paper and the appendix, we initially tune the length scale parameter $\sigma$ from the candidate set $\\{0.7, 1, 1.2\\}$. To further investigate sensitivity to this hyperparameter, we expand the range to $\\{0.1, 0.7, 1, 1.2, 2.5\\}$, where we additionally include a small value of $\sigma=0.1$ and a large value of $\sigma=2.5.$ The results on the Natural datasets, reported below, indicate that the performance remains robust within a reasonable range. We also observe a slight degradation when using extremely small values (e.g., $\sigma = 0.1$), where the model tends to overfit, or large values (e.g., $\sigma = 2.5$), where the model tends to underfit.

• Measurement of sharpness on the RKHS: Since the transportation function $f$ governs the updates of the particles, its sharpness in the RKHS governs the sharpness of individual particles. For this reason, we decide to report the sharpness of each particle, as it more directly correlates with the model’s predictive behavior and generalization ability. As demonstrated in the ablation studies, reducing the sharpness in the RKHS leads to a corresponding reduction in the sharpness of every particle, thereby improving the generalization ability of the ensemble.

We appreciate the reviewer’s thoughtful response. As rightly pointed out, we’ve acknowledged in the Limitations section that FHBI, like other particle-based Bayesian inference methods, shares the common drawback of requiring multiple models to be retained during training. This makes it less practical for training from scratch. However, FHBI remains a strong fit for fine-tuning scenarios, where the trainable components are typically lightweight. Looking ahead, an exciting direction for future work is to bring the sharpness concept in RKHS to recent variational inference (VI) methods, which do not suffer from the same memory overhead. This could offer a compelling approach in enhancing performance and robustness, while maintaining memory efficiency.