Improving Generalization with Flat Hilbert Bayesian Inference

## Abstract

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 seven 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. Our code is available at https://anonymous.4open.science/r/Flat-Hilbert-Variational-Inference-008F/ .

## 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 issue is variational inference, which seeks to approximate the true posterior distribution with simpler forms, known as approximate posteriors while optimizing a variational lower bound. Several techniques have been developed in this area, including those by Kingma & Welling (); Kingma et al. (), and Blundell et al. (), who extended the Gaussian variational posterior approximation for neural networks, as well as Gupta & Nagar (), who enhanced the flexibility of posterior approximations. In addition to variational methods, various particle sampling techniques have been proposed for Bayesian inference, especially in scenarios requiring multiple models. Notable particle sampling methods include Hamiltonian Monte Carlo (HMC) (Neal, ), Stochastic Gradient Langevin Dynamics (SGLD) (Welling & Teh, ), Stochastic Gradient HMC (SGHMC) (Chen et al., ), and Stein Variational Gradient Descent (SVGD) (Liu & Wang, ). Each method contributes to a deeper understanding and more practical application of Bayesian inference in deep learning.

Besides quantifying uncertainty, tackling overfitting is a major 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 effectively improve model generalization (Keskar et al., ; Kaddour et al., ; Li et al., ). Among the flat minimizers, Sharpness-Aware Minimization (SAM) (Foret et al., ) has emerged as a practical approach by concurrently minimizing the empirical loss and reducing the sharpness of the loss function. Recently, SAM has demonstrated its versatility and effectiveness across a wide range of tasks, including meta-learning (Abbas et al., ), federated learning (Qu et al., ), vision models (Chen et al., ), and language models (Bahri et al., ).

Figure: Figure 1: Schematic of SAM w. independent particles (green), SVGD (orange), and our FHBI (Algorithm 1) (black) updates. SAM’s particles are not aware of other’s trajectories. SVGD only seeks the modes and promotes spatial diversity. FHBI seeks the modes, minimizes sharpness, and promotes spatial and angular diversity.
Refer to caption: x1.png

#### Contribution.

We bridge the gap between the flat minimizers and particle sampling to introduce a Bayesian inference framework with improved generalization ability. To accomplish this, we first present Theorem 1, which strengthens prior generalization bounds from finite-dimensional Euclidean spaces to the reproducing kernel Hilbert spaces (RKHS), which are broader and more general functional spaces that are typically infinite-dimensional. 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 2 translates these notions of functional sharpness and generalization in RKHS into the context of Bayesian inference. This analysis establishes a connection between the general and empirical KL loss, providing a strategy to enhance generalization by minimizing the general KL loss. Motivated by these two theorems, we derive Flat Hilbert Bayesian Inference (FBVI), a practical algorithm that employs a dual-step functional sharpness-aware update procedure in RKHS. This approach improves the generalization of sampled particles, thereby enhancing the quality of the ensemble. Overall, our contributions are as follows:

- 1.
We present a theoretical analysis that characterizes generalization ability over the functional space. This analysis generalizes prior works from the Euclidean space to infinite-dimensional functional space, thereby introducing the notion of functional sharpness i.e., the sharpness of the functional spaces.
- 2.
Building on this theoretical foundation, we propose a practical particle-sampling algorithm that enhances the generalization ability over existing methods. We conducted extensive experiments comparing our Flat Hilbert Bayesian Inference (FHBI) algorithm with seven baselines on the VTAB-1K benchmark, which includes 19 datasets across various domains and semantics. Experimental results demonstrated that our algorithm outperforms these baselines by notable margins.

The paper is structured as follows: Section 2 reviews the related works on Bayesian inference and the development of flat minimizers. Section 3 provides the necessary background and notations. Section 4 discusses the motivation and theoretical development behind our sharpness-aware particle-sampling approach. Section 5 presents experimental results, comparing our algorithm against various Bayesian inference baselines across diverse settings. Section 6 offers a deeper analysis of FHBI’s behavior to gain further insight into its effectiveness over the baseline methods.

## 2 Related Works

#### Sharness-aware minimization.

Flat minimizers have been shown to be more robust to the shifts between training and test losses, thereby enhancing the generalization ability of neural networks  (Jiang et al., ; Petzka et al., ; Dziugaite & Roy, ). The relationship between generalization and the width of minima has been studied both theoretically and empirically in several prior works  (Hochreiter & Schmidhuber, ; Neyshabur et al., ; Dinh et al., ; Fort & Ganguli, ). Consequently, a variety of methods have been developed to search for flat minima  (Pereyra et al., ; Chaudhari et al., ; Keskar et al., ; Izmailov et al., ).

Among the flat minizer, Sharpness-Aware Minimization (SAM), introduced by Foret et al. (), has gained significant attention due to its effectiveness and scalability. SAM’s versatility has been leveraged across a wide range of tasks and domains, including domain generalization (Cha et al., ; Wang et al., ; Zhang et al., ), federated learning (Caldarola et al., ; Qu et al., ), Bayesian networks (Nguyen et al., ; Möllenhoff & Khan, ), and meta-learning (Abbas et al., ). Moreover, SAM has demonstrated its ability to enhance generalization in both vision models (Chen et al., ) and language models (Bahri et al., ).

Nevertheless, these studies are constrained to finite-dimensional Euclidean spaces. In this work, we strengthen these generalization principles to infinite-dimensional functional spaces and propose a particle-sampling method grounded in this theoretical framework.

#### Bayesian Inference.

Two main strategies were widely employed in the literature of Bayesian inference. The first paradigm is Variational Inference, which aims to approximate a target distribution by selecting a distribution from a family of potential approximations and optimizing a variational lower bound. Graves () introduced the use of a Gaussian variational posterior approximation for neural network weights, which was later extended in Kingma & Welling (); Kingma et al. (); Blundell et al. () with the reparameterization trick to facilitate training deep latent variable models. Louizos & Welling () proposed using a matrix-variate Gaussian to model entire weight matrices (Gupta & Nagar, ) to increase further the flexibility of posterior approximations, which offers a novel approach to approximate the posterior. Subsequently, various alternative structured forms of the variational Gaussian posterior were proposed, including the Kronecker-factored approximations (Zhang et al., ; Ritter et al., ; Rossi et al., ), or non-centered or rank-1 parameterizations (Ghosh et al., ; Dusenberry et al., ).

The second paradigm in the literature 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 Hamiltonian Monte Carlo (HMC) (Neal, ). However, HMC requires the computation of full gradients, which can be computationally expensive. To address this, Stochastic Gradient Langevin Dynamics (SGLD) (Welling & Teh, ) integrates first-order Langevin dynamics within a stochastic gradient framework. Stochastic Gradient HMC (SGHMC) (Chen et al., ) 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, ), closely related to our work, uses a set of particles that converge to the target distribution. It is also theoretically established that SGHMC, SGLD, and SVGD asymptotically sample from the posterior as the step sizes approach zero.

## 3 Backgrounds and Notations

#### Bayesian Inference.

Consider a family of neural networks $f_{\bm{\theta}}(x)$, where the random variable ${\bm{\theta}}$ represents the model parameters and takes values in the model space $\Theta\subset\mathbb{R}^{d}$. We are given a training set ${\mathcal{S}}=\{(x_{i},y_{i})\}_{i=1}^{n}$ of $n$ i.i.d observations from the data space ${\mathcal{X}}\times{\mathcal{Y}}$, and the prior distribution of the parameters $p({\bm{\theta}})$. In the literature on Bayesian inference problems, prior works typically focus on approximating the empirical posterior $\mathbb{P}_{\mathcal{S}}$, whose density function $p({\bm{\theta}}|{\mathcal{S}})$ is defined as:

$$ $p({\bm{\theta}}|{\mathcal{S}})\propto p({\bm{\theta}})\prod_{i=1}^{n}p(y_{i}|x _{i},{\mathcal{S}},{\bm{\theta}}),$ $$

where the prior distribution $\mathbb{P}$ has the density function $p({\bm{\theta}})$. The likelihood term is proportional to

$$ $p(y|x,{\mathcal{S}},{\bm{\theta}})\propto\exp\Bigg{(}-\frac{1}{|{\mathcal{S}}| }\ell(f_{\bm{\theta}}(x),y)\Bigg{)}=\exp\Bigg{(}-\frac{1}{n}\ell(f_{\bm{\theta }}(x),y)\Bigg{)},$ $$

with the loss function $\ell$. Then, the empirical posterior is given by:

$$ $p({\bm{\theta}}|{\mathcal{S}})\propto\exp\Bigg{(}-\frac{1}{n}\sum_{i=1}^{n} \ell(f_{\bm{\theta}}(x_{i}),y_{i})\Bigg{)}p({\bm{\theta}}).$ (1) $$

More formally, the empirical posterior is equal to:

$$ $p({\bm{\theta}}|{\mathcal{S}})=\exp\Bigg{(}-\frac{1}{n}\sum_{i=1}^{n}\ell(f_{ \bm{\theta}}(x_{i}),y_{i})\Bigg{)}p({\bm{\theta}})/Z_{{\mathcal{S}}},$ (2) $$

where $Z_{\mathcal{S}}$ is the normalizing constant. We define the general and empirical losses as follows:

$$ $\displaystyle{\mathcal{L}}_{\mathcal{D}}({\bm{\theta}})=\mathbb{E}_{(x,y)\sim{ \mathcal{D}}}[\ell(f_{\bm{\theta}}(x),y)],$ $\displaystyle{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}})=\mathbb{E}_{(x,y)\sim{ \mathcal{S}}}[\ell(f_{\bm{\theta}}(x),y)]=\frac{1}{n}\sum_{i=1}^{n}\ell(f_{\bm {\theta}}(x_{i}),y_{i}).$ $$

The general 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 ${\mathcal{S}}$. Based on these definitions, the empirical posterior in Eq. 2 can be written as:

$$ $p({\bm{\theta}}|{\mathcal{S}})=\exp(-{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}} ))p({\bm{\theta}})/Z_{\mathcal{S}}.$ $$

Intuitively, models with parameters ${\bm{\theta}}$ that fit well to the training set ${\mathcal{S}}$ lead to lower empirical loss values, resulting in higher density in the empirical posterior. However, simply fitting to the training samples can lead to overfitting. To improve generalization, we are more concerned with performance over the entire data distribution ${\mathcal{D}}$ rather than just the specific sample ${\mathcal{S}}$. Accordingly, we define the general posterior as $\mathbb{P}_{\mathcal{D}}$ whose density is given by:

$$ $p({\bm{\theta}}|{\mathcal{D}})=\exp(-{\mathcal{L}}_{\mathcal{D}}({\bm{\theta}} ))p({\bm{\theta}})/Z_{\mathcal{D}},$ $$

with the normalizing constant $Z_{\mathcal{D}}$. This general posterior is more general than the empirical posterior, as it captures the true posterior of the parameters under the full data distribution. However, understanding the general posterior is particularly challenging because we can only access the empirical loss ${\mathcal{L}}_{\mathcal{S}}({\bm{\theta}})$, not the general loss ${\mathcal{L}}_{\mathcal{D}}({\bm{\theta}})$. In this paper, we deviate from prior approaches that primarily focus on approximating the empirical posterior and instead propose a particle-sampling method to approximate the general posterior.

#### Reproducing Kernel Hilbert Space (RKHS).

Let $k({\bm{\theta}},{\bm{\theta}}^{\prime}):\Theta\times\Theta\to\mathbb{R}$ be a positive definite kernel operating on the model space. The reproducing kernel Hilbert space (RKHS) ${\mathcal{H}}$ of $k({\bm{\theta}},{\bm{\theta}}^{\prime})$ is the closure of the linear span $\{f:f(\cdot)=\sum_{i}a_{i}k(\cdot,{\bm{\theta}}_{i}),a_{i}\in\mathbb{R},{\bm{
\theta}}_{i}\in\Theta\}$. For $f({\bm{\theta}})=\sum_{i}a_{i}k({\bm{\theta}},{\bm{\theta}}_{i})$ and $g({\bm{\theta}})=\sum_{j}b_{j}k({\bm{\theta}},{\bm{\theta}}_{j})$, ${\mathcal{H}}$ is equipped with the inner product defined by $\left<f,g\right>_{\mathcal{H}}=\sum_{ij}a_{i}b_{j}k({\bm{\theta}}_{i},{\bm{
\theta}}_{j})$. For all ${\bm{\theta}}\in\Theta$, there exists a unique element $K_{\bm{\theta}}\in{\mathcal{H}}$ with the reproducing property that $f({\bm{\theta}})=\left<f,K_{\bm{\theta}}\right>_{\mathcal{H}}$ for any $f\in{\mathcal{H}}$.

Given that $\mathcal{H}$ is a scalar-valued RKHS with kernel $k(\theta,\theta^{\prime})$, ${\mathcal{H}}^{d}={\mathcal{H}}\times{\mathcal{H}}\times\cdots\times{\mathcal{
H}}$ is a vector-valued RKHS of functions ${\bm{f}}=[f_{1},f_{2},\cdots,f_{d}]$ corresponding to the kernel $K({\bm{\theta}},{\bm{\theta}}^{\prime})=k({\bm{\theta}},{\bm{\theta}}^{\prime}
){\bm{I}}$. ${\mathcal{H}}^{d}$ is equipped with the inner product $\left<{\bm{f}},{\bm{g}}\right>_{{\mathcal{H}}^{d}}=\sum_{i=1}^{d}\left<f_{i},g
_{i}\right>_{\mathcal{H}}$.

Let ${\bm{F}}[{\bm{f}}]$ be a functional on ${\bm{f}}\in{\mathcal{H}}^{d}$. Similar to the definition by Liu & Wang (), the (functional) gradient of ${\bm{F}}$ is defined as a function $\nabla_{\bm{f}}{\bm{F}}[{\bm{f}}]\in{\mathcal{H}}^{d}$ such that for any ${\bm{g}}\in{\mathcal{H}}^{d}$ and $\epsilon\in\mathbb{R}$

$$ ${\bm{F}}[{\bm{f}}+\epsilon{\bm{g}}]={\bm{F}}[{\bm{f}}]+\epsilon\left<\nabla_{{ \bm{f}}}{\bm{F}}[{\bm{f}}],g\right>_{{\mathcal{H}}^{d}}+{\mathcal{O}}(\epsilon ^{2}).$ (3) $$

#### Stein Variational Gradient Descent (SVGD).

Given a general target distribution $p(\bm{\theta})$, SVGD (Liu & Wang, ) aims to find a flow of distributions $\{q^{(k)}\}_{k}$ that minimizes the KL distance to the target distribution. Motivated by the Stein identity and Kernelized Stein Discrepancy, SVGD proposes the update $q^{(k+1)}=q^{(k)}_{[{\bm{T}}]}$, in which ${\bm{T}}:\Theta\to\Theta$ is a smooth one-to-one push-forward map of the form ${\bm{T}}({\bm{\theta}})={\bm{\theta}}+\epsilon\phi^{*}_{p,q}({\bm{\theta}})$ in which:

$$ $\displaystyle\phi^{*}_{p,q}(\cdot)=\mathbb{E}_{{\bm{\theta}}\sim q}[{\mathcal{ A}}_{p}k({\bm{\theta}},\cdot)]$ and $\displaystyle{\mathcal{A}}_{p}\phi({\bm{\theta}})=\phi({\bm{\theta}})\nabla_{ \bm{\theta}}\log p({\bm{\theta}})^{\top}+\nabla_{\bm{\theta}}\phi({\bm{\theta} }).$ $$

Here, ${\mathcal{A}}_{p}$ is known as the Stein operator, which acts on $\phi$ and produces a zero-mean function ${\mathcal{A}}_{p}\phi({\bm{\theta}})$ when ${\bm{\theta}}\sim p$. Notably, while SVGD is designed for general target distributions $p$, in the context of Bayesian inference, it is only applicable to the empirical posterior rather than the general posterior, which we will discuss in detail in the next section.

## 4 Flat Hilbert Bayesian Inference (FHBI)

We focus on the Bayesian inference problem of approximating a posterior distribution based on a given dataset. In prior works, such as SVGD (Liu & Wang, ), when applying to the context of Bayesian inference, the methods are only applicable to the empirical posterior $p({\bm{\theta}}|{\mathcal{S}})\propto\exp(-{\mathcal{L}}_{\mathcal{S}}({\bm{
\theta}}))p({\bm{\theta}})$ because we only have access to the empirical loss. It is evident that when sampling a set of $m$ particle models ${\bm{\theta}}_{1:m}$ from $p({\bm{\theta}}|{\mathcal{S}})$, these particles congregate in the high-density regions of the empirical posterior $p({\bm{\theta}}|{\mathcal{S}})$, corresponding to the areas with low empirical loss ${\mathcal{L}}_{\mathcal{S}}({\bm{\theta}})$. However, to avoid overfitting, it is preferable to sample the particle models ${\bm{\theta}}_{1:m}$ from the general posterior $p({\bm{\theta}}|{\mathcal{D}})\propto\exp(-{\mathcal{L}}_{\mathcal{D}}({\bm{
\theta}}))p({\bm{\theta}})$, as this approach directs the particle models ${\bm{\theta}}_{1:m}$ towards regions with low values of the general loss ${\mathcal{L}}_{\mathcal{D}}({\bm{\theta}})$, thereby enhancing generalization ability.

### 4.1 Theoretical analysis

Motivated by this observation, we seek to advance prior work by extending our focus beyond the empirical posterior to tackle the problem of approximating the general posterior. Specifically, to enhance generalization ability, our objective is to approximate the target general posterior distribution $p({\bm{\theta}}|{\mathcal{D}})$ using a simpler distribution $q^{*}({\bm{\theta}})$ drawn from a predefined set of distributions ${\mathcal{F}}$. This is achieved by minimizing the KL divergence:

$$ $q^{*}=\operatorname*{arg\,min}_{q\in{\mathcal{F}}}D_{\mathrm{KL}}\Bigg{(}q({ \bm{\theta}})\|p({\bm{\theta}}|{\mathcal{D}})\Bigg{)}.$ (4) $$

Ideally, the set ${\mathcal{F}}$ should be simple enough for simple solution and effective inference while sufficiently broad to approximate a wide range of target distributions closely. Let $q({\bm{\theta}})$ be the density of a reference distribution. We define ${\mathcal{F}}$ as the set of distributions for random variables of the form $\mathbf{\vartheta}={\bm{T}}({\bm{\theta}})$, where ${\bm{T}}:\Theta\to\Theta$ is a smooth, bijective mapping, and ${\bm{\theta}}$ is sampled from $q$. By variable change, the density of $\mathbf{\vartheta}$, denoted as $q_{[{\bm{T}}]}(\cdot)$, is expressed as follows:

$$ $q_{[{\bm{T}}]}(\mathbf{\vartheta})=q({\bm{T}}^{-1}({\bm{\theta}}))|\det(\nabla _{\mathbf{\vartheta}}{\bm{T}}^{-1}(\mathbf{\vartheta}))|.$ $$

We restrict the set of the smooth transformations ${\bm{T}}$ to the set of push-forward maps of the form ${\bm{T}}({\bm{\theta}})={\bm{\theta}}+{\bm{f}}({\bm{\theta}})$, where ${\bm{f}}\in{\mathcal{H}}^{d}$. When $\|{\bm{f}}\|_{{\mathcal{H}}^{d}}$ is sufficiently small, the Jacobian of ${\bm{T}}={\bm{I}}+{\bm{f}}$ is full-rank where ${\bm{I}}$ denotes the identity map, in which case ${\bm{T}}$ is guaranteed to be a one-to-one map according to the inverse function theorem. Under this restriction, the problem is equivalent to solving an optimization problem over the RKHS:

$$ $\displaystyle{\bm{f}}^{*}=\operatorname*{arg\,min}_{{\bm{f}}\in{\mathcal{H}}^{ d},\|{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq\epsilon}D_{\mathrm{KL}}\Bigg{(}q_{[{ \bm{I}}+{\bm{f}}]}({\bm{\theta}})\|p({\bm{\theta}}|{\mathcal{D}})\Bigg{)}.$ $$

The challenge with this optimization problem lies in our lack of access to the general loss function ${\mathcal{L}}_{\mathcal{D}}({\bm{\theta}})$ and the general posterior distribution $p({\bm{\theta}}|{\mathcal{D}})$. We present our first theorem to address this issue, which characterizes generalization ability in the functional space ${\mathcal{H}}^{d}$. The proof of this theorem can be found in Appendix A.1.

###### Theorem 1 (Informal) .

Let $\tilde{\ell}:{\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{Y}}\to\mathbb
{R}^{+}$ be a loss function on the RKHS ${\mathcal{H}}^{d}$ and the data space. Define $\tilde{L}_{\mathcal{D}}({\bm{f}})=\mathbb{E}_{(x,y)\sim{\mathcal{D}}}[\tilde{
\ell}({\bm{f}},x,y)]$ and $\tilde{L}_{\mathcal{S}}({\bm{f}})=\frac{1}{n}\sum_{i=1}^{n}\tilde{\ell}({\bm{f
}},x_{i},y_{i})$ be the corresponding general and empirical losses. Then for any $\rho>0$ and any distribution ${\mathcal{D}}$, with high probability over the choice of the training set ${\mathcal{S}}\sim{\mathcal{D}}^{n}$, we have:

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle\leq\max_{{\bm{f}}^{\prime}\in{\mathcal{H}}^{d},\|{\bm{f}}^{ \prime}-{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq\rho}\tilde{L}_{\mathcal{S}}({\bm{f} }^{\prime})+h(1/\rho^{2}),$ $$

where $h:\mathbb{R}^{+}\to\mathbb{R}^{+}$ is an increasing function.

This theorem extends prior results, such as the generalization bounds established by Foret et al. () and Kim et al. (), 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 generalization bounds rely on the dimensionality of the domain, while the RKHS is typically infinite-dimensional for many widely used kernels such as the RBF kernels (Aronszajn, ). Building on the first theorem, we present the second theorem, which directly addresses the general posterior and serves as the primary motivation for our method. The proof of this theorem can be found in Appendix A.2.

###### Theorem 2 (Informal) .

Assume that $q$ is any distribution. For any $\rho>0$, with high probability over the training set ${\mathcal{S}}$ generated by distribution ${\mathcal{D}}$, we have:

$$ $D_{\mathrm{KL}}\Big{(}q_{[{\bm{I}}+{\bm{f}}]}||p({\bm{\theta}}|{\mathcal{D}}) \Big{)}\leq\max_{{\bm{f}}^{\prime}\in{\mathcal{H}}^{d},\|{\bm{f}}^{\prime}-{ \bm{f}}\|_{{\mathcal{H}}^{d}}\leq\rho}D_{\mathrm{KL}}\Big{(}q_{[{\bm{I}}+{\bm{ f}}^{\prime}]}||p({\bm{\theta}}|{\mathcal{S}})\Big{)}+h(1/\rho^{2}).$ $$

Our objective is to learn the function ${\bm{f}}^{*}\in{\mathcal{H}}^{d}$ that minimizes $D_{\mathrm{KL}}(q_{[{\bm{I}}+{\bm{f}}]}\|p({\bm{\theta}}|{\mathcal{D}}))$. Motivated by Theorem 2, we propose to implicitly minimize $D_{\mathrm{KL}}(q_{[{\bm{I}}+{\bm{f}}]}||p({\bm{\theta}}|{\mathcal{D}}))$ by minimizing the right-hand side term $\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq\rho}D_{\mathrm{KL
}}\Bigg{(}q_{[{\bm{I}}+{\bm{f}}^{\prime}]}||p({\bm{\theta}}|{\mathcal{S}})
\Bigg{)}$. For any ${\bm{f}}\in{\mathcal{H}}^{d}$, let ${\bm{F}}[{\bm{f}}]=D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{I}}+{\bm{f}}]}\|p({\bm{
\theta}}|{\mathcal{S}})\Bigg{)}$ and ${\bm{f}}^{\prime}={\bm{f}}+\rho\hat{{\bm{f}}}$, it follows that:

$$ $\displaystyle\operatorname*{arg\,max}_{\|{\bm{f}}^{\prime}-{\bm{f}}\|_{{ \mathcal{H}}^{d}}\leq\rho}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{I}}+{\bm{f}}^{\prime }]}||p({\bm{\theta}}|{\mathcal{S}})\Bigg{)}$ $\displaystyle=\operatorname*{arg\,max}_{\|\hat{{\bm{f}}}\|_{{\mathcal{H}}^{d}} \leq 1}{\bm{F}}[{\bm{f}}+\rho\hat{{\bm{f}}}]$ $\displaystyle=\operatorname*{arg\,max}_{\|\hat{{\bm{f}}}\|_{{\mathcal{H}}^{d}} \leq 1}{\bm{F}}[{\bm{f}}]+\rho\left<\hat{{\bm{f}}},\nabla_{\bm{f}}{\bm{F}}[{ \bm{f}}]\right>_{{\mathcal{H}}^{d}}+{\mathcal{O}}(\rho^{2})$ $\displaystyle\approx\operatorname*{arg\,max}_{\|\hat{{\bm{f}}}\|_{{\mathcal{H} }^{d}}\leq 1}\left<\hat{{\bm{f}}},\nabla_{\bm{f}}{\bm{F}}[{\bm{f}}]\right>_{{ \mathcal{H}}^{d}}.$ $$

Let ${\bm{g}}=\nabla_{\bm{f}}{\bm{F}}[{\bm{f}}]\in{\mathcal{H}}^{d}$. The Cauchy-Schwarz inequality on the Hilbert spaces (Kreyszig, ) implies that:

$$ $\Bigg{|}\left<\hat{{\bm{f}}},{\bm{g}}\right>_{{\mathcal{H}}^{d}}\Bigg{|}\leq \left<\hat{{\bm{f}}},{\bm{g}}\right>_{{\mathcal{H}}^{d}}\leq\|\hat{{\bm{f}}}\| _{{\mathcal{H}}^{d}}\|{\bm{g}}\|_{{\mathcal{H}}^{d}}\leq\|{\bm{g}}\|_{{ \mathcal{H}}^{d}}.$ $$

In turn, the solution $\hat{{\bm{f}}}^{*}$ that solves the maximization problem above is given by:

$$ $\hat{{\bm{f}}}^{*}=\frac{{\bm{g}}}{\|{\bm{g}}\|_{{\mathcal{H}}^{d}}}=\frac{ \nabla_{\bm{f}}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{I}}+{\bm{f}}]}\|p(\cdot|{ \mathcal{S}})\Bigg{)}}{\Bigg{\|}\nabla_{\bm{f}}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm {I}}+{\bm{f}}]}\|p(\cdot|{\mathcal{S}})\Bigg{)}\Bigg{\|}_{{\mathcal{H}}^{d}}}.$ (5) $$

Recall that our goal is to find a sequence of functions $\{{\bm{f}}_{k}\}_{k}\subset{\mathcal{H}}^{d}$ that converges toward the optimal solution ${\bm{f}}^{*}$. With the sequence $\{{\bm{f}}_{k}\}_{k}$, we can obtain the flow of distributions $\{q^{(k)}\}_{k}$, in which $q^{(k)}=q_{[{\bm{I}}+{\bm{f}}_{k}]}$, that gradually approaches the optimal solution of Eq. 4. Motivated by Eq. 5, we propose the following functional sharpness-aware update procedure:

$$ $\displaystyle\hat{{\bm{f}}}^{*}_{k}$ $\displaystyle=\rho\frac{\nabla_{{\bm{f}}}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{I}}+{ \bm{f}}]}\|p(\cdot|{\mathcal{S}})\Bigg{)}\Bigg{|}_{{\bm{f}}={\bm{f}}_{k}}}{ \Bigg{\|}\nabla_{{\bm{f}}}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{I}}+{\bm{f}}]}\|p( \cdot|{\mathcal{S}})\Bigg{)}\Bigg{|}_{{\bm{f}}={\bm{f}}_{k}}\Bigg{\|}_{{ \mathcal{H}}^{d}}},$ (Functional Ascend step) (6) $\displaystyle{\bm{f}}_{k+1}$ $\displaystyle={\bm{f}}_{k}-\epsilon\nabla_{{\bm{f}}}D_{\mathrm{KL}}\Bigg{(}q_{ [{\bm{I}}+{\bm{f}}]}\|p(\cdot|{\mathcal{S}})\Bigg{)}\Bigg{|}_{{\bm{f}}={\bm{f} }_{k}+\hat{{\bm{f}}}^{*}_{k}},$ (Functional Descending step) (7) $\displaystyle q^{(k+1)}$ $\displaystyle=q_{[{\bm{I}}+{\bm{f}}_{k+1}]}.$ (Distributional Transformation) (8) $$

To implement this iterative procedure, we must work with the functional gradient terms. For this, we rely on the following lemma, with the proof provided in Appendix B of (Liu & Wang, ):

###### Lemma 1 .

Let ${\bm{F}}[{\bm{f}}]=D_{\mathrm{KL}}(q_{[{\bm{I}}+{\bm{f}}]}\|p(\cdot|{\mathcal{
S}}))$. When $\|{\bm{f}}\|$ is sufficiently small,

$$ $\displaystyle\nabla_{{\bm{f}}}{\bm{F}}[{\bm{f}}]$ $\displaystyle=-\mathbb{E}_{q}[\nabla_{\bm{\theta}}\log p({\bm{\theta}}+{\bm{f} }({\bm{\theta}})|{\mathcal{S}})k({\bm{\theta}},\cdot)+(I+\nabla_{\bm{\theta}}{ \bm{f}}({\bm{\theta}}))^{-1}\nabla_{\bm{\theta}}k({\bm{\theta}},\cdot)]$ (9) $\displaystyle\approx-\mathbb{E}_{q}[\nabla_{\bm{\theta}}\log p({\bm{\theta}}+{ \bm{f}}({\bm{\theta}})|{\mathcal{S}})k({\bm{\theta}},\cdot)+\nabla_{\bm{\theta }}k({\bm{\theta}},\cdot)]\overset{\underset{\mathrm{def}}{}}{=}{\bm{D}}({\bm{f }}).$ (10) $$

Substituting equation (10) into equations (6) and (7), the iterative procedure described from equations (6)-(7) becomes:

$$ $\displaystyle\hat{{\bm{f}}}^{*}_{k}$ $\displaystyle=\rho\frac{{\bm{D}}({\bm{f}}_{k})}{\|{\bm{D}}({\bm{f}}_{k})\|_{{ \mathcal{H}}^{d}}},$ $\displaystyle{\bm{f}}_{k+1}$ $\displaystyle={\bm{f}}_{k}-\epsilon{\bm{D}}({\bm{f}}_{k}+\hat{{\bm{f}}}^{*}_{k }),$ $\displaystyle q^{(k+1)}$ $\displaystyle=q_{[{\bm{I}}+{\bm{f}}_{k+1}]}.$ $$

Even though we do not have access to $p({\bm{\theta}}|{\mathcal{S}})$, we can compute $\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\mathcal{S}})$ because $\nabla_{\bm{\theta}}\log p({\bm{\theta}}|{\mathcal{S}})=\nabla_{\bm{\theta}}
\log p({\bm{\theta}})-\nabla_{\bm{\theta}}{\mathcal{L}}_{\mathcal{S}}({\bm{
\theta}})$. To implement the procedure above, we first draw a set of $m$ particles $\{{\bm{\theta}}_{i}^{(0)}\}^{m}_{i=1}$ on the model space from the initial density, and then iteratively update the particles with an empirical version of ${\bm{D}}({\bm{f}})$. Consequently, we obtain the practical procedure summarized in Algorithm 1, which deterministically transports the set of particles to match the empirical posterior distribution $p({\bm{\theta}}|{\mathcal{S}})$, therefore match the general posterior $p({\bm{\theta}}|{\mathcal{D}})$ as supported by Theorem 2. It is noteworthy that FHBI is a generalization of both SVGD and SAM. In particular, if we set $\rho=0$, we get SVGD; when $m=\textsc{\#particles}=1$, we obtain SAM.

Figure: Algorithm 1 Flat Hilbert Bayesian Inference (FHBI)

Interactive gradient directions and Connections to SAM. To gain further insight into the mechanism of FHBI and its underlying connections to SAM, consider the term $\nabla_{{\bm{\theta}}_{j}}\log({\bm{\theta}}_{j}+\hat{\varepsilon}_{i})$ in the descending step, which is related to $\nabla_{{\bm{\theta}}_{j}}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j}+\hat{
\varepsilon}_{i})$.

The perturbed loss can be approximated as:

$$ ${\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j}+\hat{\varepsilon}_{i})\approx{ \mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j})+\hat{\varepsilon}_{i}\nabla_{{ \bm{\theta}}_{j}}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j}),$ $$

where $\hat{\varepsilon}_{i}$ involves the average $\sum_{k=1}^{m}k({\bm{\theta}}_{k},{\bm{\theta}}_{j})\nabla_{{\bm{\theta}}_{k}}
{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{k})$. Consequently, the gradient of this perturbed loss indicates a direction that simultaneously minimizes $\|\nabla_{{\bm{\theta}}_{j}}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j})\|^{2}$ - which approximates the sharpness of the $j-$th particle, as discussed by Foret et al. () - and $\nabla_{{\bm{\theta}}_{j}}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{j})\cdot
\nabla_{{\bm{\theta}}_{k}}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}_{k})$ for all $j,k$, which reflects the angular similarity 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, promoting more diverse traveling directions for the particles. Besides, as discussed by Liu & Wang (), the second term acts as a spatial repulsive force, driving the particles apart to prevent them from collapsing into a single mode. Consequently, FHBI is not merely an extension of SAM to multiple independent particles; it enables the sharpness and gradient directions of the particles to interact with one another. This insight about the mechanism underlying 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 collapsing into a single mode and improve the final performance as presented in Section 5.

###### Theorem 1 (Informal) .

###### Theorem 2 (Informal) .

###### Lemma 1 .

## 5 Experiments

**Table 1: VTAB-1K classification accuracy results. All the methods are applied to finetune the same set of LoRA parameters on ViT-B/16 pre-trained with ImageNet-21K dataset.**
|  | Natural | Specialized | Structured |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | CIFAR100 | Caltech101 | DTD | Flower102 | Pets | SVHN | Sun397 | Camelyon | EuroSAT | Resisc45 | Retinopathy | Clevr-Count | Clevr-Dist | DMLab | KITTI | dSpr-Loc | dSpr-Ori | sNORB-Azim | sNORB-Ele | 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 | 68.9 | 98.1 | 90.1 | 84.5 | 54.2 | 84.1 | 94.9 | 84.4 | 73.6 | 82.9 | 69.2 | 49.8 | 78.5 | $\bf{75.7}$ | 47.1 | 31.0 | $\bf{44.0}$ | 72.0 |
| SAM | 72.7 | 90.3 | 71.4 | 99.0 | 90.2 | 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 |
| DeepEns | 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 | 82.8 | 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 |
| 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 |
|  | $\mathbf{74.1}$ | $\mathbf{93.0}$ | $\mathbf{74.3}$ | $\mathbf{99.1}$ | $\mathbf{92.4}$ | 87.3 | $\mathbf{56.5}$ | $\mathbf{85.3}$ | 95.0 | $\mathbf{87.2}$ | $\mathbf{79.6}$ | 80.1 | $\bf{72.3}$ | $\bf{52.2}$ | $\bf{80.4}$ | 72.8 | 51.2 | $\bf{31.9}$ | 41.3 | $\bf{73.7}$ |
| FHBI | (.17) | (.42) | (.15) | (0.20) | (0.21) | (.52) | (.12) | (.31) | (.57) | (.21) | (.20) | (.16) | (.27) | (.47) | (.31) | (.50) | (.32) | (.36) | (.59) |  |

#### Applications to Model Fine-tuning.

Bayesian inference methods have promising applications in model finetuning. In standard finetuning scenarios, we are given a pre-trained model $\Phi$. The objective is to find the optimal parameters ${\bm{\theta}}=\Phi+\mathbf{\beta}$, where $\mathbf{\beta}$ represents an additional module, often lightweight and small relative to the full model. Several parameter-efficient finetuning strategies have been developed, including LoRA (Hu et al., ), Adapter (Houlsby et al., ), and others. Our experiments focus on finetuning the ViT-B/16 architecture (Dosovitskiy et al., ), pre-trained with the ImageNet-21K dataset (Deng et al., ), where $\beta$ is defined by the LoRA framework. For the Bayesian approaches, we aim to learn $m$ LoRA particles $\beta^{(i)}$ to obtain $m$ model instances ${\bm{\theta}}^{(i)}$. The final output is then computed as the average of the outputs from all these model instances.

#### Experimental Details.

To assess the effectiveness of FHBI, we conduct experiments on the VTAB-1K benchmark (Zhai et al., ), a diverse and challenging image classification/prediction suite 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 Structured domains. Each dataset includes 1,000 training examples, with an official 80/20 train-validation split. We compared FHBI against seven baselines with three deterministic finetuning strategies including full finetuning, AdamW, and SAM, and four Bayesian inference techniques including Bayesian Deep Ensembles (Lakshminarayanan et al., ), BayesTune (Kim & Hospedales, ), Stochastic Gradient Langevin Dynamics (SGLD) (Welling & Teh, ), and Stein Variational Gradient Descent (SVGD) (Liu & Wang, ).

We used ten warm-up epochs, a batch size 64, 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 $\epsilon$, ascent step size $\rho$, and kernel width $\sigma$. We tuned these hyperparameters using the provided validation set, where the candidate sets are formed as $\epsilon\in\{0.15,1,1.5,2.5\},\rho\in\{0.01,0.03,0.05\},\sigma\in\{0.7,1,1.2\}$. Detailed chosen hyperparameters and data augmentations for each dataset are reported in Appendix B. For each experiment, we conducted five runs of FHBI and reported the mean and standard deviation. All methods except AdamW and full finetuning were trained with four particles on the same set of LoRA parameters.

#### Experimental Results.

To evaluate the effectiveness of our method, we first present the classification accuracy results in Table LABEL:table:_accuracy. FHBI notably improves compared to the baselines, outperforming them in most settings. Compared to other particle sampling methods, including SGLD and SVGD, FHBI consistently performs better across all settings. Moreover, FHBI improves upon SAM by a margin of $3.2\%$, highlighting the advantages of using multiple particles with the underlying interactive gradient directions as previously discussed in Section 4.1. Additionally, as illustrated in Figure 2, FHBI shows the highest performance across all three domains, further solidifying its advantage over the Bayesian inference baselines.

Figure: (a) Natural
Refer to caption: x2.png

To further assess the robustness of FHBI, we evaluate the Expected Calibration Error (ECE) of each setting. This score measures the maximum discrepancy between the model’s accuracy and confidence. As indicated in Table LABEL:table:_ece, 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.

**Table 2: VTAB-1K results evaluated on the Expected Calibration Error (ECE) metric. All methods are applied to finetune the same set of LoRA parameters on ViT-B/16 pre-trained with ImageNet-21K dataset.**
|  | Natural | Specialized | Structured |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | CIFAR100 | Caltech101 | DTD | Flower102 | Pets | SVHN | Sun397 | Camelyon | EuroSAT | Resisc45 | Retinopathy | Clevr-Count | Clevr-Dist | DMLab | KITTI | dSpr-Loc | dSpr-Ori | sNORB-Azi | sNORB-Ele | 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 | $\mathbf{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 | $\bf{0.11}$ | 0.24 | 0.22 |
| 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 | $\bf{0.12}$ |

## 6 Ablation Studies

### 6.1 Effect of #particles

Figure: Figure 3: Runtime by #particles.
Refer to caption: x5.png

To understand the impact of varying the number of particles, we conducted experiments on the seven Natural datasets, reporting both accuracy and per-epoch runtime. We compared FHBI with SVGD and SAM. Figure 3 and Table 3 indicate that multiple particles result in significant performance improvements compared to a single particle. However, while increasing the number of particles enhances performance, it introduces a tradeoff regarding runtime and memory required to store the models. Based on these observations, we found that using $\textsc{\#particles}=4$ provides an optimal balance between performance gains and computational overhead.

**Table 3: Accuracy by #particles.**
| #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 |

### 6.2 Particles Sharpness and Gradient Diversity

Figure: Figure 4: Gradients angular similarities with $m=4$. Lower values indicates greater angular diversity.
Refer to caption: x6.png

As discussed in Section 4.1 and Section 5, FHBI shares implicit connections with SAM by minimizing particle-wise sharpness and diversifying particle travel directions, improving the final performance. To empirically verify this hypothesis about the behavior of our algorithm, we contrast FHBI with SVGD on the KITTI dataset. Four particles are initialized at the same location. We measured: 1) the evolution of sharpness of each particle, defined as $\max_{\|\varepsilon\|\leq\rho}{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}}+
\varepsilon)-{\mathcal{L}}_{\mathcal{S}}({\bm{\theta}})$ according to Foret et al. (), and 2) the evolution of gradients angular diversity, quantified as the Frobenius norm of the covariance matrix formed by the particle gradients. As shown in Figure 5, FHBI not only results in significantly lower and more stable sharpness evolution but also encourages less congruent gradient directions, promoting particles to explore diverse trajectories. Hence, FHBI effectively reduces particle sharpness while promoting angular diversity, improves generalization ability and avoids overfitting by collapsing into a single mode.

Figure: (a) Particle 1
Refer to caption: x7.png

## 7 Conclusion

We introduce Flat Hilbert Bayesian Inference (FHBI), a particle-sampling method designed to enhance generalization ability beyond previous Bayesian inference approaches. This algorithm is based on a theoretical framework that extends generalization principles from Euclidean spaces to the infinite-dimensional RKHS. In our experiments on the VTAB-1K benchmark, FHBI consistently demonstrated performance improvements over six baseline methods by notable margins.

#### Limitations and Future Directions.

Similar to other particle-sampling methods, FHBI needs to store multiple models. Although it remains well-suited for fine-tuning since the additional modules are typically lightweight, this requirement is a memory bottleneck for larger models. Given that the variational inference (VI) approaches can alleviate this issue, an avenue for future research is to extend the concept of sharpness over functional spaces introduced by our theorems to the VI techniques to improve the generalization ability of these methods without storing multiple models.

## Reproducibility Statement

We open-source our implementation and provide the configs and log files at [https://anonymous.4open.science/r/Flat-Hilbert-Variational-Inference-008F/](https://anonymous.4open.science/r/Flat-Hilbert-Variational-Inference-008F/)

## Appendix A Missing proofs

We introduce a few additional notations for the sake of the missing proofs of the main theoretical results. Given a RKHS ${\mathcal{H}}$ equipped with the inner product $\left<\cdot,\cdot\right>_{\mathcal{H}}$ and the norm operator $\|\cdot\|_{\mathcal{H}}$. We define the single-sample loss function on the functional space ${\mathcal{H}}$ to be a map:

$$ $\displaystyle\tilde{\ell}:{\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{ Y}}$ $\displaystyle\to\mathbb{R}$ $\displaystyle({\bm{f}},x,y)$ $\displaystyle\mapsto\tilde{\ell}({\bm{f}},(x,y)).$ $$

Define the general functional loss $\tilde{L}_{\mathcal{D}}({\bm{f}})=\mathbb{E}_{(x,y)\sim{\mathcal{D}}}[\tilde{
\ell}({\bm{f}},(x,y))]$ and the empirical functional loss $\tilde{L}_{\mathcal{S}}({\bm{f}})=\sum_{i=1}^{n}\tilde{\ell}({\bm{f}},(x_{i},y
_{i}))$. Throughout the proof, we assume that the parameter space is bounded by $\|{\bm{\theta}}\|\leq H$, and the data is al bounded that $\|x\|\leq R,y\leq R$ for some $R,H\in\mathbb{R}$.

We introduce the following lemmas that will be used throughout the proof of our main theorems.

###### Lemma 2 (Approximation of RKHS functionals) .

Let $d\in\mathbb{N},{\mathcal{X}}=[-R,R]^{K}$ for some $K\in\mathbb{R}$. Consider ${\mathcal{K}}=\{f\in{\mathcal{H}}:\|f\|_{{\mathcal{H}}}\leq 1\}$ with ${\mathcal{H}}$ induced by some Mercer kernel which is $\alpha$-Holder continuous for $\alpha\in(0,1)$ with constant $C_{K}\geq 0$. Suppose $F$ is $s-$Holder continuous for $s\in(0,1]$ with constant $C_{f}\geq 0$. There exists some $M_{0}\in\mathbb{N}$ such that for every $M\in\mathbb{N}$ with $M>M_{0}$, by taking some fixed $\overline{t}=\{t_{i}\}$ with $N\in\mathbb{N}$, we have a tanh neural network $\hat{G}$ with two hidden layers of widths at most $N(M-1)$ and $3\frac{N+1}{2}(5M)^{N}$ parameters satisfying

$$ $\sup_{f\in{\mathcal{K}}}|F(f)-\hat{G}(f(\overline{t}))|\leq RC_{F}(\epsilon_{K }(\overline{t}))^{s}+\frac{7N^{2}RC_{G}}{M},$ (11) $$

with

$$ $C_{G}=C_{F}(1+\|K[\overline{t}]^{-1}\|_{op}\sqrt{N}C_{K}(h_{\overline{t}})^{ \alpha})^{s},$ $$

where $K[\overline{t}]$ is the Gram matrix of $\overline{t}$.

###### Proof.

The proof can be found in Zhou et al. ()
∎

###### Lemma 3 (Product of RKHSs) .

Given $n$ RKHSs ${\mathcal{H}}_{1},{\mathcal{H}}_{2},\cdots,{\mathcal{H}}_{n}$, each defined on corresponding sets $X_{1},X_{2},\cdots,X_{n}$ with kernels $k_{1}(x_{1},y_{1}),\cdots k_{n}(x_{n},y_{n})$ respectively. Then, ${\mathcal{H}}=\bigotimes_{i=1}^{n}={\mathcal{H}}_{1}\times{\mathcal{H}}_{2}
\times\cdots\times{\mathcal{H}}_{n}$ is also an RKHS, with kernel $K$ that is the product of the individual kernels.

###### Proof.

The product space ${\mathcal{H}}=\bigotimes_{i=1}^{n}{\mathcal{H}}_{i}$ consists of tuples of functions $(f_{1},f_{2},\cdots,f_{n})$. Firstly, we define the inner product in ${\mathcal{H}}$ as:

$$ $\left<(f_{1},f_{2},\cdots,f_{n}),(g_{1},g_{2},\cdots,g_{n})\right>_{\mathcal{H }}=\sum^{n}_{i=1}\left<f_{i},g_{i}\right>_{{\mathcal{H}}_{i}}.$ $$

This definition naturally defines a Hilbert space structure on ${\mathcal{H}}$ since each ${\mathcal{H}}_{i}$ is a Hilber space, and the sum of inner products is linear and positive definite. Now we define the kernel for the product space:

$$ $k((x_{1},x_{2},\cdots,x_{n}),(y_{1},y_{2},\cdots,y_{n}))=\prod_{i=1}^{n}k(x_{i },y_{i}).$ $$

Notice that the pointwise product of positive definite kernels is a positive definite kernel, hence this kernel is valid.

We now verify the reproducing property of ${\mathcal{H}}$. Consider a function $f=(f_{1},f_{2},\cdots,f_{n})\in{\mathcal{H}}$, and evaluate the function at a point $(x_{1},x_{2},\cdots,x_{n})\in\bigotimes_{i=1}^{n}X_{i}$.

The reproducing property in each individual RKHS ${\mathcal{H}}_{i}$ implies that:

$$ $f_{i}(x_{i})=\left<f_{i},k_{i}(x_{i},\cdot)\right>_{{\mathcal{H}}_{i}}.$ $$

Hence, for the function $f=(f_{1},\cdots,f_{n})$, we get:

$$ $\displaystyle f((x_{1},x_{2},\cdots,x_{n}))$ $\displaystyle=(f_{1}(x_{1}),f_{2}(x_{2}),\cdots,f_{n}(x_{n}))$ $\displaystyle=(\left<f_{1},k_{1}(x_{1},\cdot)\right>_{{\mathcal{H}}_{1}},\left <f_{2},k_{2}(x_{2},\cdot)\right>_{{\mathcal{H}}_{2}},\cdots,\left<f_{n},k_{n}( x_{n},\cdot)\right>_{{\mathcal{H}}_{n}})$ $\displaystyle=\left<(f_{1},f_{2},\cdots,f_{n}),(k_{1}(x_{1},\cdot),k_{2}(x_{2} ,\cdot),\cdots k_{n}(x_{n},\cdot))\right>_{\mathcal{H}}.$ $$

Thus, the reproducing property holds for the product space ${\mathcal{H}}$. Since ${\mathcal{H}}$ is a Hilbert space and the kernel $k$ satisfies the reproducing property, we conclude that ${\mathcal{H}}=\bigotimes^{n}_{i=1}{\mathcal{H}}_{i}$ is another RKHS.
∎

###### Lemma 2 (Approximation of RKHS functionals) .

###### Proof.

###### Lemma 3 (Product of RKHSs) .

###### Proof.

### A.1 Proof of Theorem 1

###### Theorem 3 .

For any $\rho>0$ and any distribution ${\mathcal{D}}$, with probability $1-\rho$ over the choice of the training set ${\mathcal{S}}\sim{\mathcal{D}}^{n}$,

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq \rho}\tilde{L}_{\mathcal{S}}({\bm{f}}^{\prime})+RC_{F}(\epsilon_{K}(\overline{ t}))^{s}+\frac{7N^{2}RC_{G}}{M}$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}.$ $$

###### Proof.

$\tilde{\ell}$ is a functional that maps from ${\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{Y}}$ to $\mathbb{R}$. Notice that ${\mathcal{H}}^{d}$ is a RKHS, ${\mathcal{X}}=\mathbb{R}^{a}$ and ${\mathcal{Y}}=\mathbb{R}^{b}$ for some $a,b\in\mathbb{Z}$ are Euclidean spaces, which are also instances of RKHS. Moreover, the product of RKHS’s is also a RKHS according to Lemma 3. Hence, ${\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{Y}}$ is also a RKHS. According to Lemma 2, there exists $N$ points $\overline{{\bm{\theta}}}=\{{\bm{\theta}}_{i}\}^{N}_{i=1}\subset\Theta$, and a two-layer neural network $G_{\bm{W}}$ parameterized by ${\bm{W}}$ so that

$$ $|\tilde{\ell}({\bm{f}},x,y)-G_{\bm{W}}({\bm{f}}(\overline{{\bm{\theta}}}),x,y) |\leq RC_{F}(\epsilon_{K}(\overline{t}))^{s}+\frac{7N^{2}RC_{G}}{M},$ $$

for every $({\bm{f}},x,y)\in{\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{Y}}$. Consider ${\bm{f}}^{\prime}\in{\mathcal{H}}^{d}$ so that $\|{\bm{f}}^{\prime}-{\bm{f}}\|\leq\rho$, it implies $|{\bm{f}}(\overline{{\bm{\theta}}})-{\bm{f}}^{\prime}(\overline{{\bm{\theta}}}
)|\leq P\|{\bm{f}}-{\bm{f}}^{\prime}\|_{{\mathcal{H}}^{d}}\leq P\rho$. Denote $\tilde{{\bm{\theta}}}={\bm{f}}(\overline{{\bm{\theta}}})\in\mathbb{R}^{N^{
\prime}}$ for some $N^{\prime}\in\mathbb{Z}$, by invoking the inequality from Foret et al. (), let $\rho^{\prime}=\rho P$, it follows that:

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle=\mathbb{E}_{(x,y)\sim{\mathcal{D}}}[\tilde{\ell}({\bm{f}},x,y)] \leq\mathbb{E}_{(x,y)\sim{\mathcal{D}}}[G_{\bm{W}}({\bm{f}}(\overline{{\bm{ \theta}}}),x,y)]+RC_{F}(\epsilon_{K}(\overline{t}))^{s}+\frac{7N^{2}RC_{G}}{M}$ $\displaystyle=\mathbb{E}_{(x,y)\sim{\mathcal{D}}}[G_{\bm{W}}(\tilde{{\bm{ \theta}}},x,y)]+RC_{F}(\epsilon_{K}(\overline{t}))^{s}+\frac{7N^{2}RC_{G}}{M}$ $\displaystyle\leq\max_{\|\tilde{{\bm{\theta}}}^{\prime}-\tilde{{\bm{\theta}}} \|^{2}_{2}\leq\rho^{\prime}}\frac{1}{n}\sum_{i=1}^{n}G_{\bm{W}}(\tilde{{\bm{ \theta}}}^{\prime},x,y)+h(M,N)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{\prime 2}} \Bigg{(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log \frac{n}{\rho^{\prime}}+8\log(6n+3k)}.$ $$

By definition, a RKHS is a closed Hilbert space. Then, there exists a sequence $\{{\bm{f}}^{\prime}_{n}\}$ so that ${\bm{f}}^{\prime}_{n}(\overline{{\bm{\theta}}})$ that gets arbitrarily close to $\tilde{{\bm{\theta}}}^{\prime}$. Then, for any $\epsilon>0$, it follows:

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle\leq\max_{\|\tilde{{\bm{\theta}}}^{\prime}-\tilde{{\bm{\theta}}} \|^{2}_{2}\leq\rho^{\prime}}\frac{1}{n}\sum_{i=1}^{n}G_{\bm{W}}(\tilde{{\bm{ \theta}}}^{\prime},x,y)+h(M,N)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{\prime 2}} \Bigg{(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log \frac{n}{\rho^{\prime}}+8\log(6n+3k)}$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}(\overline{{\bm{\theta}}})-{\bm{f}}( \overline{{\bm{\theta}}})\|^{2}_{2}\leq\rho P}\frac{1}{n}\sum_{i=1}^{n}G_{\bm{ W}}({\bm{f}}^{\prime}(\overline{{\bm{\theta}}}),x,y)+h(M,N)+\epsilon{\mathcal{ O}}(1)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|^{2}_{2}\leq\rho}\frac{1} {n}\sum_{i=1}^{n}G_{\bm{W}}({\bm{f}}^{\prime}({\bm{\theta}}),x,y)+h(M,N)+ \epsilon{\mathcal{O}}(1)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|^{2}_{2}\leq\rho}\tilde{L }_{\mathcal{S}}({\bm{f}}^{\prime})+h(M,N)+\epsilon{\mathcal{O}}(1)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}.$ $$

This is true for any $\epsilon>0$. Hence, it implies

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|^{2}_{2}\leq\rho}\tilde{L }_{\mathcal{S}}({\bm{f}}^{\prime})+RC_{F}(\epsilon_{K}(\overline{t}))^{s}+ \frac{7N^{2}RC_{G}}{M}$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)},$ $$

which concludes our proof.
∎

###### Theorem 3 .

###### Proof.

### A.2 Proof of Theorem 2

Now we can prove the Theorem 2. We restate the theorem

###### Theorem 4 .

For any target distribution $p$, reference distribution $q$, and any $\rho>0$, we have the following bound between the general KL loss and the empirical KL loss

$$ $\displaystyle D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{f}}]}||p\Big{(}{\bm{\theta}}|{ \mathcal{D}}\Big{)}\Bigg{)}$ $\displaystyle\leq\max_{{\bm{f}}^{\prime}\in{\mathcal{B}}_{\rho}({\bm{f}})}D_{ \mathrm{KL}}\Bigg{(}q_{[{\bm{f}}^{\prime}]}||p\Big{(}{\bm{\theta}}|{\mathcal{S }}\Big{)}\Bigg{)}+RC_{F}(\epsilon_{K}(\overline{t}))^{s}+\frac{7N^{2}RC_{G}}{M}$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}.$ $$

###### Proof.

Consider the left-hand side, we have:

$$ $\displaystyle D_{\mathrm{KL}}(q_{[{\bm{f}}]}\|p({\bm{\theta}}|{\mathcal{D}}))$ $\displaystyle=\int q_{[{\bm{f}}]}({\bm{\theta}})\Bigg{(}{\mathcal{L}}_{ \mathcal{D}}({\bm{\theta}})+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{p({\bm{ \theta}})}+\log Z_{\mathcal{D}}\Bigg{)}d{\bm{\theta}}$ $\displaystyle=\int q_{[{\bm{f}}]}({\bm{\theta}})\Bigg{(}\mathbb{E}_{(x,y)\sim{ \mathcal{D}}}\ell({\bm{\theta}},x,y)+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{ p({\bm{\theta}})}+\log Z_{\mathcal{D}}\Bigg{)}d{\bm{\theta}}$ $\displaystyle=\mathbb{E}_{(x,y)\sim{\mathcal{D}}}\Bigg{[}\int q_{[{\bm{f}}]}({ \bm{\theta}})\Big{(}\ell(\theta;x,y)+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{ p({\bm{\theta}})}\Big{)}d{\bm{\theta}}\Bigg{]}+\int q_{[{\bm{f}}]}({\bm{\theta }})\log Z_{\mathcal{D}}d{\bm{\theta}}.$ $$

On the other hand, we also have:

$$ $\displaystyle D_{\mathrm{KL}}(q_{[{\bm{f}}]}\|p({\bm{\theta}}|{\mathcal{S}}))$ $\displaystyle=\int q_{[{\bm{f}}]}\Bigg{(}({\bm{\theta}}){\mathcal{L}}_{ \mathcal{S}}({\bm{\theta}})+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{p({\bm{ \theta}})}+\log Z_{\mathcal{S}}\Bigg{)}d{\bm{\theta}}$ $\displaystyle=\int q_{[{\bm{f}}]}({\bm{\theta}})\Bigg{(}\frac{1}{n}\sum_{i=1}^ {n}\ell({\bm{\theta}},x_{i},y_{i})+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{p( {\bm{\theta}})}+\log Z_{\mathcal{S}}\Bigg{)}d{\bm{\theta}}$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\Bigg{[}\int q_{[{\bm{f}}]}({\bm{\theta }})\Big{(}\ell(\theta;x_{i},y_{i})+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{p( {\bm{\theta}})}+\log Z_{\mathcal{S}}\Big{)}d{\bm{\theta}}\Bigg{]}+\int q_{[{ \bm{f}}]}({\bm{\theta}})\log Z_{\mathcal{S}}d{\bm{\theta}}.$ $$

We define $\tilde{L}$ to be the functional such that:

$$ $\displaystyle\tilde{L}:{\mathcal{H}}^{d}\times{\mathcal{X}}\times{\mathcal{Y}}$ $\displaystyle\to\mathbb{R}$ $\displaystyle({\bm{f}},x,y)$ $\displaystyle\mapsto\tilde{L}({\bm{f}},x,y)=\int q_{[{\bm{f}}]}({\bm{\theta}}) \Big{(}\ell(\theta;x,y)+\log\frac{q_{[{\bm{f}}]}({\bm{\theta}})}{p({\bm{\theta }})}\Big{)}d{\bm{\theta}}.$ $$

According to Theorem 1, we have:

$$ $\displaystyle\tilde{L}_{\mathcal{D}}({\bm{f}})$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq \rho}\tilde{L}_{\mathcal{S}}({\bm{f}}^{\prime})RC_{F}(\epsilon_{K}(\overline{t }))^{s}+\frac{7N^{2}RC_{G}}{M}$ (12) $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}.$ (13) $$

Moreover, the model and data spaces are bounded, so ${\mathcal{L}}_{\mathcal{D}}({\bm{\theta}})$ and ${\mathcal{L}}_{\mathcal{S}}({\bm{\theta}})$ are bounded. Then, there exists constants $d,D$ such that $d\leq\log Z_{\mathcal{S}},\log Z_{\mathcal{D}}\leq D$, which also implies $d\leq\int q_{[{\bm{f}}]}\log Z_{\mathcal{S}}d{\bm{\theta}}\leq D$ and $d\leq\int q_{[{\bm{f}}]}\log Z_{\mathcal{D}}d{\bm{\theta}}\leq D$. It follows that for all ${\bm{f}},{\bm{f}}^{\prime}\in{\mathcal{H}}^{d}$:

$$ $\int q_{[{\bm{f}}]}({\bm{\theta}})\log Z_{\mathcal{D}}d{\bm{\theta}}\leq\int q _{[{\bm{f}}^{\prime}]}({\bm{\theta}})\log Z_{\mathcal{S}}d{\bm{\theta}}+D-d.$ (14) $$

Combining the Inequalities 13 and 14, it follows that:

$$ $\displaystyle D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{f}}]}||p\Big{(}{\bm{\theta}}|{ \mathcal{D}}\Big{)}\Bigg{)}$ $\displaystyle\leq\max_{\|{\bm{f}}^{\prime}-{\bm{f}}\|_{{\mathcal{H}}^{d}}\leq \rho}D_{\mathrm{KL}}\Bigg{(}q_{[{\bm{f}}^{\prime}]}||p\Big{(}{\bm{\theta}}|{ \mathcal{S}}\Big{)}\Bigg{)}+RC_{F}(\epsilon_{K}(\overline{t}))^{s}+\frac{7N^{2 }RC_{G}}{M}+(D-d)$ $\displaystyle+\sqrt{\frac{N^{\prime}\log\Bigg{(}1+\frac{C}{\rho^{2}P^{2}}\Bigg {(}1+\sqrt{\frac{\log(N)}{N^{\prime}}}\Bigg{)}^{2}\Bigg{)}}{n-1}+4\log\frac{n} {\rho P}+8\log(6n+3k)}.$ $$

which concludes our proof.
∎

###### Theorem 4 .

###### Proof.

## Appendix B Experimental Details

### B.1 Chosen Hyperparameters

We grid-search hyperparameters on the validation set, where the key hyperparameters are: the kernel width $\sigma$, the initial learning rate $\epsilon$, and the ascend step size $\rho$. The candidate sets are formed as $\epsilon\in\{0.15,1,1.5,2.1,2.5\},\rho\in\{0.01,0.03,0.05\},\sigma\in\{0.7,1,1
.2\}$. The chosen hyperparameters are as follows $(\epsilon,\rho,\sigma)$: CIFAR100 = $(0.15,0.03,1.2)$, Caltech101 = $(2.1,0.05,1.2)$, DTD = $(0.15,0.03,1.2)$, Flowers102 = $(0.15,0.03,1)$, Pets = $(0.15,0.03,1.2)$, SVHN = $(2.5,0.01,1)$, Sun397 = $(0.15,0.03,1.2)$, Patch-Camelyon = $(2.1,0.05,1)$, DMLab = $(2.1,0.03,1)$, EuroSAT = $(2.5,0.01,1.2)$, Resisc45 = $(1.5,0.03,1.2)$, Diabetic- Retinopathy = $(2.1,0.03,1)$, Clevr-Count = $(2.5,0.01,1)$, Clevr-Dist = $(1,0.01,1.2)$, KITTI = $(2.1,0.05,1)$, dSprites-loc = $(2.1,0.05,1)$, dSprites-ori = $(2.1,0.03,1.2)$, smallNorb-azi = $(1,0.05,1)$, smallNorb-ele = $(1,0.03,0.7)$.

### B.2 Data Augmentations

Our implementation is based on the repository [V-PETL](https://github.com/synbol/Parameter-Efficient-Transfer-Learning-Benchmark/tree/main). Similar to this repository, we use a different data augmentation among the following three augmentations for each dataset. In particular, the data augmentations that we used for each setting are:

- •
For CIFAR100, DTD, Flower102, Pets, Sun397
[⬇](data:text/plain;base64,ICAgICBzZWxmLnRyYW5zZm9ybXNfdHJhaW4gPSB0cmFuc2Zvcm1zLkNvbXBvc2UoCiAgICAgICAgWwogICAgICAgICAgICB0cmFuc2Zvcm1zLlJhbmRvbVJlc2l6ZWRDcm9wKAogICAgICAgICAgICAgICAgKHNlbGYuc2l6ZSwgc2VsZi5zaXplKSwKICAgICAgICAgICAgICAgIHNjYWxlPShzZWxmLm1pbl9zY2FsZSwgc2VsZi5tYXhfc2NhbGUpLAogICAgICAgICAgICApLAogICAgICAgICAgICB0cmFuc2Zvcm1zLlJhbmRvbUhvcml6b250YWxGbGlwKHNlbGYuZmxpcF9wcm9iKSwKICAgICAgICAgICAgdHJhbnNmb3Jtcy5Ucml2aWFsQXVnbWVudFdpZGUoKQogICAgICAgICAgICBpZiBzZWxmLnVzZV90cml2aWFsX2F1ZwogICAgICAgICAgICBlbHNlIHRyYW5zZm9ybXMuUmFuZEF1Z21lbnQoc2VsZi5yYW5kX2F1Z19uLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgc2VsZi5yYW5kX2F1Z19tKSwKICAgICAgICAgICAgdHJhbnNmb3Jtcy5Ub1RlbnNvcigpLAogICAgICAgICAgICB0cmFuc2Zvcm1zLk5vcm1hbGl6ZShtZWFuPVswLjQ4NSwgMC40NTYsIDAuNDA2XSwKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBzdGQ9WzAuMjI5LCAwLjIyNCwgMC4yMjVdKV0pLAogICAgICAgICAgICB0cmFuc2Zvcm1zLlJhbmRvbUVyYXNpbmcocD1zZWxmLmVyYXNlX3Byb2IpLAogICAgICAgIF0KICAgICkKICAgIHNlbGYudHJhbnNmb3Jtc190ZXN0ID0gdHJhbnNmb3Jtcy5Db21wb3NlKAogICAgICAgIFsKICAgICAgICAgICAgdHJhbnNmb3Jtcy5SZXNpemUoCiAgICAgICAgICAgICAgICAoc2VsZi5zaXplLCBzZWxmLnNpemUpLAogICAgICAgICAgICApLAogICAgICAgICAgICB0cmFuc2Zvcm1zLlRvVGVuc29yKCksCiAgICAgICAgICAgIHRyYW5zZm9ybXMuTm9ybWFsaXplKG1lYW49WzAuNDg1LCAwLjQ1NiwgMC40MDZdLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHN0ZD1bMC4yMjksIDAuMjI0LCAwLjIyNV0pXSksCiAgICAgICAgXQogICAgKQ==)
self.transforms_train = transforms.Compose(
[
transforms.RandomResizedCrop(
(self.size, self.size),
scale=(self.min_scale, self.max_scale),
),
transforms.RandomHorizontalFlip(self.flip_prob),
transforms.TrivialAugmentWide()
if self.use_trivial_aug
else transforms.RandAugment(self.rand_aug_n,
self.rand_aug_m),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])]),
transforms.RandomErasing(p=self.erase_prob),
]
)
self.transforms_test = transforms.Compose(
[
transforms.Resize(
(self.size, self.size),
),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])]),
]
)
- •
For Caltech101, Clevr-Dist, Dsprites-Loc, Dsprites-Ori, SmallNorb-Azi, SmallNorb-Ele:
[⬇](data:text/plain;base64,ICAgIHNlbGYudHJhbnNmb3JtX3RyYWluID0gdHJhbnNmb3Jtcy5Db21wb3NlKFsKICAgICAgICB0cmFuc2Zvcm1zLlJlc2l6ZSgoMjI0LCAyMjQpKSwKICAgICAgICB0cmFuc2Zvcm1zLlRvVGVuc29yKCksCiAgICAgICAgdHJhbnNmb3Jtcy5Ob3JtYWxpemUobWVhbj1bMC40ODUsIDAuNDU2LCAwLjQwNl0sCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgc3RkPVswLjIyOSwgMC4yMjQsIDAuMjI1XSldKQogICAgc2VsZi50cmFuc2Zvcm1fdGVzdCA9IHRyYW5zZm9ybXMuQ29tcG9zZShbCiAgICAgICAgdHJhbnNmb3Jtcy5SZXNpemUoKDIyNCwgMjI0KSksCiAgICAgICAgdHJhbnNmb3Jtcy5Ub1RlbnNvcigpLAogICAgICAgIHRyYW5zZm9ybXMuTm9ybWFsaXplKG1lYW49WzAuNDg1LCAwLjQ1NiwgMC40MDZdLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgIHN0ZD1bMC4yMjksIDAuMjI0LCAwLjIyNV0pXSk=)
self.transform_train = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
self.transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(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:
[⬇](data:text/plain;base64,ICAgIGZyb20gdGltbS5kYXRhIGltcG9ydCBjcmVhdGVfdHJhbnNmb3JtCiAgICBzZWxmLnRyYW5zZm9ybV90cmFpbiA9IGNyZWF0ZV90cmFuc2Zvcm0oCiAgICAgICAgICAgICAgICBpbnB1dF9zaXplPSgyMjQsIDIyNCksCiAgICAgICAgICAgICAgICBpc190cmFpbmluZz1UcnVlLAogICAgICAgICAgICAgICAgY29sb3Jfaml0dGVyPTAuNCwKICAgICAgICAgICAgICAgIGF1dG9fYXVnbWVudD0ncmFuZC1tOS1tc3RkMC41LWluYzEnLAogICAgICAgICAgICAgICAgcmVfcHJvYj0wLjAsCiAgICAgICAgICAgICAgICByZV9tb2RlPSdwaXhlbCcsCiAgICAgICAgICAgICAgICByZV9jb3VudD0xLAogICAgICAgICAgICAgICAgaW50ZXJwb2xhdGlvbj0nYmljdWJpYycsCiAgICAgICAgICAgICkKICAgIGF1Z190cmFuc2Zvcm0udHJhbnNmb3Jtc1swXSA9IHRyYW5zZm9ybXMuUmVzaXplKCgyMjQsIDIyNCksCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICBpbnRlcnBvbGF0aW9uPTMpCiAgICBzZWxmLnRyYW5zZm9ybV90ZXN0ID0gdHJhbnNmb3Jtcy5Db21wb3NlKFsKICAgICAgICAgICAgdHJhbnNmb3Jtcy5SZXNpemUoKDIyNCwgMjI0KSksCiAgICAgICAgICAgIHRyYW5zZm9ybXMuVG9UZW5zb3IoKSwKICAgICAgICAgICAgdHJhbnNmb3Jtcy5Ob3JtYWxpemUobWVhbj1bMC40ODUsIDAuNDU2LCAwLjQwNl0sCiAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgc3RkPVswLjIyOSwgMC4yMjQsIDAuMjI1XSldKQ==)
from timm.data import create_transform
self.transform_train = create_transform(
input_size=(224, 224),
is_training=True,
color_jitter=0.4,
auto_augment=’rand-m9-mstd0.5-inc1’,
re_prob=0.0,
re_mode=’pixel’,
re_count=1,
interpolation=’bicubic’,
)
aug_transform.transforms[0] = transforms.Resize((224, 224),
interpolation=3)
self.transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - The method is tested only in fine-tuning regime. Is there any reason for that? 
Testing the proposed models trained from scratch would strengthen the empirical significance of the proposed method. I am not sure which datasets would be suitable, but datasets like ImageNet and its variants (ImageNet-A [1], and ImageNet-C[2]), and scientific machine learning benchmarks might be good datasets as uncertainty prediction would be critical. 

[1] https://arxiv.org/abs/1907.07174
[2] https://arxiv.org/abs/1903.12261
soundness: 3
strengths: - The work extends generalization bounds from finite-dimensional parameter spaces to functional spaces, which leads to FHBI, a theoretically grounded Bayesian inference algorithm.
- The performance improvement by the proposed method is validated through extensive comparisons with previous works.
summary: This work presents an algorithm called Flat Hilbert Bayesian Inference (FHBI), which incorporates Sharpness-aware minimization (SAM) technique in Bayesian inference. Specifically, together with SAM, the authors perform Stein variational gradient descent (SVGD) as the dynamics of the model parameters, which leads the particles (models) to flat and diverse modes. FHBI is tested on VTAB-1K, a collection of various classification tasks, and achieves better performance on average among different Bayesian NN methods.
weaknesses: - I don't see any particular weaknesses in this paper.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Have you compared SVGD+SAM to SGLD+SAM [1, 2]? 


[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] 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.
soundness: 2
strengths: 1. Incorporating SAM into SVGD can somehow improve the stability. The idea generally makes sense. 
2. The paper is clearly written overall, and it is easy to capture the motivation.
3. The empirical results are superior and tested on various datasets.
summary: This paper combines sharpness aware minimization (SAM) with SVGD in RKHS as the Flat Hilbert Bayesian Inference (FHBI). It extends the proof of [1] to infinite-dimensional functional space.
Empirical validations show better performance of FHBI compared to previous Bayesian inference approaches on various datasets.

[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.
weaknesses: Even though I agree that the proposed method may be effective in some conditions, I am not convinced by the theory in this paper. Following are my concerns. 
1. Notations.  
    * Empirical posterior should be $\mathbb{P}\_{\theta|S}$ not $\mathbb{P}\_S$, which is the data distribution. And the prior distribution should be $\mathbb{P}_{\theta}$.
    * $p(\theta|S)\propto p(\theta)p(S|\theta) = p(\theta)\prod_{i=1}^n p(y_i, x_i|\theta)$
     The exponential average loss should be related to a specific distribution. I do not know how you can obtain this form directly. 
     * What is "general loss"? I have never heard this terminology. $\mathcal{L}_{\mathcal{D}}(\theta)$ is often called population loss/ true error/ generalization error in different papers or learning theory books. 
2. For Theorem 2, what is the exact definition of $h(1/\rho^2)$? It should be clearly presented in the main paper. In your derivation, $\mathcal{O}(\rho^2)$ is simply ignored, which means $\rho$ should be very small, which in turn gives a large $h(1/\rho^2)$.
3. You are actually deriving an upper bound of the true risk, which is the empirical risk plus some complexity term. However, the sample complexity is not directly reflected in the bounds presented in the main paper. You claim that you can approximate the true posterior $p(\theta|\mathcal{D})$. This is impossible if you only have limited samples $n$. 

4. Experiments
	* How does data augmentation affect the empirical results? Have you used it in all baselines or just your method? 
	* As shown in Figure 3, the runtime of the FHBI increases fast w.r.t the number of particles. Compared to SVGD, the margin also grows with the number of particles. Where is the computation bottleneck? Can you find any way to reduce it? 

### Minor: 
1. Is the template used correctly? There is no line number. 
2. Typo. In the related work section, sharness -> sharpness.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - In my understanding, Theorem1 expresses the tradeoff relation between (1) worst-case empirical loss  that scales with the size of the neighborhood (radius of the step size)   (2) wellness of the upperbound approximation based on the empirical loss.  The reviewer is making in his mind a super-rough analogy of "local linear approximation of a complex function", whose worst case error scales with the size of the neighborhood, and "functional ascending step" is analogous to choosing the direction of worst error in this approximation. 

If this analogy is correct,  the neighborhood radius $\rho$ as well as the update size $\epsilon$ would be roughly analogous to a stepsize in Euler-Maruyama simulation of the ODE, except that in the context of FBHI the system to be simulated is infinite dimensional.  Would the method improve its performance by choosing $\rho$ and $\epsilon$ to be small, and running the iterations for greater number of times, as in ODE simulation?    Also, with this intuition it feels as if $\rho$ shall be similar to $\epsilon$;  why are they chosen quite differently in the experimental section? 

While the reviewer is not too confident of the intuition based on this super-rough analogy, the reviewer also believes that it is crucial that the paper provides some explanation (either numerical or experimental, or at the very least, the heuristic with intuitive explanation) regarding the choice of $\rho$ and $\epsilon$, both for the sake of the future practical user of FBHI and the for the sake of the future schemes that will possibly branch out from FBHI.

- In the similar note as above, because the "wellness" of RKHS method generally depends on the affinity between the choice of the kernel (and the hence the nature of the continuity of functions in the space) and the dataset.   In the case of  this research,  the choice, I believe, is indirectly related to dataset because RKHS is a space of functions on "parameters".  Is there any ablation study regarding this choice, or at the very least, a good heuristics that would help the user to choose appropriate Kernel in the applications? 

- It is explained in the paper that $\rho=0$ would correspond to  SVGD, and "in equation" it seems so. However, because the paper emphasizes its connection to SVGD, the reviewer wishes to see why this comes about in connection to Theorem 1 and the derivation of SVGD. 

- Is there comparisons against SAM in the setting of section 6.2?
soundness: 3
strengths: - This research presents a scheme that generalizes both SVGD and SAM in the context of Bayesian Inference. 
The idea of FBHI itself is very clearly stated and convincing, and its efficacy are validated with ample experiments. 
The layout of the derivation is very instructive as well. 
The claimed sharpness advantage (that FBHI would is more sharpness-aware) is also validated in experiments, empirically  
showcasing the mechanism behind the advantage of the method.
summary: This paper presents Flat Hilbert Bayesian Inference(FHBI),   a method to compute a posterior inference by integrating along the flow of distributions that converges to the target posterior.   By formulating  the derivative of the flow with the pushforward function that lives in a vector valued RKHS,  the framework naturally establishes the "flow of infinite number of particles" which allows one to target the "general posterior" as opposed to "empirical" posterior in the setting of Bayesian Inference.   Importantly, the method differs from Stein Variational Gradient Descent in that it derives the infinitesimal pushforward function that reduces the "worst upper bound" of the KL divergence in a way that resembles adversarial training.

The propsoed  paradigm is made implementable by the computable form of $\lambda p(\theta | S)$, together with the upper bound of the KL divergence  to general posterior defined with the KL divergence to empirical posterior.  
The Efficacy of the algorithm is verified through experiments and its scaling properties are investigated and compared with respect to SVGD.
weaknesses: - While the relation of FHBI to SAM and the interpretation with spatial and angular repulsive force is insightful,  the reviewer was a little confused in the introduction and in the multiple reference to SAM that sounded as if SAM would be the major idea on which to develop FHBI (which, in retrospect, is not "directly" so?).  Meanwhile, it is  true in the algorithm that when m=1, the algorithm agrees with SAM in the end.  In the current presentation the connection to SAM seems something a posteriori.

If indeed the sharpness-aware philosophy is indeed the "motivation" of FHBI (which is, unfortunately, not yet clearly conveyed to this reviewer), the reviewer would like to see more analytical connection to SAM's derivation.  

- In the similar note, as the reviewer will post in the "Questions" section, the reviewer feels that the supposedly the important connections to SVGD and SAM are not clearly explained through objective function and derivations (**more than $\rho=0$ and $m=1$ )

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Major:
* Something seems amiss in going from the empirical posterior $p(\theta|\mathcal S)$ to what you call the general posterior $p(\theta|\mathcal D)$. Let $L_n(\theta) = \frac{1}{n} \sum_{i=1}^n -\log p(y_i|x_i,\theta)$. Then the empirical posterior is $p(\theta|\mathcal S) = \exp(-n L_n(\theta)) p(\theta)$. The population counterpart to this is $\exp(-n L(\theta)) p(\theta)$. But this is *not* your general posterior which no longer has any dependence on sample size $n$. 

Minor:
* Currently, the way the loss $\ell$ shows up in Section 3 might give the wrong impression that it is something one can freely choose. In fact the loss has to be (proportional) to the negative log likelihood. I think you need to write out explicitly what $\ell(f_\theta(x),y)$ is. 
* The sentence above Eqn (5), "In turn, the solution $\hat f^*$ that solves the maximization problem above is given by". Which maximization problem above? There are so many approximations here, please use \label and \ref to refer to exact maximization problem. (I also doubt this is a true statement, (5) is not really a solution but an approximation of a solution right?)
* Would you consider labeling the iterative procedure right below Lemma 1 and making it painfully clear how that turns into Algorithm 1?
soundness: 2
strengths: * Originality: 
I'm uncertain about the level of originality here. Variational families within an RKHS framework have, to my knowledge, been explored previously. 
* Quality: The work appears to be of high quality. The results seem technically correct, and while I haven’t meticulously checked every mathematical detail, the derivations appear consistent with expectations.
* Clarity: Overall, the paper is well-written. I appreciated the clear, step-by-step walkthrough of the optimization problem and the detailed exposition of various bounds and relaxations necessary to implement the proposed approach.
* Significance: The proposed method, Flat Hilbert Bayesian Inference (FHBI), is presented as a generalization of Stein Variational Gradient Descent (SVGD) and Sharpness-Aware Minimization (SAM). The experiments demonstrate promising potential for FHBI in applications like LoRA-style fine-tuning.
summary: In traditional variational inference, the goal is to approximate an intractable posterior \( p(\theta | \mathcal{S}) \) by selecting a variational distribution \( q \) from a family \( \mathcal{Q} \) that minimizes a divergence, often the KL divergence, resulting in the optimization problem:

$$
\arg \min_{q \in \mathcal{Q}} \, \text{KL}(q(\theta) \parallel p(\theta | \mathcal{S})),
$$

where \( p(\theta | \mathcal{S}) \), termed the "empirical posterior" by the authors, is the target distribution.

This paper introduces a variational family within a Reproducing Kernel Hilbert Space (RKHS) framework and further reformulates the optimization problem to focus on approximating the "general posterior," \( p(\theta | \mathcal{D}) \), over the dataset \( \mathcal{D} \). The optimization problem thus becomes:

$$
\arg \min_{f \in \mathcal{H}^d, \, \|f\| \leq \epsilon} \, \text{KL}(q_{[I + f]}(\theta) \parallel p(\theta | \mathcal{D})),
$$

where \( q_{[I + f]}(\theta) \) represents the transformed variational distribution. This reformulation leads to multiple steps of approximation, relaxation, and bounding, culminating in an iterative optimization procedure (detailed below Lemma 1).

The proposed method, Flat Hilbert Bayesian Inference (FHBI), is designed to enhance generalization in Bayesian inference by leveraging the structure of RKHS.
weaknesses: * The general posterior defined, $p(\theta|\mathcal D)$, does not appear to me to be the correct theoretical counterpart to the "empirical" posterior. See more in questions below. 
* If the paper’s originality hinges partly on targeting the "general posterior" $p(\theta|\mathcal D)$, I have reservations about its practical benefits. In addition to my concerns about the missing sample size term, targeting this posterior seems unlikely to yield practical advantages and might even be counterproductive. The authors expend considerable effort introducing approximations to recast the optimization problem in terms of the "empirical posterior," which requires potentially loose bounds. The value of this detour would be enhanced by a more thorough discussion of why these approximations are justified or necessary.
* The paper emphasizes "improving generalization" as a primary benefit of FHBI, yet this claim seems tenuous without more foundational support. It's unclear what is meant by "enhancing generalization in Bayesian inference" mathematically, as the method doesn’t inherently introduce any features that theoretically boost generalization. Rather, it appears that FHBI, when applied to fine-tuning, yielded improved generalization performance on a benchmark dataset compared to baseline methods. This distinction between observed outcomes and the initial methodological intent would be clearer if the authors clarified that FHBI’s generalization performance was more an empirical finding than a theoretically driven design choice.

Meta-review: This paper proposes Flat Hilbert Bayesian Inference (FHBI), an algorithm to enhance generalization in Bayesian inference. Similar to Sharpness-Aware Minimization (SAM), each iteration of FHBI involves an adversarial perturbation step followed by a descent step, with the key difference being that these updates occur within a Reproducing Kernel Hilbert Space (RKHS). The proposed algorithm is accompanied by a theoretical analysis that extends previous SAM-like methods from finite-dimensional Euclidean spaces to infinite-dimensional function spaces. The paper also presents an empirical evaluation of FHBI on the VTAB-1K benchmark (consisting of 19 datasets across various domains), demonstrating that FHBI consistently outperforms the baselines.  The authors also provide the code for their implementation to facilitate the reproduction of their experiments.

I find the idea of implementing SAM in function space via the RKHS framework elegant and original. The fact that the empirical assessment also shows consistent improvements further adds to the significance of this work, making it a potentially impactful development. These strengths are also reflected in the reviewers' comments. However, the proofs related to the presented theory seem underdeveloped, needing either corrections or the inclusion of missing details about the derivation. Concerns related to the proofs were mainly raised by reviewer pWGm, who thoroughly reviewed the appendix where the proofs are presented. The reviewer and authors engaged in multiple rounds of feedback and response, and as they delved deeper into the details of the proofs, more missing details surfaced. While this does not necessarily indicate errors in the proofs, it clearly shows that the proofs in their current state may not be clear enough to readers, hindering a full grasp of the theoretical results and under-appreciating the paper's contributions. This suggests the paper needs careful proofreading and revision of the proof presentation to improve clarity.

Unfortunately, as the theory is one of the major contributions, its clarity and accessibility are crucial and cannot be overlooked, despite the positive reviews from other reviewers. As much as I appreciate this paper (the idea of performing SAM in RKHS, the theoretical analysis, and the positive empirical results), sadly I have to recommend resubmission after a revision. I hope the authors are not discouraged by this decision and are assured that they have produced very interesting work that just needs some polishing to be well-understood and appreciated.

As a minor comment, I echo the observation made by some reviewers that the paper uses some uncommon terminology, which can lead to confusion. I hope the authors will make these small changes in their future submission to improve clarity. In particular, I encourage them to replace "general loss" with "generalization loss" or "true risk," and perhaps "functional space" with "function space."

Decision: Reject