Improving Generalization with Flat Hilbert Bayesian Inference
Authors: Tuan Truong
Quyen Tran
Ngoc-Quan Pham
Nhat Ho
Dinh Phung
Trung Le
Venue: ICML 2025 poster (2025)
Keywords: Bayesian Inference
Sharpness-aware Minimization
RKHS
Primary area: probabilistic_methods->variational_inference

Abstract: We introduce Flat Hilbert Bayesian Inference (FHBI), an algorithm designed to enhance generalization in Bayesian inference. Our approach involves an iterative two-step procedure with an adversarial functional perturbation step and a functional descent step within the reproducing kernel Hilbert spaces. This methodology is supported by a theoretical analysis that extends previous findings on generalization ability from finite-dimensional Euclidean spaces to infinite-dimensional functional spaces. To evaluate the effectiveness of FHBI, we conduct comprehensive comparisons against nine baseline methods on the VTAB-1K benchmark, which encompasses 19 diverse datasets across various domains with diverse semantics. Empirical results demonstrate that FHBI consistently outperforms the baselines by notable margins, highlighting its practical efficacy.

Improving Generalization with Flat Hilbert Bayesian Inference
Tuan Truong Quyen Tran
Ngoc-Quan Pham
Nhat Ho Dinh Phung Trung Le 3
Abstract
area, including those by Kingma & Welling (2014); Kingma et al. (2015), and Blundell et al. (2015), who extended the Gaussian variational posterior approximation for neural net- works, as well as Gupta & Nagar (2018), who enhanced the flexibility of posterior approximations. In addition to varia- tional methods, various particle sampling techniques have been proposed for Bayesian inference, especially in scenar- ios requiring multiple models. Notable particle sampling methods include Hamiltonian Monte Carlo (HMC) (Neal, 1996), Stochastic Gradient Langevin Dynamics (SGLD) (Welling& & Teh, 2011), Stochastic Gradient HMC (SGHMC) (Chen et al., 2014), and Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016). Each method contributes to deeper understanding and more practical application of Bayesian inference deep learning.
We introduce Flat Hilbert Bayesian Inference (FHBI), an algorithm designed to enhance gener- alization in Bayesian inference. Our approach involves an iterative two-step procedure with an adversarial functional perturbation step and a functional descent step within a reproducing kernel Hilbert space. This methodology is sup- ported by a theoretical analysis that extends pre- vious findings on generalization ability from finie-dimeenioonl Euclidean spaces to infinite- dimensional functional spaces. To evaluate the effectiveness of FHBI, we conduct comprehen- sive comparisons against nine baseline methods the VTAB-1K benchmark, which encompasses 19 diverse datasets across various domains with diverse semantics. Empirical results demonstrate that FHBI consistently outperforms the baselines by notable margins, highlighting its practical effi- cacy.
Besides quantifying uncertainty, tackling overfitting is significant challenge in machine learning. Overfitting often occurs when the training process gets stuck in local minima, leading to a model that fails to generalize well to unseen data. This problem is mainly due to loss functions' high- dimensional and non-convex nature, which often exhibit multiple local minima in the loss landscape. In standard deep network training, flat minimizers improve model gen- eralization (Keskar et al., 2017; Kaddour et al., 2022; Li et al., 2022). Among the flat minimizers, Sharpness-Aware Minimization (SAM) (Foret et al., 2021) has emerged as a practical approach by concurrently minimizing the empirical loss and reducing the sharpness of the loss function. Re- cently, SAM has demonstrated its versatility and effective- ness across a wide range of tasks, including meta-learning (Abbas et al., 2022), vision models (Chen et al., 2021), and language models (Bahri et al., 2022).
1. Introduction
Quantifying and tackling uncertainty in deep learning is one of the most challenging problems, mainly due to the inherent randomness of the real world and the presence of noisy data. Bayesian inference provides a robust framework for understanding complex data, allowing for probabilistic interpretation of deep learning models and reasoning under uncertainty. This approach not only facilitates predictions but also enables the quantification of uncertainty. A primary challenge in this domain is the computation and sampling from intricate distributions, mainly when dealing with deep learning models. One effective strategy to tackle this is- sue is variational inference, which seeks to approximate the true posterior distribution with simpler forms, known as approximate posteriors, while optimizing variational lower bound. Several techniques have been developed in this
Contribution. We combine flat minimizers and particle- based Bayesian methods to introduce a novel Bayesian infer- ence framework with improved generalization ability. To ac- complish this, we first present Theorem 4.2, which strength- ens prior generalization bounds from finitedimensional Euclidean spaces the reproducing kernel Hilbert spaces (RKHS), which are broader functional spaces and are typi- cally infinite-dimensionall Notably, this theorem introduces the notion of functional sharpness that offers an insight to improve the generalization ability of current particle- sampling methods. Subsequently, Theorem 4.3 translates
Equal contribution 'Movian AI, Vietnam "The University of Texas at Austin, USA 'Monash University, Australia. Correspon- dence to: Trung Le <trunglm @ monash.edu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Improving Generalization with Flat Hilbert Bayesian Inference.
against various Bayesian inference baselines across diverse settings. Section 6 offers deeper analysis of FHBI's be- havior to gain further insight into its effectiveness over the baselines.
0 (++1)
0 (4)
2. Related Works
0
0 11++)
Sharpness-aware minimization. Flat minimizers are more robust to the shifts between training and test losses, thereby enhancing the generalization ability of neural net- works (Jiang et al., 2020; Petzka et al., 2021; Dziugaite & Roy, 2017). The relationship between generalization and the width of minima has been studied both theoretically and em- pirically in several prior works (Hochreiter Schmidhuber, 1994; Neyshabur et al., 2017; Dinh et al., 2017; Fort & Gan- guli, 2019). Consequently, a variety of methods have been developed to search for flat minima (Pereyra al., 2017; Chaudhari et al., 2017; Keskar et al., 2017; Izmailov et al., 2018). Among the flat minimizers, Sharpness-Aware Mini- mization (SAM) (Foret et al., 2021) has gained significant attention due to its effectiveness. SAM has been leveraged across a wide range of tasks and domains, including do- main generalization (Cha al., 2021; Wang et al., 2023; Zhang et al., 2023), federated learning (Caldarola al., 2022; Qu et al., 2022), Bayesian networks (Nguyen et al., 2023; Möllenhoff & Khan, 2023), and meta- meta-learning (Abbas et al., 2022). Moreover, SAM has demonstrated its ability to enhance generalization in both vision models (Chen et al., 2021) and language models (Bahri et al., 2022).
Figure 1. Schematic of SAM w. independent particles (green), SVGD (orange), and our FHBI (black) updates. SAM's particles are not aware of others' trajectories. SVGD only seeks the modes and promotes spatial diversity. FHBI seeks the modes, minimizes sharpness, and promotes spatial and angular diversity.
these notions of functional sharpness and generalization from the context of RKHS to the context of Bayesian infer- ence. This analysis establishes a connection between the population and empirical KL loss, providing a strategy to enhance generalization by minimizing the population KI. loss. Motivated by these two theorems, we present Flat Hilbert Bayesian Inference (FBVI), practical algorithm that employs a dual-step functional Charpness-awar aware update procedure in RKHS. This approach improves the general- ization of sampled particles, thereby enhancing the quality of the ensemble. Overall, our contributions are as follows:
Nevertheless, these studies are constrained to finite- dimensional Euclidean spaces. In this work, strengthen these generalization principles to infinite-dimensional func- tional spaces and propose particle-saplingg method grounded in this theoretical framework.
1. We present a theoretical analysis that characterizes generalization ability over the functional space. This analysis generalizes prior works from the Euclidean space to introducing the notion of functional sharpness i.e., the sharpness of the functional spaces.
functional space, thereby
Bayesian Inference. Two main strategies were widely employed in the literature of Bayesian inference. The first paradigm is Variational Inference, which aims to approxi- mate a target distribution by selecting a distribution from a family of potential approximations and optimizing vari- ational lower bound. Graves (2011) introduced the use of Gaussian variational posterior approximation for neural network weights, which was later extended by Kingma & Welling (2014); Kingma et al. (2015); Blundell et al. (2015) with the reparameterization trick to facilitate training deep latent variable models. Louizos & Welling (2017) proposed using matrix-variate Gaussian model entire weight ma- trices (Gupta & Nagar, 2018) to increase further the flex- ibility of posterior approximations, which offers novel approach to approximate the posterior. Subsequently, vari- ous alternative structured forms of the variational Gaussian posterior were proposed, including the Kronecker-factton approximations (Zhang et al., 201 8; Ritter al., 2018; Rossi
2. Building on this theoretical foundation, we propose practical purticle-samping algorithm that enhances the generalization ability over existing methods. We conducted extensive experiments comparing our Flat Hilbert Bayesian Inference (FHBI) algorithm with nine baselines on the VTAB-1K benchmark, which includes datasets across various domains and semantics. Ex- perimental results demonstrated that our algorithm out- performs these baselines by notable margins.
Organization. The paper structured as follows: Sec- tion 2 reviews the related works on Bayesian inference and the development of flat minimizers. Section 3 provides the necessary background and notations. Section 4 dis- cusses the motivation and theoretical development behind our sharpness aware murticle-sampling approach. Section presents experimental results, comparing our algorithm
2

---

Improving Generalization with Flat Hilbert Bayesian Inference.
al., 2020), or non-centered or rank-1 parameterizations (Ghosh al., 2018; Dusenberry et al., 2020). Recently, Pham et al. (2025) proposes to incorporate sharpness and dis- tributional robustness to Bayesian Inference, thereby mod- elling the interactions between particles and allowing a more diverse set of model particles.
More formally, the empirical posterior is equal to:
n Σ n
p(0|S) = exp
((o((i);) p(0) (2)
The second paradigm of Bayesian inference is Markov Chain Monte Carlo (MCMC), which involves sampling multiple models from the posterior distribution. MCMC has been applied to neural network inference, such as Hamil- tonian Monte Carlo (HMC) (Neal, 1996). However, HMC requires the computation of full gradients, which can be computationally expensive. To address this, Stochastic Gra- dient Langevin Dynamics (SGLD) (Welling & Teh, 2011) integrates first-order Langevin dynamics within a stochastic gradient framework. Stochastic Gradient HMC (SGHMC) (Chen et al., 2014) further incorporates stochastic gradients into Bayesian inference, enabling scalability and efficient exploration of different solutions. Another critical approach, Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016), closely related our work, uses a set of particles that converge the target distribution. It also theoretically established that SGHMC, SGLD, and SVGD asymptotically sample from the posterior as the step sizes approach zero.
where Zs the normalizing constant. We define the popu- lation and empirical losses as follows:
(D( =
"
Ls(0) = E, -(..))ss(((0(()))
(ff((:)) u.).
The population loss is defined as the expected loss over the entire data-label distribution. In contrast, the empirical loss is the average loss computed over a given training set S. Based on these definitions, the empirical posterior in Eq. 2 can be written as:
p(0 S) - (xs]=cccc<((()p000))
Intuitively, models with parameters O that fit well to the training set S lead to lower empirical loss values, resulting in higher density in the empirical posterior. This formu- lation has appeared in prior works on Bayesian Inference, such as (Nguyen et al., 2023). However, simply fitting to the training samples can lead to overfitting. To improve gen- eralization, we are more concerned with performance over the entire data distribution 'D rather than just the specific sample S. Accordingly, we define the population posterior as IP whose density is given by:
3. Backgrounds and Notations
Bayesian Inference. Consider a family of neural networks fo(z), where the random variable 0 represents the model parameters and takes values in the model space θ Rdd We are given a training set S = {[ y:) of n i.i.d observations from the data space X' x V. and the prior distribution of the parameters p(0). In the literature on Bayesian inference problems, prior works typically focus on approximating the empirical posterior Pois. whose density function p( S) is defined as:
p = cxp(-Co(0))p(00/1
(3)
with the normalizing constant Zp. This population posterior is more general than the empirical posterior, as it captures the true posterior of the parameters under the full data dis- tribution. However, understanding the population posterior is particularly challenging because we can only access the empirical loss (O), not the population loss (D))). In this paper, we deviate from prior approaches that primarily focus on approximating the empirical posterior and instead propose particle-sampling method to approximate the pop- ulation posterior.
P(O[S) & p() II plu,
where the prior distribution Po has the density function p(0). The likelihood term is proportional to
plyl,, S, 0) x exp (-)c(66or.))
Reproducing Kernel Hilbert Space (RKHS). Let O 0') θ x θ + I. be a positive definite kernel op- erating on the model space. The reproducing kernel Hilbert space (RKHS) H of k(0, 0') is the closure of the linear span {f f()) = Σ, a,k(,, 0,), a, E R, O E }). For f() = Σ, a,k 0.) and g = Σ, byk (e000,)). H is equipped with the inner product defined by (f. 9. = Enj k,,, 0,). For all O E O, there exists a unique
with some loss function f, such as the CrossEntropy loss, and a sufficiently expressive model fg. Then, the empirical posterior is:
(-1
1 Σ {(f(xx))) p(0). (1) n
P(0[S) x exp
3

---

Improving Generalization with Flat Hilbert Bayesian Inference.
element Kg_ E H with the reproducing property that f(0) (J. Ko) n. for any f E H.
Proposition 4.1. Consider the problem of finding the distri- bution Q that solves:
Given that H is a scalar-valuee
RKHS with kernel k(0, 8), H. = H x H. x ... x H. is a vector-valued RKHS of functions f = f2. ... ff] corresponding to the ker- nel K (0,, 0¹) = k(0,000000 1. Hd is equipped with the inner
Q' = min Q
BenalCD(OO] + DKL(Q|IP)))
(5)
where we search over Q absolutely contimuous w.n.t Fg, and the second term is the regularization term. The closed-form solution this problem is exactly the population posterior defined in Eq. 3.
product (f.g) H4a = Ex=n 9.)
Let F[f] be a functional on f E Hd. Similar to the defi- nition by (Liu Wang, 2016), the (functional) gradient of F is defined
function V F[f Hd such that for any
g Hd and E IR
In this proposition, aim to identify the posterior distribu- tion that minimizes the expected population loss, where the expectation is taken over the entire parameter space with i ~ Q while maintaining proximity to the prior distribu- tion to ensure simplicity. With access to this posterior Q" we can sample a set of particles whose average performance optimally minimizes the population loss. Since the solution to this optimization problem corresponds precisely to the population posterior, the ensemble of the particles sampled from Q = p(O[D) effectively minimizes the average value of the population loss. This is because Q' is explicitly cho- sen to minimize the expected value of the population loss Cp. which means the ensemble fits the whole data distribu- tion instead of overfitting to the specific dataset S, therefore establishes improved generalizability. Consequently, this proposition theoretically asserts that sampling from p(0(D) enhances the generalizability of the ensemble.
F + g (V ( (4)
+
Stein Variational Gradient Descent (SVGD). Given a general target distribution O)), SVGD (Liu & Wang. 2016) aims to find flow of distributions (q & that minimizes the KI. distance to the target distribution. Motivated by the Stein identity and Kernelized Stein Discrepancy, SVGD
proposes the update (k+11) " in which T O is a smooth one-to-one push- forwaforward map of the form T(0)-
+ eppp. (0) which:
(.) Ea~ (A,
(0) = ( vogg + Veσ(O).
Ap is known as the Stein operator, which acts on á and produces a zero-mean function (,) when 0 ~ p. While SVGD is designed for general target distributions p. in the context of Bayesian inference, it only applicable to the em- pirical posterior rather than the population posterior, which will discuss in detail in the next section.
4.1. Theoretical analysis
Motivated by this observation, we advance prior work by approximating the population posterior. Specifically, to improve generalizability, our objective is to approximate the target population posterior distribution p(0[D) using a simpler distribution q" (0)) drawn from a predefined set of distributions F. This is achieved by minimizing the KI.
4. Flat Hilbert Bayesian Inference (FHBI)
divergence:
Consider the Bayesian inference problem of approximating a posterior distribution. prior works, such as SVGD (Liu & Wang. 2016), when applying to the context of Bayesian inference, the methods are only applicable the empiri- cal posterior p(0 p(O)S) because we only have access to the empirical loss. It is evident that when sampling a set of m particle models 01.m. from p(0) S)) these particles congre- gate in the high-density regions of the empirical posterior (())S) corresponding to the areas with low empirical loss (O). However, to avoid overfitting, it is preferable to sample the particle models 11:m. from the population poste- rior p(0 D) x cxp(-Lx (0))p(()) as this approach directs the particle models 0p.m towards regions with low values of the population loss (0), thus improving generalization ability. To better understand this motivation from a theoreti- cal perspective, consider the following proposition, with the proof provided in Appendix A.I:
9º arg min DKL ((0))1((1))) qEFF
(6)
Ideally, the set F should be simple enough for a simple solution and effeectvee inference while sufficiently broad to
approximate a wide range of target distributions closely. Let (0) be the density of a reference distribution. We define F as the set of distributions for random variables of the form & T()), where T: : θ θ is smooth, bijective mapping, and 0 is sampled from g. By variable change, the density of 0, denoted q (-), is expressed follows:
Lizy (1) = g ( de,
We restrict the set of the smooth transformations T to the set of push-forward maps of the form T(0) = 0 f(),, where

---

Improving Generalization with Flat Hilbert Bayesian Inference.
I Hd. When x is sufficiently small, the Jacobian of T I+ f is full-rank where I denotes the identity map, in which case T's is guaranteed to be one-to-one map according to the inverse function theorem. Under this restric- tion, the problem is equivalent to solving an optimization problem over the RKHS:
Our objective is to learn the function f" Hd that minimizes (xu[41=)]]l((00))) Theorem 4.3, we propose to implicitly minimize DKL ((ts+1)ll)(() by minimizing the right-hand side term max p' (T-fllaass DKL (91+s)11p((1))) For any
Motivated by
f E H., let Fif) DEL (((+5))(((0))) and f
arg min SeRelllllyee
DEL (1=7)())r((000))
f + pf. it follows that:
arg max DKL (11ts)lle((1s))
The challenge with this optimization problem lies in our lack of access to the population loss function (o) and the population posterior distribution p(0(D). We present our first theorem to address this issue, which characterizes generalization ability in the functional space Hd. The proof of this theorem can be found in Appendix A.2.
11'- SlaSe
arg max Flf + Pfl x
(8)
arg max F(f) + (f. +O(
(9)
Theorem 4.2 (Informal). Let ê H. x + R²+ be loss function the RKHS H. and the data space. Define Lp(f) = B(ayy)=(((( (((f x, v) and Ls(f) = - [{ 3i, be the corresponding population and empirical losses. Then for any p > 0 and any distribution D, with probability of 1 8 over the choice of the training set S ~ DD, we have:
arg max (f, (f.V) 1f
(10)
Let g - F[J] E Hd. The Cauchy-Schre Schwarz inequality on Hilbert spaces (Kreyszig, 1978) implies:
(f.9) 1e 3 (f.9) ne S {1 S |191|
Lo(f) $
Ls(f')
max
P'EH11If" f1l Sp
In turn, the solution f that solves the maximization prob- lem in Eq. 10 is given by:
log(1 } + log (3)
+0
Vy DEL ( (11+))1l(((1)))
f =
9
(11)
1
V DKL (qr4n)ll((s)))
This theorem extends prior results, such as the generaliza- tion bounds established by Foret et al. (2021) and Kim et al. (2022), from Euclidean space to a broader, more general reproducing kernel Hilbert space. It is noteworthy that this is not a straightforward extension, as the previous general- ization bounds rely on the dimensionality of the domain, while the RKHS is typically dnfiniteddimensiona for many widely used kernels such as the RBF kernels (Aronszajn, 1950). Building on the first theorem, we present the second theorem, which directly addresses the population posterior and serves the primary motivation for our method. The proof of this theorem can be found in Appendix A.3.
H
Recall that our goal is to find sequence of functions (fickk C Hd that converges toward the optimal solution f With the sequence {fa],, we can obtain the flow of distributions {q(*) k.. in which g -II+ful that gradu- ally approaches the optimal solution of Eq. 6. Motivated by Eq. 11, we propose the following functional sharpness- aware update procedure:
V y DKL. ((L+))ll((s)))
Theorem 4.3 (Informal). Let q be any distribution and dyc denotes the VC dimension of he hypothesis space F = (so E θ}. For any p > O, with probability of 1 - 8 over the training set S generated by distribution D, we have:
ft
g-
(12)
V DKL q (T=s))l(((5))
s=52 Hd
fati fn - v DKT. ((I+)1lP(11S)) f f
(13)
DKL (1L+)llo(01)))
(k+1) =qi++f+t)
(14)
S max
DKL. (1I+):111(005))
To implement this iterative procedure, we must work with the functional gradient terms. For this, we rely on the fol- lowing lemma, with the proof provided in Appendix B of Liu Wang (2016):
Vdve log 2en dvc
log(1 l ( + + log (#)
+ +0
5v2n
5

---

Improving Generalization with Flat Hilbert Bayesian Inference.
Lemma 4.4. Let F(f) = DaL Dka(vr-r)ll(((S) When ||
Algorithm 1 FLAT HILBERT BAYESIAN INFERENCE (FHBI)
is sufficiently small,
Input: Initial particles (0!")) ym= number of epochs N,
V F[
(15)
step size p > 0
- ((0)S)(00)) Vok(O, -)].
Output: A set of particles {0.}m² that approximates the population posterior distribution p() for iteration do (k)
(16)
Denote the right-hand side of the equation above D(f). Substituting into Eq. (12) and (13), the iterative procedure described from equations (12)-(13) becomes:
where x(0)
- 1M (0 O
(b)
V of logp(0(*)(s)
"
V k0,0
((0) (m))
(f+
ft = | Rd
(17)
0 (1++1)
( )(())
fk+1 f&- - -e f2),
(18)
where
$0 e)
(k+)) -qu+te++1*
(19)
" b [1 0
oge log ((g)(
+ E/5) +
V of (0,
Even though we do not have access to p(0)S). we can compute Vo log p(OS)) because Vo log (0)S) = Vo log (O). To implement the procedure above, we first draw a set of m particles (0 M== on the model space from the initial density, and then iteratively update the particles with an empirical version of Dff) Consequently, we obtain the practical procedure summa- rized in Algorithm I, which deterministically transports the set of particles to match the empirical posterior distribution p(0]S), therefore match the general posterior p(f(D) as supported by Theorem 4.3, In Algorithm I, at each iteration k, we have m particles (0) (0(4)))7 y Eq. 17 computes the m ascend steps e(<);; then, Eq. 18 and Eq. 19 use these ascend steps to transport the m model particles to (0) J==1- It is noteworthy that FHBI is generalization of both SVGD and SAM. In particular, we set p = 0, we get SVGD; when m #PARTICLES = 1, we obtain SAM.
end for
spatial repulsive force, driving the particles apart to prevent them from collapsing into a single mode. Consequently, FHBI is not a straightforward extension of SAM to multiple independent particles; it enables the sharpness and gradi- ent directions of the particles interact with one another. This insight on our algorithm is summarized in Figure 1. In Section 6, we empirically demonstrate that, compared to SVGD, FHBI not only effectively minimizes particle-wise sharpness and loss values but also fosters greater diversity in the travel directions of the particles during training. This increased directional diversity, combined with the kernel gradient term, further mitigates the risk of particles collaps- ing into single mode and improve the final performance as presented in Section 5.
Interactive gradient directions and Connections to SAM. To gain further insight into the mechanism of FHBI and its connections to SAM, consider the term Vo, log @ + & in the descending step, which is related to Ls + E.).
5. Experiments
Applications to Model Fine-tuning, Bayesian inference methods have promising applications in model fine-tuning. We are given a pre-trained model Φ in standard fine-tuning scenarios. The objective is to find the optimal parameters 0 = + + 3, where / represents an additional module, often lightweight and small relative to the full model. Several parameter-efficient finetuning strategies have been devel- oped, including LoRA (Hu et al., 2022), Adapter (Houlsby et al., 2019). and others. Our experiments focus on fine- tuning the ViT-B/16 architecture (Dosovitskiy, 2021), pre- trained with the ImageNet-21r 22 dataset (Deng et al., 2009), where /3 is defined by the LoRA framework. For the Bayesian approaches, we aim to learn m LoRA particles 300 to obtain m model instances (()) The final output is then computed as the average of the outputs from all these model instances.
The perturbed loss can be approximated as:
Os + &4) R (0,) + & Cs
=== where & involves the average ΣxL k(Ox, 0,) Ls (Ox).
Consequently, the gradient of this perturbed loss indicates a direction that simultaneously minimizes || (s,, || which approximates the sharpness of the j-th particle, as discussed by Foret et al. (2021) and Co,Cs Vo, (0x) for all j, k. which reflects the angular simi- larity in the directions of the two particles. Thus, in addition to minimizing the sharpness of each particle, the first term of the descent step acts as an angular repulsive force, promot- ing diverse traveling directions for the particles. Besides, as discussed by Liu & Wang (2016), the second term acts as a
6

---

Improving Generalization with Flat Hilbert Bayesian Inference.
Table VTAB-1K results evaluated on Top-1 accuracy. All methads are applicd finetune the of LoRA parameters on VYT-B/16 pre-trained with InageNet-22 -21K dataset.
Natural
Specialized
Structured
Retinoparthy
Clev- Count
SNORB- Adl
NNORB-ED Elo
CIFARIOD
Calteh1101
DT1D
Flower
SWHN
Sun:397
Camelpo
EuroSAA
Resise45
Clevn-Dot
DMLab
KITTI
dSpe Lac
dSpe-On
Method
AVG
FFT
68.9
87.7
64.3
97.2
86.9
87.4
38.8
79.7
95.7
84.2
73.9
56.3
58.6
41.7
65.5
57.5
46.7
25.7
29.1
65.6
AdamW
67.1
90.7
84.5
54.2
84.1
94.9
84.4
82.9
49.8
78.5
75.7
47.1
SAM
72.7
90.3
71.4
84.4
52.4
82.0
92.6
84.1
74.0
76.7
68.3
47.9
74.3
71.6
43.4
26.9
39.1
70.5
Deeplins
69.1
88.9
67.7
98.9
90.7
85.1
54.5
82.6
94.8
82.7
75.3
46.6
47.1
47.4
68.2
71.1
36.6
30.1
35.6
67.0
BayesTune
67.2
91.7
69.5
99.0
90.7
86.4
54.7
84.9
95.3
84.1
75.1
12.88
68.9
49.7
79.3
74.3
46.6
30.3
42.8
72.2
SGLD
68.7
91.0
67.0
98.6
89.3
83.0
51.6
81.2
93.7
83.2
76.4
80.0
70.1
48.2
76.2
71.1
39.3
31.2
38.4
70.4
SADA-JEM
91.9
70.2
54.7
94.1
77.0
79.9
72.1
40.1
SA-BNN
65.1
91.5
71,0
89.4
89.3
55.2
83,2
94.5
86.4
75.2
61,4
63.2
40.0
71.3
64.5
34.5
27.2
31,2
SVGD
71.3
90.2
71,0
98.7
90.2
84.3
52.7
83,4
93.2
86.7
75.1
75,8
70,7
49.6
79.9
69.1
41.2
30,6
33.1
70.9
FHBI
74.1
93.0
74.3
99
92 4
87.3
56.5
85.3
95.0
87.2
79.6
80.1
72.3
52.2
80
70.8
51.2
31.9
41.3
73.7
(17)
(42)
(.15)
(0.20)
(0.21)
(.52)
(12)
(.31)
(.57)
(221)
(.20)
(.16)
(.27)
(.47)
(.31)
(.50)
(.32)
(.36)
(.59)
87
60
Score (96)
Score (%)
55
Score
80
85
50
78
83
45
SGLD
DeepEns
BayesTune
SVGD
FHBI
SGLD
DeepEns
BayesTune
SVGD
FHBI
SGLD
DeepEns
Buyes
SVGD
FHBI
Figure 2. Domain-wise -iis average scores on Natural (left), Specialized (middle), and Structured (right) datasets. FHBI performs best in all three domains compared to the Bayesian inference baselines.
Experimental Details. To assess the efficacy of FHBI, we experiment VTAB-1K (Zhai et al., 2020), challenging image classification//rcdiitiooo benchmark consisting of 19 datasets from various domains. VTAB-1K covers various tasks across different semantics and object categories. The datasets are organized into Natural, Specialized, and Struc- tured domains. We compared FHBI against nine baselines with three deterministic finetuning strategies, including full finetuning, A, and SAM, and four Bayesian inference techniques, including Bayesian Deep Ensembles (Lakshmi- narayanan et al., 2017), Bayes'Tune (Kim & Hospedales, 2023), Sharpness-Awart Bayesian Neural Network (SA- BNN) (Nguyen et al., 2023), Sharpness-aware Joint Energy- based Model (SADA-JEM) (Yang et al., 2023), Stochas- tic Gradient Langevin Dynamies (SGLD) (Welling & Teh, 201 and Stein Variational Gradient Descent (SVGD) (Liu & Wang, 2016). The experiments were conducted on sin- gle Tesla V100 GPU. Appendix C provides more details regarding the chosen hyperparameters.
baselines, outperforming them in most settings. Compared other particle sampling methods, including SGLD and SVGD, FHBI consistently performs better across all set- tings. Moreover, FHBI improves upon SAM by margin of 3.2%, highlighting the advantages of using multiple parti- cles with the underlying interactive gradient directions as previously discussed in Section 4. Additionally, as illus- trated in Figure 2, FHBI shows the highest performance across all three domains, further solidifying its advantage over the Bayesian inference baselines. Besides, FHBI also outperforms SA-BNN, which resembles the combination of SVGD and SAM, as well as SADA-JEM, which resembles SGLD and SAM, highlighting that FHBI is substantially dif- ferent from a combination SVGD and SAM as discussed in Section 4. and in the ablation study in Section 6.
To further assess the robustness of FHBI, evaluate the Ex- pected Calibration Error (ECE) of each setting. This score measures the maximum discrepancy between the model's accuracy and confidence. As indicated in Table 2, even though there is typically a trade-off between accuracy and ECE, our approach achieves a good balance between the ECE and the classification accuracy.
Experimental Results. We present the classification ac- curacy results in Table FHBI notably improves upon the

---

Improving Generalization with Flat Hilbert Bayesian Inference.
Table VTAB-1K results evaluated the Expected Calibration Error (ECF) metric. All methods are applied finetune the of
LoRA parameters on VIT-B/16 pre- vreirrined with ImageNet- 21K dataset.
Natural
Specialized
Structured
CIFAR100
Caltcc111
DTD
Flower 102
Pets
SVHN
Sun399
Camelyon
EuroSAT
Resise45
Retinopally
Clevr- Count
Clevs- Dist
DMLab
KITTI
dSpr-Loc
dSpe-O6i
3N23NO AVORB-AD
SNORB-ER
Method
AVG
FFT
0.29
0.23
0.20
0.13
0.27
0.19
0.45
0.21
0.13
0.18
0.17
0.41
0.44
0.42
0.22
0.14
0.23
0.24
0.40
0.26
AdamW
0.38
0.19
0.18
0.05
0.09
0.10
0.14
0.11
0.09
0.12
0.11
0.12
0.19
0.34
0.18
0.14
0.21
0.18
0.31
0.17
SAM
0.21
0.25
0.20
0.11
0.12
0.15
0.14
0.17
0.16
0.14
0.09
0.12
0.17
0.24
0.16
0.21
0.19
0.13
0.16
0.16
DeepEns
0.24
0.12
0.22
0.04
0,10
0.13
0.23
0.16
0.07
0.15
0.21
0.31
0.32
0.36
0.13
0.32
0.31
0.16
0.29
0.20
BayesTune
0.32
0.08
0.20
0.03
0.85
0.12
0.22
0.13
0.07
0.13
0.22
0.12
0.23
0.30
0.24
0.28
0.28
0.31
0.26
0.23
SGLD
0.26
0.20
0.17
0.05
0.18
0.14
0.23
0.18
0.09
0.12
0.32
0.26
0.29
0.21
0.26
0.42
0.39
0.11
0.24
0.22
SADA-JEM
0.22
0.11
0.20
0.05
0.13
0.16
0.18
0.15
0.21
0.23
0.26
0.19
0.20
0.25
0.27
0.35
0.20
0.14
0.13
0.19
SA-BNN
0.22
0.08
0.19
0.15
0.12
0.12
0.24
0.13
0.06
0.12
0.18
0.14
0.21
0.22
0.24
0.25
0.41
0.46
0.34
0.20
SVGD
0,20
0.13
0.19
0.04
0.16
0.09
0.20
0.15
0.11
0.13
0.12
0.17
0.21
0.30
0.18
0.21
0.25
0.14
0.26
0.18
FHBI
0.19
0.10
0.16
0.06
0.06
0.09
0.16
0.09
0.05
0.12
0.08
0.14
0.15
0.21
0.15
0.16
0.18
0.11
0.07
0.12
Gradients ancular similaribes
3.75
FHBI
PH1 SVGD
3.50
SVGD
0.35
Sturseess
3,25
0.20
V wmmm
0.22
3.00
mmm
60
200
100
2.50
Epach
Epoch
2.25
(a) Particle
(b) Particle 2
EDoch
FHRRI
Shurpnccs 5/4
0.30
Shapnus
Figure 3. Gradients angular similarities with m = 4. Lower values indicate greater angular diversity.
0.20
A potential concern regarding our algorithm is the additional computational overhead compared to SVGD due to the gra- dient ascent step. Nevertheless, Appendix B.I, we have conducted an experiment to assess the immpact of differ- number of particles. Figure 5 and shows that with negotiable tradeoff in runtime, multiple particles result in significant performance improvements compared to a single particle, and we found that #Particles 4 provides an optimal balance between performance gains and compu- tational overhead for this particular application. Another potential concern is the effect of the kernel choice k. Indeed, have chosen the RBF kernel due its effectiveness in the literature of kernel methods. As have shown in Ap- pendix B.2, the performance gap between different kernel choices is insignificant, showing the robustness of FHBI with respect to the kernel choice.
Epoch
Epoch
(c) Particle
(d) Particle 4
Figure 4. Evolution of sharpness of four particles over 100 epochs with SVGD (blue) and FHBI (red)
improving the final performance. To empirically verify this hypothesis on the behavior of our algorithm, con- trast FHBI with SVGD the KITTI dataset. Four par- ticles are initialized at the same location. We measured: 1) the evolution of sharpness of each particle, defined as max jellSp (s + e) - Ls(0) according to Foret et al. (2021), and 2) the evolution gradients angular diversity, quantified as the Frobenius norm of the covariance matrix formed by the particle gradients. Figure shows that FHBI results in significantly lower and more stable sharpness evo- lution, encouraging less congruent gradient directions and promoting particles to explore diverse trajectories. Hence, FHBI reduces particle sharpness while promoting angu- lar diversity, thereby improving generalization ability and avoiding overfitting by collapsing into single mode.
6. Ablation Studies: Particles Sharpness and Gradient Diversity
As discussed in Section 4.1 and Section 5, FHBI shares implicit connections with SAM by minimizing particle- wise sharpness and diversifying travel directions, thereby
8

---

Improving Generalization with Flat Hilbert Bayesian Inference.
7. Conclusion
Caldarola, D., Caputo, B., and Ciccone, M. Improving generalization in federated learning by seeking flat min- ima. In European Conference Computer Vision, pp. 654-672. Springer, 2022.
We introduce Flat Hilbert Bayesian Inference (FHBI) to enhance generalization ability beyond previous Bayesian inference approaches. This algorithm based on the- oretical framework that extends generalization principles from Euclidean spaces to the infiniteddimensioni RKHS. Empirically, FHBI consistently demonstrated significant performance improvements over nine baseline methods.
Cha, J., Chun, S., Lee, K., Cho. H.-C., Park, S., Lee, Y., and Park, S. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405-2244 2021.
Limitations and Future Directions. Similar to other particle-based methods, FHBI needs to store multiple mod- els. Although it remains well-suited for fine-tuning since the additional modules typically lightweight, this require- ment is memory bottleneck for larger models. Given that the variational inference approaches can alleviate this issue, an avenue for future research is to extend the concept of sharpness over functional spaces introduced by our theory to the VI techniques to improve the generalization of these methods without storing multiple models.
Chaudhari, P., Choromańska, A., Soatto, S., LeCun, Y., Baldassi, C.,, Borgs, C., Chayes, J. T., Sagun, I.., and Zecchina, R. Entropy-sgd: biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019, 2017.
Chen, T., Fox, E.,, and Guestrin, C. Stochastic gradient Hamiltonian Monte Carlo. In International conference on machine learning, pp. 1683-1691. PMLR, 2014.
Chen, X., Hsieh, C.-J., and Gong, B. When vision transform- ers outperform resnets without pre-training or strong data augmentations. In International Conference Learning Representations (ICLR), 2021.
Acknowledgments
Trung Le was partly supported by ARC DP23 grant DP23010117 176 and by the Air Force Office of Scientific Research under award number FA9550-23-S-0001.
Deng, J., Dong, W., Socher, R., Li, L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vi- sion and Pattern Recognition, pp. 248-255, 2009. doi: 10.1
Impact Statement
This paper presents work to advance the field of Machine Learning. There are potential societal consequences of our work, none of which we feel must be highlighted.
Dinh, L.,, Pascanu, R., Bengio, S., and Bengio, Y. Sharp min- ima can generalize for deep nets. In International Con- ference on Machine Learning, pp. 1019-1028. PMLR, 2017.
References
Dosovitskiy, A. An image is worth 16x16 words: Trans- formers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
Abbas, M., Xiao, Q. Chen, L., Chen, P.-., and Chen, T. Sharp-maml: Sharpness-awen aware model-agnostic meta learning. arXiv preprint arXiv:2206.0339 2022.
Aronszajn, N. Theory of reproducing kernels. Transactions of the American Mathematical Society, 68(3):337-404, 1950.
Dusenberry, M., Jerfel, G., Wen, Y., Ma, Y., Snoek, J., Heller, K., Lakshminarayanan, B., and Tran, D. Efficient and scalable Bayesian neural nets with rank-1 factors. In International conference on machine learning, pp. 2782- 2792. PMLR, 2020.
Bahri, D., Mobahi, H., and Tay, Y. Sharpness-aware mini- mization improves language model generalization. In Pro- ceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers),
Dziugaite, G. K. and Roy, D. M. Computing nonvacuous generalization bounds for deep (stochastic) neural net- works with many more parameters than training data. In UAI. AUAI Press, 2017.
pp. 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 10.18633/v1/002.. acl-long.508. URL https::t/// http://.....ttttt/aaaaaaaaaaaaa. http::::///aaaatoooogo..ooo 2022. ac1- 508.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-aware minimization for efficiently improving generalization. In 9th International Conference on Learn- ing Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: / //openrevis ew net forum? d= 30-6mmmmmmssssmmmmmmm
Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622, PMLR, 2015.

---

Improving Generalization with Flat Hilbert Bayesian Inference.
Fort, S. and Ganguli, S. Emergent properties of the lo- cal geometry of neural loss landscapes. arXiv preprint arXiv:1910.05929, 2019.
Kim, M. and Hospedales, T. Bayestune: Bayesian sparse deep model fine-tuning. In Advances in Neu- ral Information Processing Systems 36 (NeurIPS 2023), volume 36, pp. 55317-65366 Curran Associates Inc, December 2023. URL ht.pps: /tees///eerriss cc/ Conferences/ / 223.. Thirty-Seventh Conference on Neural Information Processing Systems, NeurIPS 2023 Conference date: 10-12-2023 Through 16- 12-2023.
Ghosh, S., Yao, J., and Doshi- Velez, Structured varia- tional learning of Bayesian neural networks with horse- shoe priors. In International Conference on Machine Learning, pp.
PMLR, 2018.
Graves, A. Practical variational inference for neural net- works. Advances in Neural Information Processing Sys- tems, 24, 2011.
Kim, M., Li, D., Hu, S., and Hospedales, Fisher sam: Information geometry and sharpness aware minimisa- tion. In Chaudhuri, K., Jegelka, S., Song, L., Szepes- vari, C., Niu, G., and Sabato. S. (eds.), Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, pp- 11148- 1 11161. PMLR, July 2022. URL ht / /icml.m.../// cc/ Con ferenceseeeeeee 2.22. 39th International Conference on Machine Learning, ICML 2022 Conference date: 17-07-2022 Through 23-07-2022.
Gupta, A. K. and Nagar, D. K. Matrix variate distributions. Chapman and Hall/CRC, 2018.
Hochreiter, S. and Schmidhuber, J. Simplifying neural nets by discovering flat minima. In NIPS, pp. 529-536. MIT Press. 1994.
Kingma, P. and Welling, M. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Houlsby, N., Giurgiu, A., Jastrzebski, S., Morrone, B., De Laroussilhe, Q.. Gesmundo, A., Attariyan, M., and Gelly, S. Parameter-efffiiient transfer learning for NLP. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 2790-2799. PMIR, 09-15 Jun 2019. URL https: http://oooednnnngg/e mlr, press/v97/ houlsbyl 9a. html.
Kingma, D. P., Salimans, T., and Welling, M. Variational dropout and the local reparameterization trick. Advances in Neural Information Processing Systems, 28, 2015.
Kreyszig, E. Introductory Functional Analysis with Appli- cations. John Wiley & Sons, 1978.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. In ICLR. OpenReview.net, 2022. URL http:///blp uni-trier de/dd / htw.:22222.2mmmm..mmmmmmmmmmm1 HWw.AAAAWAAAAAAA22222
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles, 2017. URL
aba/
1612.01477
Li, Z., Wang, Z., and Li, J. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability, 2022.
Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D. P., and Wilson, A. G. Averaging weights leads to wider optima and better generalization. In UAI, pp. 876 885. AUAI Press, 2018.
Liu, Q. and Wang, D. Stein variational gradient descent: A general purpose Bayesian inference algorithm. Advances in Neural Information Processing Systems, 29, 2016.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. ICLR. OpenReview.net, 2020.
Louizos. C. and Welling, M. Multiplicative normalizing flows for variational Bayesian neural networks. In In- ternational Conference on Machine Learning, pp. 2218- 2227. PMLR, 2017.
Kaddour, J., Liu, L., Silva, R., and Kusner, M. J. Questions for flat-minima optimization of modern neural networks. CoRR, htbp//222.00661... 2022. URL https::/arxiv. org/abs/ 2202.22222 00661.
Möllenhoff, T. and Khan, M. E. SAM as an optimal re- laxation of bayes. In The Eleventh International Confer- ence on Learning Representations, 2023, URL https: / //openreveew.oo net for
Keskar, N. S., Nocedal, J., Tang, P. T. P., Mudigere, D., and Smelyanskiy, M. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, 2017.
Neal, R. M. Bayesian Learning for Neural Networks. Springer- Ppringer-Verlag, Berlin, Heidelberg, 1996. ISBN 0387947248.
10

---

Improving Generalization with Flat Hilbert Bayesian Inference.
Neyshabur, B., Bhojanapalli, S., McAllester, D., and Srebro, N. Exploring generalization in deep learning. Advances in Neural Information Processing Systems, 30, 2017.
Machine Learning, ICML'11, pp. 681-688, Madison, WI, USA, 2011. Omnipress. ISBN 97814 9781450306195.
Yang, X., Su, Q.. and Ji, S. Towards bridging the perfor- mance gaps of joint energy- models. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15732-15741, June 2023.
Nguyen, V.-A,, Vuong, T.-L., Phan, H., Do. T.-I., Phung, D., and Le, T. Flat seeking Bayesian neural networks. In Advances in Neural Information Processing Systems, 2023.
Pereyra, G., Tucker, G., Chorowski, J., Kaiser, L.,, and Hinton, G. E. Regularizing neural networks by penaliz- ing confident output distributions. In ICLR (Workshop). OpenReview.net, 2017.
Zhai, X., Puigcerver, J., Kolesnikov, A., Ruyssen, P., Riquelme, C., Lucic, M., Djolonga, Pinto, A. S.. Neumann, M., Dosovitskiy, A., Beyer, L., Bachem, O., Tschannen, M., Michalski, M., Bousquet, O., Gelly, S., and Houlsby, N. A large-scale study of representation learning with the visual task adaptation benchmark, 2020. URL https: arxiv. orq/ abs aa//aba/1999 /1900. 04867.
Petzka, H., Kamp, M., Adilova, L., Sminchisescu, C., and Boley, M. Relative flatness and generalization. In NeurlPS, pp. 18420-1844 2021.
Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. Noisy natural gradient as variational inference. In Interna- tional Conference on Machine Learming, pp. 5852-5861. PMLR, 2018.
Pham, N.-Q., Truong, T., Tran, Q.. Nguyen, T. M., Phung, D., and Le, T. Promoting ensemble diversity with inter- active bayesian distributional robustness for fine-tuning foundation models. In Proceedings of the 42st Interna- tional Conference on Machine Learning (ICML), Proceed- ings of Machine Learning Research, Vancouver, Canada, 2025. PMLR.
Zhang, X., Xu, R., Yu, H., Dong, Y., Tian, P., and Cui, P. Flatness-aware minimization for domain generalization. In Proceedings of the IEEE/CVF International Confer- ence on Computer Vision pp. 5189-5202, 2023.
Qu, Z., Li, X., Duan, R., Liu, Y., Tang, B., and Lu, Z. Gen- cralized federated learning via sharpness aware minimiza- tion. In Chaudhuri, K., Jegelka, S., Song, L.,, Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings Machine Learning Research, pp. 18250-18288 PMLR, 17-23 Jul 2022. URL https: proceedi ngs ml.. press /v 62 qu22a html.
Zhou, T.-Y,, Suh, N., Cheng, G., and Huo, X. Approxima- tion of RKHS functionals by neural networks, 2024. URL attps://arxie wttpi//.rggaab///0//0/ 121877
Ritter, H., Botev, A., and Barber, D. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018- Conference Track Proceedings, volume 6. International Conference on Representation Learning, 2018
Rossi, S., Marmin, S., and Filippone, M. Walsh-Hadee Hadamard variational inference for Bayesian deep learning. Ad- vances in Neural Information Processing Systems, 33: 9674-96866 2020.
Shalev-Shwartz, S. and Ben-David, S. Understanding Ma- chine Learning: From Theory to Algorithms. Cambridge University Press, 2014. ISBN 978-1-107-05513-5
Wang, P., Zhang, Z., Lei, Z., and Zhang, L. Sharpness- aware gradient matching for domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3769-3778, 2023.
Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient Langevin dynamics. Proceedings of the 28th International Conference on International Conference on

---

Supplement to "Improving Generalization with Flat Hilbert Bayesian Inference"
The Appendix will be structured as follows: Section presents the proofs to our theoretical results. Next, Section B is dedicated for the additional experiments to assess the effectiveness of our method, including the effect of the number of particles, and the effect of kernel choices. Then, Section discusses the related experimental details, including the chosen hyperparameters data augmentations used in our experiments.
A. All Proofs
We introduce a few additional notations for the sake of the missing proofs of the main theoretical results. Given a RKHS H equipped with the inner product (ts t and the norm operator || functional space H to map:
1/x. We define the single-sample loss function on the
i: Hd x X x y
(f.I,)) + ((f((xx)))
Define the general functional loss LD(f) B(x,))=D(EC (x, 2))] and empirical functional loss Ls(f) ((f, Throughout the proof, we assume that the parameter space is bounded by ||0]] S H, the data is bounded that ||>||| < R. < R for some R, H. IR, and the loss function i is o-Hölder continuous. Indeed, this HöHder continuity property holds for wide range of loss functions, including the MSE loss, Huber loss, Hinge loss, Smooth LI loss, and Logistic loss.
We introduce the following lemmas that will be used throughout the proof of our main theorems.
Lemma A.1 (Approximation of RKHS functionals), Let d E
R]KK for some K E R. Consider K. = (s H ll < 1} with H. induced by some Mercer kernel which is a-Holder continuous for E (0, 1) with constant Ox 2 0. Suppose is s- Holder continuous for E (0, 1] with constant Cj 2 0. There exists some My E N such that for every M E N with M >> Mo, by taking some fixed T. = (t,) with N E N, we have tanh neural network G with two hidden layers
of widths at most N (MM - 1) and 3 N+1 (5M) parameters satisfying
sup [F()) G(f(F) (C)( 7N RCG rek
(20)
M
with
= Cx | (10
where K[t the Gram matrix of F.
Proof. The proof can be found in (Zhou et al., 2024)
Lemma A.2 (Product of RKHSs). Given n. RKHSs H4, H2, ... Hrw each defined on corresponding sets X1, X2...., ,,,, with kernels (1 ..).... (n Vn) respectively. Then, H = 10_1 = H2 x H2 x ... x Hr is also an RKHS, with kernel K that is the product of the individual kernels.
Proof. The product space H = =x" 2t=120 H- consists of tuples of functions product H as:
... fn). f..
Firstly, we define the inner
(21.f2,-... (n)) 92.
Inm) 7 ..)))..)))))).............)).)))
Σ (f..))..
1m1
12

---

Improving Generalization with Flat Hilbert Bayesian Inference.
This definition naturally defines Hilbert space structure H since each H, is Hilbert space, and the of inner products is linear and positive definite. Now we define the kernel for the product space:
, "2,"" " ; In); 1/, ... ynn) = II ((11,3)).
Notice that the pointwise product positive definite kernels is a positive definite kernel, hence this kernel is valid.
We now verify the reproducing property of H. Consider function f2. .fn) E H, and evaluate the function point (2, ..n) X
The reproducing property in each individual RKHS H, implies that:
f(.))) kx Mn. n.
Hence, for the function f .., (n)) get:
0(1 ,k,) (2(22), (n(san) ((fi,ki )m,, (J2, k2 (2 "oa " (Sn, kin (In HH (2,,, ... Sn), (21 ) (x2.))) "). kin (2 1))
Thus, the reproducing property holds for the product space H. Since H is a Hilbert space and the kernel satisfies the reproducing property, we conclude that 7 H4 is another RKHS.
A.1. Proof of Proposition 4.1
Proposition A.3. Consider the problem of finding the distribution Q that solves:
Q-=minn Eo-e DKI Qo
(21)
where we search over Q absolutely continuos w.r. Po. and the second term is the regularization term. The closed-form solution to this problem is the population posterior whose density has the form:
q" (0) exp(
Proof. This proposition is the general case of Theorem 3.1 by (Nguyen et al., 2023). Denote g(-) as the density function of Q We have:
Eg Ee,c(Cr(()) + (Ql(())P)) / (D()))(0)d0 / ( log (0) d0
p(0)
The Lagrangian is given by:
Mgg ()) J (o())) + / (()) ((0)log log (())
do al a(0)do - 1). p(0)
Taking derivative with respect to (0),, it follows
Lo + (0) + 1 p00) + 0,
( = (xp(=fo()))pp()ooom exp( - 1),
which implies that
(0) x exp( (x)(=C)())))())
Then, the optimal solution is the population posterior ()))S), which concludes the proof.
13

---

Improving Generalization with Flat Hilbert Bayesian Inference.
A.2. Proof of Theorem 4.2 Theorem A.4. For any p > O and any distribution D, with probability 1 over the choice of the training S D".
Lo(f) & max (f(t')
11'' 10
N' log
c +
pep2
1
log((N) Nr
+ 4 log 3 + 8 (log(6n + 3k)
Proof. is functional that maps from Hd x x X/ to R. Notice that 74 is RKHS, X Ra and y RRb for some a, b Z are Euclidean spaces, which are also instances of RKHS. Moreover, the product of RKHS's is also RKHS according to Lemma A.2. Hence, 71t x X x y is also RKHS. According to Lemma A.1, there exists N points 0 [0:1/= C O, and two-layer neural network Gw parameterized by W so that
RN2RCG
( (()) S. (C)(Ee(t))
M
for every (f,I,y) E Hd x X x y. Consider f E Hd so that |f - -fll S 0, RR and
Recall that we have 0 = =10 where each
and f(0) [f(0)). .f(0x)] N where N' Nd. Then, have:
Σ ( 1/2
II((0) ('
Now let = [2,,,,,, [s] and f' = [f., S'll where fú, fJz E H. we have:
d Σ
d (k(0,,)) - 17) S
-1, (
-
k(00,
IIII, - fill k(0, 0,)10 8'll
1=1
#=1
Therefore, it follows:
I1(00) s')) S Σ k(O 0.) 1/2 1- 1'11-
Under the assumption that k) < C (this is true for the widely used kernels, for example, C' = for the RBF kernel) implies (0) f() < Pllf f'llza I Pp where = CN1/22 Denote = f(0) E R for some N' E Z, by invoking inequality from (Foret et al., 2021), let p' = pP, it follows that:
Lp(f) = ((a.)) v)1 s (e RCr(e((()) + RC
M
7N2
-((a))>DGGw(l v) + (Cr((c(())
M
max 1 Sp
Gw (w()) h(M,N)
i=1
N' log
C 1
pd
1+
N) Nr
+ 4log 3 + 8 logg 6n + 3k)
where we defined NMM,N) = RCF(ex(())*+ RCC M
By definition, RKHS is closed Hilbert space. Then, there
exists sequence { so that f%, (0) that gets arbitrarily close to 0'. Then, for any > 0, follows:
14

---

Improving Generalization with Flat Hilbert Bayesian Inference.
n Σ
[D(f) =
max ||ō
Gw ,2,1) N(M,N)
p n
i=1
2
N' log
1+ p
1+
log(N) N
+loog 38 + 1og((6 + 3k)
max
Σ
Gw (f ((M c0
-
n i=1
N' log
c 1+
1 /
log(N) Nr
+ 4 4 + 8 8log(6n + 3k)
p2p2
n
Σ Gw (f) + +(,,N)+(011
max
f11ssp
N' log
1 p-pz
1
log N) N
+ 4 log 3 llg(6nn 3k)
- +s(f")+hMMM)) c0(1)
max
I fllisp
N' log
1
pppz
C
1
(og( Nr
+ 4 log + 8 Slog(6n + 3k)
71
This is true for any e > Moreover, we can choose ex and M to be arbitrarily small so that h(M, N) - 0. Hence, implies
Lp(f) v max Ls(F')
Il fllss
N'
C 1
p2p2
1
log(N) Nr
+ 4 log + +8 88log66n + 3k)
n
which concludes our proof.
A.3. Proof of Theorem 4.3
Now we can prove the Theorem 4.3. We restate the theorem
Theorem A.5. For any target distribution p. reference distribution g, and any p > we have the following bound between
15

---

Improving Generalization with Flat Hilbert Bayesian Inference.
the general KL loss and the empirical KL loss
DRL. (anll((11)))
(1))ll((1s))
max f f'e8,(f)
DrL.
N' log 1+
c
1
log(N)
+ 4 log 7/ + 8 8log 66n + 3k)
p-p2
Proof. Consider the left-hand side, we have:
DKL(g1)]lo((1D)) - f (n)
(D( + +log (0) p(0)
+ log ZD d0
f (0)
9 9(5( (0)
(00,) + log
log ZD
d0
p(0)
/
(1) p(0) d0
--aaan (n) log Zdd0,
log
where Zd is the normalizing constant. On the other hand, also have:
DKL (11(91)111(01S) / |6و /anl
(0) + log
(0) p(0)
log Zs d0
= (
Σ
n
l(0, Fz wi) + log
(0)
log Zs d0
p(0)
/
(1): v1) log
(0) p(0)
+ log Zs ) + +
/ (1))(()
log Zsd0,
n
where Zg is the normalizing constant. We define to be the functional such that:
Ż Hd x X x x -
(0) (E(0,x.)) + log (n)
(f,r,y) i(f,x.))
p(0) d0.
According to Theorem 4.2, we have:
(o(f) II -f
max
Ls(f')
(22)
2
N'log
1 pepz
1 V
log(N)
+ 4 log 2 logg66 3k)
(23)
n
Moreover, notice that we have Zp = Elexp(Co(0))pp()) Zs = (0))()))(0)))())) Let us denote VCdim(F) = dvc
where F (fo E 0). According to Theorem 6.11 in (Shaley-Shwartz & Ben-David, 2014) (Page 51), with the
16

---

Improving Generalization with Flat Hilbert Bayesian Inference.
probability least 1 6, we have
4 4+ v log TH (2m)
sup AE
LD (0) - Ls (0)
8v 2m
dvc
Note that the growth function Th (2n) _
2em dvc
we obtain
+ dvc log dve
sup 0c0
LD (0) Ls (0)
8v2n
Therefore, for all 0 E O we gain
2en 4 v dyc log
dvc
exp (0)) I exp (Cs (0)) exp
8v2n
It follows that Which is equivalent to Therefore, have
Zen 4 dyv dve
exp (CD (0) do S exp
exp (Cs (0)) do.
8v2n
2en 4 / dve log
dvo
ZD S exp
Zg
8v2n
Zen + ddv log dvc
log ZD < log Zs +
8v
Combining the Inequalities 23 and A.3, it follows that:
Dkr. ( (anill((o1))
DKL. ((Lrill((1s))
<
max Sp
N' N' log
1+ pppz
1
log(N) NN
log + 1lgg(6nn + +3k)
2
Zen 4 v d log dvc
8/22
which concludes proof.
B. Additional Experiments
B.1. Effect of #particles
SVGD
To understand the impact of varying the number of particles, we ex- periment on the seven Natural datasets, reporting both accuracy and per-epoch runtime. We compared FHBI with SVGD and SAM. Figure and Table indicate that multiple particles result in significant perfor- mance improvements compared to single particle. However, while increasing the number of particles enhances performance, it introduces tradeoff regarding runtime and memory required to store the models. Based on these observations, we found that using #PARTICLES = 4 provides an optimal balance between performance gains and compu- tational overhead.
Runtime (s)
100
10
Figure 5, Runtime by #PARTICLES.
17

---

Improving Generalization with Flat Hilbert Bayesian Inference.
#PARTICLES
CIFAR100
Caltech101
DTD
Flower102
Pets
SVHN
Sun397
1p (SAM)
72.7
90.3
71.4
99.0
90.2
84.4
52.4
4p
74.8
93.0
74.3
99.4
92.4
87.5
56.5
10p
75.0
93.2
75.0
99.1
92.4
87.9
58.3
Table 3. Accuracy by #PARTICLES.
B.2. Effect of kernel choice
The implementation of FHBI relies on the choice of the kernel k. In our experiments, we selected the RBF kernel due to its widespread use in the kernel methods literature, known for its strong representational capabilities and its ability to balance underfitting and overfitting through the kernel width parameter o. To evaluate the impact of different kernel choices, we tested our method on the four Specialized datasets using the polynomial kernel degree 10 as a comparison. The results, summarized Table 4, indicate that while the polynomial kernel slightly underperforms relative the RBF kemel,, the difference is minimal, with performance gap of less than 0.3%.
Table 4. Classification accuracy on the Specialized datasets with diffcrent kernel choices
Kernel
Camelyon
EuroSAT
Resise45
Retinopathy
AVG
RBF
85.3
95.0
87.2
79.6
86.8
Polynomial (d=10)
85.0
94.9
86.8
79.2
86.5
B.3. Memory consumption and scalability to larger models
To assess the scalability to larger models, we conduct an additional experiment on the larger ViT-B/16. In Tables 5 and 6, we report the memory usage and training time for the CIFAR- -00, Caltech- -101, and Pat ch ch-Came Camelyon datasets.
Table VRAM consumption (MB) ViT-B/16 and ViT-I/16
Architechture
CIFAR-100
Caltech- 101
Patch-Camelyon
ViT-B/16
12541
12539
12535
ViT-L/16
33426
33684
32126
B.4. Sensitivity to the RBF kernel length scale
As noted before, initially tune the length scale parameter a from the candidates set {0.7, 1, To further investigate sensitivity to this hyperparameter, value of 0.1 and a large value of 2.5. The results the Natural datasets, reported in Table 7, indicate that the performance remains robust within reasonable range. We also observe slight degradation when using extremely small values (c.g., σ 0.1), where the model tends to overfit, or large values (e.g., = 2.5), where the model tends to underfit.
expand the range to (0.1, 0.7, 1, 1.2, 2.5), where additionally include small
C. Experimental Details
C.1. Chosen Hyperparameters
For each experiment, we conducted five runs of FHBI and reported the mean and standard deviation. All Bayesian methods were trained with four particles the same set of LoRA parameters. We used ten warm-up epochs, batch size 64, the Gaussian kernel, and the cosine annealing learning rate scheduler for all settings. The experiments were run with PyTorch on a Tesla V100 GPU with 40GB of RAM. FHBI involves three hyperparameters: the learning rate ascent step size and kernel width a. To tune these hyperparameters, we grid-seardh hyperparameters on the validation set, where the key hyperparameters are: the kernel width a, the initial learning rate e, and the ascent step size p. The candidate sets are formed as E (0.15, 1, 1.5, 1.7, 1.9, 2,1, 2.3, 2.5), (0.01, 0.03, 0.05). E 10.7,1,1.21
18

---

Improving Generalization with Flat Hilbert Bayesian Inference
Table VRAM consumption (MB) ViT-B/16 and ViT-I/16
Architechture
CIFAR-100 Caltech- 101
Patch-Camelyon
ViT-B/16
5.55
5.66
5.46
ViT-L/16
17.45
17.35
17.02
Table 7. Sensitivity the RBF length scale
CIFAR-100
Caltech-101
DTD
Flowers
Pets
SVHN
Sun397
0.1
72.1
91.6
74.0
97.9
90.2
85.3
52.2
0.7
73.8
91.8
73.3
98.7
92.4
86.7
56.1
1
73.6
92.7
72.7
99.1
91.9
87.3
54.3
1.2
74.1
93.0
74.3
98.3
92.4
86.4
56.5
2.5
69.2
90.9
69.4
97.5
90.9
84.6
52.6
C.2. Data Augmentations
Our implementation is based on the repository V-PETL. Similar this repository, we use different data augmentation among the following three augmentations for each dataset. The data augmentations that we used for each setting are:
For CIFAR100, DTD, Flower102, Pets, Sun397
sel cranssorssstttraaiio. tranaforms.com Compose
transforms transforms. RandomResizeedcog
(self. sclf.size),
scale= (self self.max_scaleg
) transforms. cransforrs..anndomooooooooo waa.oomloo1oooooaalliiii.oe (sl...ffippp...bbb trans www.frmmm..iiialllaameeettt.....
if self.use uelfiuse_trivial aug
else trans RandAugment
self.rand_aug_m)
transforms.. PoTensor (),
MoansformsNNoraaliii (mean- [8.485, 0.456, 0.406), transforms Normalorizziiaoo
s0d=(0.22) 229, 0.224, 0.2251)1),
transforms. Randomnrasing (prealdeeppp111ffeeee__e_ppp
sel f. frtrsssforms_tttt transforms.
[
transforms ( (self.size, self.size),
),
transforms PoTensor ( (),
transforms Normalize (mean= (mean=[0.485, 48.485, 0.456, 0.406),
atd- (2.229, 0.224, 0.225])1),
For
Caltech101 01, Clevr-Dist, Dsprites-loc, Dsprites-Ori, Smal NNrrb-Norrttto Azi,
Smal SmallNorb-ENM Ele:
self. crafforranffonttrrain transforms. coansfornsscompoee Compose [[
transforms. ResioorrmmsReeiiieeee (224, 2
transforms ToTensor Tooeenso (),
transforms. Maaassformss.orma.iee.. (mean= 10.485, 0.456, 0.406],
(t. 0.224, 0.225])])
19

---

Improving Generalization with Flat Hilbert Bayesian Inference.
self. cransform_teett transforns Compose
transforms zesizze (224, 224)), transforms ToTensor (), transforms. Normal ize (mean- [0.485, 0.456, 0.406], std=(0.229, 0.224, 0.225])])
For
Clevr-Count, DMLab, EuroSAT, KITTI, Patch Camelyon, Resisc45, SVHN, Diabetic Retinopathy:
from timm.data import greate_transtin
traff..raanformmtraaaaa - create_transsorm
input_size- (224, 224),
is_training=True,
color_jitte==0.44
o re_prob=0.0,
re_mode-/pixel',
re_count=l,
interpo' attpppo:laaioon"iibibiiii
aug.transfom.trassforr [0] transforms Resize (224, 224),
interpolation-3)
telf.transform_test transforms Compose
transforms Resize ( (224, 224)),
transforms. ToTensor (),
transforms. Normallze (mean- [0. 485, 0.456, 0.406], std=10.229) 0.224, 0.22511
20


Reviews:
[{"review_id": "jSHdao5vb1", "reviewer": "Reviewer_JWDR", "summary": "The paper introduces Flat Hilbert Bayesian Inference (FHBI), a novel algorithm designed to enhance generalization in Bayesian inference by extending principles from finite-dimensional Euclidean spaces to infinite-dimensional reproducing kernel Hilbert spaces (RKHS). FHBI employs an iterative two-step procedure involving adversarial functional perturbation and functional descent within RKHS, supported by a theoretical framework that analyzes generalization in infinite-dimensional spaces. Empirically, FHBI is evaluated on the VTAB-1K benchmark, which includes 19 diverse datasets across various domains, where it consistently outperforms nine baseline methods by significant margins, demonstrating its practical efficacy and potential for real-world applications.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and would like to address the concerns as follows:\n\n+ **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.\n\nVRAM consumption (MB):\n| Architechture | CIFAR100 | Caltech101 | Patch Camelyon |\n| ------------- | -------- | ---------- | -------------- |\n| ViT-B/16              |     12541     |     12539       |   12535             |\n| ViT-L/16          | 33426     | 33684       |      32126          |\n\nTraining time (s/iter):\n| Architechture | CIFAR100 | Caltech101 | Patch Camelyon |\n| ------------- | -------- | ---------- | -------------- |\n| ViT-B/16          | 5.55     | 5.66       |      5.46          |\n| ViT-L/16              |     17.45     |    17.35      |   17.02             |\n\n\n+ **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.\n\n| $\\sigma$ | CIFAR100 | Caltech101 | DTD | Flowers102 | Pets | SVHN | Sun397 |\n| -------- | -------- | --------- |-------- |-------- |-------- |-------- |-------- |\n| 0.1     | 72.1     | 91.6     | 74.0    |97.9     |90.2     |85.3     |52.2     |\n| 0.7     | 73.8     | 91.8     | 73.3    |98.7     |92.4     |86.7     |56.1     |\n| 1       | 73.6    | 92.7    |72.7     | **99.1**     |91.9    | **87.3**     |54.3     |\n| 1.2     | **74.1**     | **93.0**        | **74.3**     |98.3    | **92.4**     |86.4     | **56.5** |\n| 2.5    | 69.2     | 90.9    |69.4    |97.5     |90.9     |84.6     |52.6    |\n\n+ **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.", "claims_and_evidence": "Appropriate", "methods_and_evaluation_criteria": "The evaluation method refers to the work of BayesTune, which I believe is acceptable in this work.", "theoretical_claims": "All proofs except those in the supplementary materials have been checked. It can be considered that there are no significant errors.", "experimental_designs_or_analyses": "The experimental designs and analyses are appropriate", "relation_to_broader_scientific_literature": "This paper is a combination of work on flat minimizers and particle-based Bayesian methods, and a detailed investigation has been conducted on the historical work of both.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "1.\tAlthough four particles were chosen as the equilibrium point, the actual training time and resource consumption (e.g., GPU memory) have not been quantified. Particularly, the scalability on large models (e.g., ViT-L/16) has not been verified.\n2.\tThe sensitivity of the results to the length-scale of the RBF kernel was not analyzed.\n3.\tAlthough the theory proposes \"functional sharpness,\" the experiments only indirectly assess it through empirical metrics, without designing a direct measurement method for the RKHS space.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sPfgTEwLXo", "reviewer": "Reviewer_SGnT", "summary": "The paper introduces a novel Bayesian inference method designed to improve generalization by leveraging functional sharpness-aware particle sampling in RKHS. The key innovation lies in combining sharpness aware minimization (SAM) with SVGD in infinite-dimensional RKHS to form the proposed FHBI. The authors extend theoretical generalization bounds from finite-dimensional Euclidean spaces to infinite-dimensional functional spaces. Empirical evaluations on the VTAB-1K benchmark show that FHBI outperforms several SAM, SVGD and their combinational baselines.\n\n## update after rebuttal\nThanks to the authors for the clarification. After reading the initial response to my questions, as well as the other reviews and replies in general, I will keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's feedback. We would like to address the concerns as follows: \n\n+ **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. \n\n+ **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.\n\n+ **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.", "claims_and_evidence": "I think the claims made in the paper are clear, and supported by both theoretical proofs and experimental results.", "methods_and_evaluation_criteria": "The methodology builds upon: 1) generalization theory in RKHS, extending flatness-based optimization(SAM) beyond Euclidean settings. 2) particle optimization of Stein Variational Gradient Descent (SVGD). The idea of combining SAM and SVGD in RHKS is very clear and makes sense to me.", "theoretical_claims": "The theoretical contributions are significant:\n1. Theorem 4.2 extends Euclidean space generalization bounds to RKHS.\n2. Theorem 4.3 bridges functional sharpness minimization with Bayesian inference. This establishes a connection between empirical and population KL loss.\n\nI haven’t carefully checked every mathematical detail, but the proof appears correct.", "experimental_designs_or_analyses": "The experimental design is strong. It evaluates FHBI on the VTAB-1K classifications and compares it with several baselines, including SVGD, SAM-based methods, SGLD, SVGD+SAM, SGLD+SAM, and deep ensembles. Moreover, it conducts a wide ablation study about the number of particles, sharpness, and gradient diversity to further demonstrate the validity of FHBI.", "relation_to_broader_scientific_literature": "The paper is situated within the fields of Bayesian deep learning (variational methods, particle-based sampling), sharpness-aware optimization (SAM), and kernel methods in function space (RKHS).", "essential_references_not_discussed": "I think the paper sufficiently cites prior related work.", "other_strengths_and_weaknesses": "Strengths:\n1. The idea of combining SAM with SVGD in RKHS to improve generalization is novel and generally makes sense.\n2. Novel theoretical extension of generalization bounds in Euclidean space to the function space.\n3. Strong empirical performance across diverse benchmarks.\n\nWeaknesses:\n1. I'm a little confused about the relationship between FHBI and SAM. It is not clearly demonstrated in the objective function.\n2. Some notation is confusing, e.g., general posterior and population posterior. Are they referring to the same thing?\n3. The runtime comparison in Fig. 3 only with SVGD.", "other_comments_or_suggestions": "There is a typo in the line 204: the change of variable $q(T^{-1}(\\theta))$ should be $q(T^{-1}(\\vartheta))$", "questions_for_authors": "1. How does FHBI compare computationally to other baselines in terms of runtime and memory overhead besides SVGD?\n2. As you conduct the ablation study in Sec.6 to compare the particle sharpness of FHBI and SVGD, can FHBI and SAM be compared in detail using some metric to illustrate the advantages of particle-based sampling techniques?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3DC5Haw0RU", "reviewer": "Reviewer_nMfi", "summary": "This paper proposes a sharpness-aware version of stein variational gradient descent and applies it to neural network fine tunings.\nThe goal is to learn a map (in a RKHS) that transforms from the reference distribution to the true posterior by minimizing the KL divergence.\nTo motivate the sharpness-aware optimization for Bayesian inference, the authors prove a generalization bound, which upper bounds the true KL divergence by a worse-case KL divergence on the training data in a small neighborhood of RKHS.\nThen the algorithm generally resembles the idea of Foret et al. (2021), but the difference is that the optimization is carried out in the RKHS by functional gradient descent.\nThe authors evaluate the propose method in neural network fine tuning, where they demonstrate improve accuracy and uncertainty estimate (measured by the ECE score).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.", "claims_and_evidence": "I find the claims (both theoretical and empirical) in the paper are generally well supported.", "methods_and_evaluation_criteria": "Yes, they make sense.", "theoretical_claims": "I skimmed through the proof of Theorem 4.2, but I did not check the proof of Theorem 4.3.\nThe authors claim that these results are non-straightforward extensions of prior results of Foret et al. (2021) because RKHS are typically infinite dimensional.\nHowever, the proof seems to a simple combination of existing results.\nThey first construct a two-layer neural network that approximates the RKHS function and then the theorem is proved by invoking the results of Foret et al. (2021).", "experimental_designs_or_analyses": "I read through the experiments in Section 5.", "relation_to_broader_scientific_literature": "This paper is a sharpness-aware variant of SVGD (Liu and Wang, 2016) by adapting the idea of SAM (Foret et al., 2021).\nWhen applied to training neural networks, the method inherits the benefits from the both worlds---better generalization and better uncertainty estimate.", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "1. The method has large memory consumption because it needs to maintain several copies of the neural network parameters.\nThus, the authors have focused on fine tuning as opposed to training from scratch.\nI assume that it is very hard to apply this method to training neural networks from scratch.\nBut that is a common limitation of all particle based Bayesian inference methods.\n1. The writing is of high quality.\n1. The experiments are comprehensive across many datasets and baselines.\nExperiments demonstrate that the proposed method improves not only the accuracy but also uncertainty estimates.\nAdditional ablation experiments demonstrate that the proposed method indeed yields more flat solutions.", "other_comments_or_suggestions": "1. The equations in Algorithm 1 are hard to parse of the line breaks.\nMaybe use a double-column algorithm environment.", "questions_for_authors": "None.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces Flat Hilbert Bayesian Inference (FHBI), a novel algorithm that incorporates flatness into Bayesian inference. The method combines an adversarial functional perturbation step with a functional descent step in a reproducing kernel Hilbert space (RKHS), and is theoretically grounded to improve generalization. Notably, the theoretical development is also novel, extending prior analysis of generalization from finite-dimensional Euclidean spaces to infinite-dimensional function spaces.

While the concepts of flatness and sharpness-aware minimization have been explored extensively in the community, there have been few successful attempts to incorporate flatness-seeking behavior into Bayesian posterior inference. In this regard, FHBI fills an important gap. All reviewers are positive about the novelty and significance of the contribution, and I share their assessment.
Decision: Accept (poster)