Improving Generalization of Complex Models under Unbounded Loss Using PAC-Bayes Bounds
Authors: Xitong Zhang
Avrajit Ghosh
Guangliang Liu
Rongrong Wang
Venue: Accepted by TMLR ()

Abstract: Previous research on PAC-Bayes learning theory has focused extensively on establishing tight upper bounds for test errors. A recently proposed training procedure called PAC-Bayes training, updates the model toward minimizing these bounds. Although this approach is theoretically sound, in practice, it has not achieved a test error as low as those obtained by empirical risk minimization (ERM) with carefully tuned regularization hyperparameters. Additionally, existing PAC-Bayes training algorithms often require bounded loss functions and may need a search over priors with additional datasets, which limits their broader applicability. In this paper, we introduce a new PAC-Bayes training algorithm with improved performance and reduced reliance on prior tuning. This is achieved by establishing a new PAC-Bayes bound for unbounded loss and a theoretically grounded approach that involves jointly training the prior and posterior using the same dataset. Our comprehensive evaluations across various classification tasks and neural network architectures demonstrate that the proposed method not only outperforms existing PAC-Bayes training algorithms but also approximately matches the test accuracy of ERM that is optimized by SGD/Adam using various regularization methods with optimal hyperparameters.

Published in Transactions on Machine Learning Research (09/2024)
Improving Generalization of Complex Models under Un- bounded Loss Using PAC-Bayes Bounds
Xitong Zhang
shangnit
Deportment of Computational Mathematics, Science and Engineering Michigon State University
Avrajit Ghosh
ghoshaur &msu. edu
Department of Compulational Mathematics, Science and Engineering Michigon State University
Guangliang Liu
Kuguan Sarguan5@mss.. @msu edu
Department of Computer Science and Engineering Michigan State University
Rongrong Wang
wangron63msu..cm
Department of Compulational Mathematics, Science and Engineering Department of Mothematics Michigan State University
Reviewed on OpenReview: https: // openreview. net/ forum? id= MP8bmavii6
Abstract
Previous research on PAC-Bayes learning theory has focused extensively on establishing light upper bounds for test errors. A recently proposed training procedure called PAC-Bayes training, updates the model toward minimizing these bounds. Although this approach is theoretically sound, in practice, it has not achieved test error as low as those obtained by empirical risk minimization (ERM) with carefully tuned regularization hyperparameters. Additionally, existing PAC-Bayes training algorithms (e.g., Pérez-Ortiz el al. (2021)) often require bounded loss functions and may need a search over priors with additional datasets, which limits their broader applicability. In this paper, we introduce new PAC-Bayes training algorithm with improved performance and reduced reliance on prior tuning. This is achieved by establishing new PAC-Bayes bound for unbounded loss and theoretically grounded approach that involves jointly training the prior and posterior using the same dataset. Our comprehensive evaluations across various classification tasks and neural network architectures demonstrate that the proposed method not only out performs existing PAC-Bayes training algorithms but also approximately matches the test accuracy of ERM that is optimizod by SGD/ Adam using various regularization methods with optimal hyperparameters.
1 Introduction
The PAC-Bayes bound plays vital role in assessing generalization by estimating the upper limits of test errors without using validation data. It provides essential insights into the generalization ability of trained models and offers thcoretical backing for practical training algorithms (Shawe-Taylor & Williamson, 1997). For example, PAC-Bayes bounds highlight the discrepancy between the training and generalization errors, indicating the need to incorporate regularizers in empirical risk minimization and explaining how larger datascts contribute to improved generalization. Furthermore, the effectiveness of the PAC-Bayes bounds in estimating the generalization capabilities of machine learning models has been supported by extensive experiments across different generalization metrics (Jiang et al. 2019).

---

Published in Transactions on Machine Learning Research (09/2024)
Traditionally, PAC-Bayes bounds have been primarily used for quality assurance model selection (McAllester
1998 1999:
Herbrich & Graepel 2000), particularly with smaller machine learning models. Recent work has
introduced a framework that minimizes a PAC-Bayes bound during training large neural networks (Dziugaite
& Roy 2017). Idcally, the gencralization performance of deep neural networks could be enhanced by directly minimizing its quantitative upper bounds, specifically the PAC-Bayes bounds, without incorporating any other regularization tricks. However, the effectiveness of applying PAC-Bayes training to deep neural networks is challenged by the well-known issue that PAC-Bayes bounds can become vacuous in highly over-parameterized settings (Livni & Moran 2020). Additionally, selecting suitable prior, which should be independent of training samples, is critical yet. challenging. This often leads to conducting a parameter search for the prior using separate datasets (Pérez-Ortiz et al., 2021). Furthermore, existing PAC-Bayes training methods are typically tailored for bounded loss (Dziugaite & Roy 2017 2018 Pérez-Ortiz el al. 2021 limiting their straightforward application to popular losses like Cross-Entronn Entropy.
On the other hand, the prevalent training methods for neural networks, which involve minimizing empirical risk with AddD,,mm, achieve satisfactory lest performance. However, they oflen require integration with various regularization techniques to optimize generalization performance. For instance, research has shown that
factors such as larger learning rates (Cohen et al. 2021 Barrett & Dherin 2020). momentum (Ghosh et al.
2022; Cattanco et al.] 2023). smaller batch sizes (Lee & Jang) 2022), parameter noise injection (Neelakantan (Luo et al. 2018) all induce higher degrees of implicit regularization, yielding better generalization. Besides, various explicit regularization techniques, such as weight decay (Loshchilov & Hutter 2017), dropout (Wei al. 2020). label noise (Damian et al. 2021 can also significantly affect generalization. While many studies have explored individual regularization techniques to identify their unique benefits, the interaction among these regularizations remains less understood. As result, in practical scenarios, one has to extensively tune the hyperparameters corresponding to each regularization technique to obtain the optimal test performance.
lt al. 2015; Orvieto et al. 2022), and batch normalization
Although further invostigation is needed to fully understand the underlying mechanisms, training models using ERM with various regularization methods remain the prevalent choice and typically deliver slate-of-the-ard test performance. While PAC-Bayes training is built upon a solid theoretical basis for analyzing generalization, its wider adoption is limited by existing assumptions about loss and challenges in prior selection. Moroover, it is still an open question regarding how to enhance PAC-Bayes training to match the performance of ERM methods with well-tuned regularizations.
In this paper, we propose a practical PAC--ayee-bouudd-beed training algorithm that nearly matches the performance of the optimally tuned ERM while being more robust to the choice of hyperparameters. The key differences between our algorithm and the previous ones are:
1. A new PAC-Bayes bound for unbounded loss is adopted for training, providing tighter numerical values for highly over wer-parametrized models.
2. The PAC-Bayes training is enhanced by an optional second stage of Bayesian training, which uses key parameters estimated from the PAC-Bayes training stage.
We provide mathematical analysis to support the proposed algorithm and conduct extensive numerical experiments to demonstrate its effectiveness.
2 Preliminaries
This section outlines the PAC-Bayes framework. For any supervised learning problem, the goal is to find a proper model h from some hepot space H, with the help of the training data S {2))m where 2; is the training pair with sample x; and its label yı- Given the loss function f(h; 2:): h + R, which measures the misfit between the true label yi and the predicted label by h, the empirical and population/ rppulationgeeneralizationnn errors are defined as:
/gencraiizaaionn
(h: = Σ
(h; 2))
[(b; D) = (s
1=1
2

---

Published in Transactions on Machine Learning Research (09/2024)
by assuming that the training and testing data are both ii... sampled from the same unknown distribution D. PAC-Bayes bounds include family of upper bounds on the generalization error of the following type.
Theorem 2.1. (Mawrer 2004) Assume the loss function E is bounded within the interval [0, 1]. Given preset prior distribution P over the model space H. and given a scalar ō E (0, 1), for any choice of i.dd m-sized training dataset S according to D, and all posterior distributions Q over H,
of S En=oe((; + v log 2 KL(QII)
2m
holds with probability at least 1- 8. Here, KL stands for the Kullback- Leibler divergence.
A PAC-Bayes bound measures the gap between the expected empirical and generalization errors. It's worth noting that this bound holds for all posterior Q for any given data-independent prior P and, which enables optimization of the bound by searching for the best posterior. In practice, the posterior mean corresponds to the trained model, and the prior mean can be set to the initial model. In this paper, we will use || to denote a generic norm, and || 2/2 to denote the L2 norm.
The focus on Neural Networks: PAC-Bayes training algorithms, including the one proposed here, can be applied to a wide range of supervised learning problems. However, in this paper, we will focus on the training of deep neural networks for the following reasons:
Over-paranmeteriieed deep neural networks present a significant challenge for PAC-Bayes training, as the Kullback-Leibler (KL) term in the PAC-Bayes bounds is believed to grow rapidly with the number of parameters, quickly leading to vacuous bounds. According to the current literature, deep networks are indeed one of the most critical models where existing FAC-Bayes training algorithms encounter difficulties.
Intuitively, minimizing the generalization bound does not require additional implicit or explicit regularization, thus reducing the need for tuning. This advantage of reduced tuning is particularly significant in the training of neural networks, where extensive tuning is typically required for ERM. However, for classical problems such as Lasso, which involve only one or two hyperparameters, this advantage of PAC-Bayes training may be less pronounced.
3 Related Work
PAC-Bayes bounds were first used to train neural networks in Dziugaite & Roy (017)) Specifically, the bound McAllester (1999) has been employed for training shallow stochastic neural networks on binary MNIST
classification with bounded 0-1 loss and has proven to be non-vacuous. Following this work, many recent
studics (Zhou et al. 2018: Letarte et al. 2019 Rivasplata et al., 2019 Pérez-Ortiz et al., 2021 Biggs & Guedj 2021 Perez- el al.J 2021 Viallard et al. 2023) expanded the applicability of PAC-Bayes bounds to wider range of neural network architectures and datasets. However, most studies are limited training shallow networks with binary labels using bounded loss, which restricts their broader application to deep network training. Although PAC- Bayes bounds for unbounded loss have been established (Audibert &
2019 Kuzborskij & Sxepesvári 2019 Haddouche et al. 2021: Rivasplata et al. 2020; [Roriguez- Gálvez et al. 2023 Casado et al. 2024), it remains unclear whether these bounds can lead to enhanced test performance in training neural networks. This uncertainty arises partly because they usually include assumptions that are difficult to validate or terms that are hard to compute in real applications. For example, Kuzborskij & Szepesvári (2019) derived a PAC-Bayes bound under the second-order moment condition of the unbounded loss. However, as mentioned in the paper, that bound is semi-empirical, in the sense that it contains the population second order moment of the loss with respect to both the posterior and the data distributions. Since conditional on the posterior, the samples are no longer i..dd., this type of bound is difficult to estimate. To the best of our knowledge, existing PAC-Bayes bounds built under the second-order moment. condition all suffer from this issue.
Catoni 2011 Alquier & Guedj] 2018 Holland]
Recently, Dziugaite et al. (2021) suggested that a tighter PAC-Bayes bound could be achieved with a data-dependent prior. They divide the data into two sets, using one to train the prior and the other train the posterior with the optimized prior, thus making the prior independent from the training dataset.
3

---

Published in Transactions on Machine Learning Research (09/2024)
for the posterior. This, however, reduces the training data available for the posterior. Dziugaite & Roy (2018) and Rivasplata et al. (2020) justified the approach of learning the prior and posterior with the same set of data by utilizing differential privacy. However, the argument only holds for priors provably satisfying the so-called DP(c))
in differential privacy, which limits their practical application. [Pérzz-
let al. 2021 also empirically shows training with Dziugaite & Roy (2018) could sacrifice lest accuracy if the bound not tight enough. In this work, we advance the PAC-Bayes training approach, enhancing its practicality and showcasing its potential in challenging settings.
4 Proposed method
4.1 Motivation
To help readers understand the challenges involved in designing practical PAC-Bayes training algorithms for deep neural networks, we begin by offering a detailed examination of the limitations inherent current PAC-training algorithms and PAC-Bayes bounds. Most popular PAC-Bayes training algorithms (Dziugaite 2017 2018 Pérez- Ortiz et al., 2021) are designed for learning problems with bounded loss. When dealing with unbounded loss, it is necessary to clip the loss finite range before training, which can result in suboptimal performance, as demonstrated in Table 1 in the numerical section.
& Roy
Some PAC-Bayes bounds for unbounded loss have been established in the literature, exlending the requirement from bounded loss to sub-Gaussian loss (Theorem 5.1 of Alquier (2021)), loss that satisfies the so-called hypothesis-dependent range (HYPE) condition (Haddouche et al. 2021), and loss controlled by some finito cumulant generaling function (CGF) (Rodríguez et al. 2023). More specifically, 2021), the loss requirement is relaxed the existence of K (h) such that sup=-pe(h, 2) S. ((),, 4h E H. This is known as the HYPE condition and is much weaker than the requirement for bounded loss. However, commonly used cross-entropy loss still does not satisfy this condition without additional assumptions the boundedness of both the input and the model. Furthermore, our numerical experiments indicate that minimizing bounds tailored for sub-Gaut Gaussinnn and CGF losses during training is also largely ineffective. Specifically, on CIFAR10 with CNN9, the bound values for sub-Gaussian and CGF losses are 3.17 and 5.08, respectively. These values almost do not decrease during training and are even higher than the initial test loss (with value of 2.30) under random weight initialization, indicating that these bounds can not contribute improving the training performance.
(Haddouche el al,
Lastly, the PAC-Bayes bound established under the (theoretically) weakest assumption is in Kuzborskij & Szepesvári (2019), where the loss is only required to have finite second-order moment. The associated bound holds with probability 1 --, 2:
1 n² En-o((h;D) Enwee(h; + + /2
I In m²3)), + +3) (KL +
(1)
where B = E [1; (b; However, this bound is not casily amendable for training the term En (h; in 3 is population second-order moment of the loss with respect to the data and the posterior distribution. Since the data is no longer i..dd when conditioned on the posterior, estimating this population's s second-order moment becomes challenging.
The above-mentioned limitations related literature motivate us to design a new PAC-Bayes bound for unbounded loss that easy apply in training.
4.2 A new PAC-Bayes Bound for Unbounded loss
We propose a training-friendly PAC-Bayes bound that holds under mild conditions.
4

---

Published in Transactions on Machine Learning Research (09/2024)
4.2.1 Condition under which the new bound will hold
The new bound we shall propose is based on the following assumption of the loss function in Definition Please note that the X in Definition 4.1 will later represent the training loss the PAC-Bayes analysis, and [[X] will represent the population loss.
1.1
Definition 4.1 (Exponential moment on finite intervals). Let X be a random variable defined on the probability space P) and S 71 S 72 S oo be two numbers. We call any K > exponential moment bound of X over the interval [7:77], when
Elexp (+([[]] X))] s exp (7²K)
(2)
holds for all y km.721
We provide a few remarks to clarify the position of this definition the literature.
Remark 4.2 (The left-side moment). Similar to certain other PAC-Bayes conditions in the literature
[Alquier 2021 this definition only requires the left-side moment to be bounded. Specifically, it requires E exep (7
X))] to be bounded for y > O, rather than for all 7 E R. The need for this left-side condition arises from the goal of establishing an upper bound on the population loss. An upper bound on the population
loss [[[] in terms of the empirical loss X translates to upper bounding E[X] - X. For positive 7, this leads
the left-side condition.
Remark 4.3 (y within a finite interval bounded away from 0). In practice, we set both Y and 72 to positive values. We observe that restricting to a finite interval h 72] with positive n. often reduces the bound K. In contrast, all previous bounds use 71 = 0.
L1, MSE, Huber loss, hinge loss, Log-cosh loss, quantile loss) are non-negative, it is of great interest to analyze the strength of Definition 4.1 under X 2 0. In this case, we can show that Definition 4.1 is weaker than the second-order moment condition.
Remark 4.4 (Non-negative loss). Since most loss functions in machine learning (e.g..
Lemma 4.5 (Comparison with the second-order- moment condition). For non-negative random variable 2 0, the existence of K on the interval 7 E [0, oo) in Definition 4.1 can be implied by the eristence of the second-order moment EX²2 <0.
The assumption X 2 0 in Lemma 4.5 can be further relaxed to X -M with M > 0, as in this case the random variable X + M is non-negative to which Lemma 4.5 can be applied.
Remark 4.6 (Comparison with the first-order-moment condition). Still under the assumption X 2 0, when the 71 in Definition 4.1 is finite (bounded away from 0), the existence of K can be implied by the existence of first-order moment. Indeed, by taking K H[] the inequality [[[] X s E[X] (assumed X 2 0) immediately implies Equation 22 However, this argument does not hold when 71 - Hence we cannot say Definition is as weak as the first-order moment condition.
71
Proof of Lemma 4.5 We show that EX²2 < oo implies Definition 4. holding for any 7 E [0, oc) with some 0c, we have (EX)² < EX² < oc. If 2 EX then it suffices take the K in
finite K. Since E.X²2 2
(Ex < 2K
to be K= EX S (EX)= K1. I[ 1. then using the inequality Ex:
5sl+rtd" V
(X) with x >EX - xX) SEX < 1, we have
Ee TRx < E(1 + <EX x)+ 7 -X3) - VarX) Var(X)
Therefore, it suffices to take K Var(X) = K2. Collecting the two cases, we see taking K max K2} would be enough for Definition 4.1 to hold with 71 0, = 0o.
5

---

Published in Transactions on Machine Learning Research (09/2024)
Now, we generalize Definition 4.1 to the PAC-Bayes setting where the random variable is parameterized by models in a hypothesis space.
Let us first explain what we mean by random variables parameterized by models in a hypothesis space. In (()z)))) y), where h is the model and is the misfit between the model output h(x)) and the data y- For a fixed model h, ((h) is a random variable whose randomness comes from the input pairs (x,y) D (D is the data distribution). Since (h) varies with h, we call it a random variable parameterized by the model h.
supervised learning, we define (h) as ((h = f(h(z))
Definition 4.7 (Exponential moment over hypotheses). Let (h) be a random variable parameterized by the hypothesis h in some hypothesis space H (i.e., h H), and fix interval [n, 72l with 0 S n < I cc. Let (Px. A E }} be family of distribution over H parameterized by X E A C R Then, we call any non-negative function (A) exponential moment bound for X Xhh over the priors (Px, X E A} and the interval [2,,7],, if the following holds
EA Blexp (r(E(X - ( S exp (7 K())),
for all 7 E (m,72), and any A E REACRO RR. The minimal such K (())
1 log( Blexp (x(EXX()) x(h))))).
Kmin (A)
sup yelm,,a
(3)
7
Definition 4.7 is a direct extension of Definition 4.1 to random variables parametrized by hypothesis space. Similar to Definition 4.1 when dealing with non-negative loss, the existence of the exponential moment bound Kmin is guaranteed, provided that the second-order moment of the loss is bounded, or provided that the first-order moment of the loss is bounded and V1 is bounded away from 0. Remark 4.8 (Dependency of K on prior parameters). The key difference in Definition 4.7 compared to prior work is that we allow the exponential moment. bound K to depend on the prior parameter X, rather than requiring a single K to hold for all possible À E A. This advantage carries over into the PAC-Bayes bound established in the next section.
In practice, there are two options to compute K(A): one is to use the upper bound derived in the proof of Lemma. 4.5 and the other is to estimate Kmin directly from the data, via Equation 3 For both options, we need to estimate the expectation of some random variable over the prior distribution Px for a given X from the training data. Fortunately, this is much easier than estimating the expectation with respect to the posterior distribution Q, as is required by previous work Kuzborskij & Szepesvári] (2019), since in our case, after conditioning the prior, the training data remains i.i.d., allowing a reliable approximation of the population mean by the empirical mean when the data is abundant.
4.2.2 A new PAC- Bayes bound
We are ready to present the PAC-Bayes bound for losses that satisfy Definition 4.7
Theorem 4.9 (PAC-Bayes bound for unbounded loss with a preset prior distribution). Given a prior distribution Px over the hypothesis space H. parametrized by X E A. Assume the loss E(h, 2;) as a random variable parametrized by h satisfies Desinition 4.7. For any 0 < 8 and 7 E [m,, 72), we have
Prs ve E Q, onh En-ee((; +
(log + -L(())))))+++) + -K(2)) 21-8
mm
where Q is the set of all probability distributions..
The proof of this theorem is available in Appendix
Remark 4.10 (Asymptotic convergence rate). By setting 7.7h (0m-111)) observe that the asymptotic behavior of this bound aligns with the 0(m-1/2) convergence rate of popular PAC-Bayes bounds in the literature.
Remark 4.11 (Applicability on CE loss). This theorem, when combined with Lemma 4.5 guarantees that the 0/2- convergence rate is achieved for CE loss under a bounded second-order moment. condition.
1/2)

---

Published in Transactions on Machine Learning Research (09/2024)
Remark 4.12 (The role of (())). The improvement of the new bound primarily stems from our definition of K (()). Specifically, Def 4.7 allows K to be dependent on both the range of 7 and the prior parameter X. This flexibility reduces the value of K. The trade-off is that 7i and 7/2 need to be selected in advance as hyperparameters, and (A) must be estimated from data and optimized during training, which slightly increases the difficulty of the optimization.
To support the claim of Remark 4.12 we provide numerical evidence of the advanlage of our K in Figure I
We use Sedtriuue--Gálvee et al. 2023 and [Alquerr (2021) as the baseline PAC-Bayes bounds for unbounded loss for comparison. First, observe that both baseline bounds can be written in the form of
En~ S E KL(OIPP) + +,K,
(4)
myy
where & hides some absolute constant. Namely, both bounds have a K term, although they follow different definitions. When the loss is assumed to be sub-Gassuain, K is defined as the sub-Gaussiag norm of the loss (Theorem 5.1 of Alquier (202T Under the CGF condition, K is calculated based on the specific choice of the CGF Goodriguez-áálver et al., 2023) (sce Appendix C.2 for our choice of CGF this experiment). For our bound, the K is chosen as our Kmin (A) for various choices of 71 and 7/. This experiment conducted CIFAR10 using CNN9. We follow standard practice in PAC-Bayes training literature by using Gaussian prior, fixing its mean to the random Kaiming initialization, and picking a universal variance for all the weights. According to Definition 4.7 our K depends on "1 and 72, and the prior parameter X, which we set to the universal standard deviation (std) of the prior weight distribution. Since we observed that the variation of K with respect to 72 is very minimal, Figure | only plots K as a function of the prior std and 71. The K values in the previous two bounds (i.c., aus-Gaussian and CGF) do not depend on these paramcters. Hence, they are represented by two horizontal planes. We observe the following:
our K values approach those in the CGF bound as 71 approaches 0 and the prior std becomes large;
our K is significantly smaller than those in the other two bounds near the region of the optimal choice of prior std and 7, which we determined by testing various pairs of nn and prior standard deviation for the combination that yields the tightest bound after optimized over the posterior. This
explains the advantage of making K function of the prior std;
since all other terms are the same, the smaller the K is, the tighter the bound is after optimizing over 7.
Figure also indicates that the value of Kmin is more stable with respect to 7i than with respect to the prior standard deviation, especially near the optimal parameter region, suggesting that choosing ?1 is easier. Therefore, in our algorithm, we fix 7n in advance while optimizing over the prior variance A during training. It turns out that the same choice of 71 works for a wide variety of architectures, as shown in the numerical section;
4.3 PAC-Bayes training based on the new bound
With the relaxed requirements on the loss function, the proposod bound offers a basis for establishing effective optimization over both the posterior and the prior. We will first outline the training process, which focuses on jointly optimizing the prior and posterior to avoid the complex typer-paramett search over the prior Pérez-Ortiz et al. (2021), followed by a discussion of its theoretical guarantees. The procedure is similar to the one in Dziugaite & Roy (017)), but has been adapted to align with our newly proposed bound.
We begin by paramcterizing the posterior distribution as Qw, where w E Rd represents the parameters of the posterior. Next, we parameterize the prior as Px, where À E lR. We operate under the assumption that the prior has significantly fewer parameters than the posterior, that is, k < d; the relevance of this assumption will become apparent upon examining Theorem 4.15 For our PAC-Bayes training, we propose optimize over all three variables: w, 7, and X:
(w, 4) arg min LPaC(W,7, 7, 7, v.
(P)
7

---

Published in Transactions on Machine Learning Research (09/2024)
1p14
ours
10
10m2s Tinoo Ttree Stmes Aones 108 1055
CGF
10-18
sub-G sub-Gaussien
prior std
10
10
10-1a
Figure 1: Our definition of K demonstrates notable advantage in terms of the achieved numerical value compared to the K in the Bab-Gaussian and CGF bounds. Our K varies with Yi and the prior standard deviation, while the previous bounds remain constant as they don't depend on these paramcters. The vertical dotted line marks the optimal choice of prior std and "r during training, which we determined by testing various pairs 7/1 and prior standard deviation for the combination that yields the smallest value of the bound after optimized over the posterior. The figure shows that near this optimal combination, our K is much smaller than those in the other two bounds. This experiment is conducted on CIFAR10 using CNN9; see details in Sec.
where
1 _ + (L(22((P)) (*K(A).
(, = En~o. (): + - (log 7m
5
Please note that here K depends on the prior parameter X, and we need to optimize it along with other terms. The probability failure should be fixed in advance.
4.4 Theoretical Analysis
Two important questions arise from the training framework based on Equation 5
there any theoretical guarantee for the performance of training over the prior?
Does the exponential moment bound function K(A) exhibit enough regularity as a function of À to be accurately estimated from the data?
This section addresses these two questions. First, provide PAC-Bayes bound that allows the prior to be data-dependent. Then, we validate the Lipschitz continuity of K(A) when the prior and posterior are both set to Gaussian distributions.
4.4.1 PAC-Bayes bound for arbitrary distribution families
We first make the following assumptions, which will be shown to automatically satisfy by the special Gaussian prior and posteriors.
Assumption 4.13 (Continuity of the KL divergence). Let Q be a family of posterior distributions, let P (Px, X e C R*} be a family of prior distributions parameterized by A. We say the KL divergence [ is continuous with respect to X over the posterior family, if there exists some non-decreasing function 7/. (x) R4 R4 with vn (0) 0, such that [KL(Q||P)) KL(Q[PP) m (IIA -(]), for all pairs
X, A A and for all Q E 2

---

Published in Transactions on Machine Learning Research (09/2024)
Assumption 4.14 (Contimuity of the (()).. Let Kmin (A) be as defined in Definition 4.7 Assume it is
Lipschitz continuous with respect to the parameter À of the prior in the sense that there exists a non-decreasing function (2)x)= : R+ + with 72 (0) = 0 such that [m (A) - Kmin (A)l S m2(1 - X(), for all X, A E A.
These two assumptions are quite weak and can be satisfied by popular continuous distributions, such as the exponential family.
We now provide an end-to-end thcorem that guarantees the performance of the PAC-Bayes training based on Equation (5) for general priors and posteriors.
Theorem 4.15 (PAC-Bayes bound for unbounded losses and trainable priors). Assume the loss E(h, 2;) as a random variable parametrized by model h satisfies Definition 4.7 Let Q be a family of posterior distribution Q = [Qw: w ] let & (Px, E A C BR*} be a fammiyy of prior distributions parameterized by X. Let n(e) (A. e))) be the covering number of the set of the prior parameters. Under Assumplion Assumption 4. the following inequality holds for the minimizer (w, 9, X) of Equation with probability as least 1 e with (n(E) + m )8:
13 and P and any 6, 8
En~ (():)) LeAc (W, (W,,, +
B C. (m]()) + na(E)) log(n(s)+ "2-) and C and B are constants depending 7, 72, and the upper bounds of the paramelers in the prior and posterior.
where 71 Be
nm
The proof is available in Appendix A.2
The theorem provides a generalization bound on the model learned as the minimizer of Equation P with
data-dependent priors. This bound contains the PAC-Bayes loss LAAC along with an additional correction term '7, which is notably absent in the traditional PAC-Bayes bound with fixed priors. Given that (w, 4, Å) minimizes ( 6), evaluating LPAC at its own minimizer ensures that the first term is small. the correction term is also small, then the test error remains low. In the next section, we will delve deeper into the condition for this term to be small. Intuitively, selecting small helps to maintain low values for the first three terms in y, Although smaller e increases the n(s) in the last term, this increase is moderated because il is inside the logarithm and divided by the size of the dataset.
4.4.2 Restricting to the Gaussian Families
For the LPAC objective to have a closed-form formula, we employ the Gaussian distribution family. For ease of illustration, we introduce a new notation for the parametrization. Consider model denoted as fe, where f represents the model architecture (e.g., VGG net), and 0 is the weight. Tn this context, fe aligns with the h discussed in carlier scctions. Moving forward, will use fo to refer to the model instead of h.
We define the posterior distribution of the weights as a Gaussian distribution centered around the trainable weight μ,, with trainable variance o,, i.e., the posterior weight distribution is (nu, diag(or)), denoted by The assumption of a diagonal covariance matrix implics the indopendence the weights. We consider two types of priors, both centered around the initial weight 04 of the model (as suggested by Dziugaite
[ m
Roy (217 but with different settings on the variance.
Scalar prior: we
universal scalar to encode the variance of all the weights in the prior, i.c., the weight distribution of Px is N N(Ho, XTa), where À is a scalar. With this prior, the KI. divergence KL (Q ||P (o..))
in Equation is:
1 - 1 log(or d(log(A)) - 1) (||r) + || moll))
2
Layerwise prior: weights in the ith layer share common variance A, but different layers could have different variances. By setting X = (Ai, )x) as the vector containing all the layerwise variances of a k-layer
A clarification of the notation: have defined several model-related notations h, w, 6, n, and To clarify, their relations are Sa, O Qu Qw. and J aj
9

---

Published in Transactions on Machine Learning Research (09/2024)
neural network, the weight distribution of prior Px is N(u,, BlockDiaggA))) where BlockDiag(A) is obtained by diagonally stacking all Aa into a d x d matrix, where d; is the number of weights of the ith layer. The KL divergence for layerwise prior is in Appendix A.3 For shallow networks, it is enough to use the scalar prior; for deep neural networks and ncural networks constructed from different types of layers, using the layerwise prior more sensible.
By plugging in the closed-form Equation (7) for (L have the following corollary that justifies the usage of PAC-Bayes bound on large neural networks with the trainable prior.
|i (P into the PAC-Baaes bound in Theorem 4.15
Corollary 4.16 (Validity of trainable Gaussian priors). Suppose the posterior and prior are Gaussion distributions as defined above. Assume all parameters for the prior and posterior are bounded, i.e., we restrict the model parameter Ja, the posterior variance a and the prior variance X, all to be searched over bounded sets, θ = {u E RR -|lel|2 I vdM}, = for E B4 11o1l11 S dT}, A {A E te=",e"1"), respectively, with fixed M, a. > 0. Then,
Assumption 4.13 1.13 holds with nn (x)x)= = L where L1 1 max ea (2VdM + +T)]
Assumption 4. 14 holds with mp(xa)= = L2x, where L2 = 7 (2dM3 + dlatb)
With high probability, the PAC-Bayes bound for the minimizer of Equation has the form
Ee- a (ffo: D) S LPAc o). 1. 6) +
where 11 Is
(1+ log B is a conslant depending on 7 6, M, d, T, a, b, m
XCC+B)AAm L L1 + L2, = max (b + a, 2(12 - 71)} C 7im + 72
1
"imm
In the bound, the term LP o). 3, A, 5) inherently minimized as it evaluates the function LPac(s.... at its own minimizer. The overall bound remains low if the correction term 17 can be deemed insignificant. The logarithm term in the definition of 1 grows very mildly with the dimension in general, so can treat it (almost) as constant. Thus, 1) 71m small if prior's degree of freedom & is substantially less than the dataset size m; 2). This bound still achieves the asymptotic rate of 0/m= after optimizing over 71- We note that even if the corollary assumes that the parameters (i.e., mean and variance) of the Gaussian distribution are bounded, the random variable itself is still unbounded, so the loss is still unbounded. The proof and more discussions can be found in Appendix A.4
I_. from which we see that 1). (and therefore the bound) would be
4.5 Training algorithm
Estimating Kmin (A): In practice, the function Kmin (A) must be estimated first Since we showed in Corollary 1.16 and Remark 4.6 that (A)n is Lipschtiz continuous and bounded, we can approximate it using piecewise- linear functions, More explicitly, we use Equation (3) and Monte Carlo Sampling to estimate Kmin on some discrete grid of A, then we interpolate Kmin(A) using piecewise linear functions. More details are in Appendix B.1
Notably, since for each fixed À A, the prior is independent of the training data, this procodure of estimating Kmin (X) can be carried out before training. Recall that after Remark 4.8 we discussed two ways to estimate Kmins using Equation
2 illustrates the advantage of the former approach. Besides having a much smaller numerical value, estimating Kuni directly from Equation (3) also involves fewer constraints compared to using the upper bound, which requires a non-negative loss with a bounded second-order moment.
(3) or using its theoretical upper bound from Lemma 4.5) Figure
Two-stage PAC-Bayes training: Algorithm I outlines the proposed PAC-Bayes training algorithm with scalar prior. The version that uses the layerwise prior is detailed in Appendix B.2 Algorithm D contains two stages. Stage 1 performs pure PAC-Bayes training, and Stage 2 is an optional Bayesian refinement stage that is only activated when the first stage does not sufficiently reduce the training loss. For Stage 1, although there are several input parameters to be specified, one can use the same choice of values across very different
28ee Appendix A.4 for the explicit form B.
10

---

Published in Transactions on Machine Learning Research (09/2024)
Algorithm PAC-Bayes training (scalar prior)
Input: initial weight E R, Ti = 500, A1 = ==-22,00 = e =0.5, 72 = 10. // Ti, \1, X2, 71 72 are
fixed in all experiments of Sec
Output: trained weight i, posterior noise level δ
f + po, v + la log( (og(2l) b t log(H ( |
Obtain (A) with A = [11 using Equation 266 (Appendix Algorithm
/*Slage /
for epoch = T do for sampling one batch s from do //Ensure non-negative variances
+ exp(b), t exp(v)
Px + NNNNo Ala), Quo μ NO0; diag(or) //Get the stochastic version of Ea- E
Draw one θ Qu.o and evaluate S)
Compute the KT. divergence as Equation
Compute 7 as Equation (8)
Compute the loss function L LPAC in Equation
//Update all parameters
Oc b- b + 1 0b v + v ov t- μ + 13p
ac
ac
end for end for //Fir the noise level from now on f- exp(v) /*Slage 2*/ // Run this slage the training has not fully converged during Slage 1 while not converge do for sampling one batch from S do
// //Noise injection
Draw one 0 ~ Qp. and evaluate S) as //Update model parameters
ac a + 0p end for end while f 4-+p μ
30
12
10
1934 tonis
Figure Comparison of the upper bound max(K,, K2} for Kmins derived in Lemma 4.5 with the data- driven estimate of Kmin obtained via Equation (3) CNN9 using the CIFAR10 dataset. with the prior parameterized as Gaussian distribution centered at the Kaiming initialization.
11

---

Published in Transactions on Machine Learning Research (09/2024)
network architectures and datasets with minor modifications. Please see Appendix C.1 for more discussions. When everything else in the PAC-Bayes loss is fixed, 7 E [2:72] has a closed-form solution,
log
+ KI, (Ho
x = min
max
72
(8)
mKmin
Therefore, we only need to perform gradient updates on the other three variables, μ4 O,O.A.
The second stage of training: Gastpar et al. (2023); Nagarajan & Kolter (2019) showed that achieving high accuracy on certain distributions precludes the possibility of getting a tight generalization bound in overparameterized settings. This implies that it is less possible to use reasonable generalization bound to fully train one overparameterizee model on a particular dataset. It is also observed in our PAC-Bayes training experiments that, oftentimes, minimizing the PAC-Bayes bound only (Stage 1) cannot make the training accuracy reach 100%. If this happens
add sccond stage to further increase the training accuracy. Specifically, in Slage 2, we continue to update the model by minimizing only Ee- ( over je, and keep all other variables (i.e., X, or) fixed to the solution found by Stage 1. This is essentially a stochastic gradient descent with noise injection, the level of which has been learned from Stage 1. The two-stage training is similar to the idea of the learning-rate scheduler (LRS). In LRS, the initial large learning rate introduces an implicit bias that guides the solution path towards a flat region (Cohen et al., 2021 Barrett Dherin 2020), and the later lower learning rate ensures the convergence to local minimizer in this region. Without the large learning rate stage, it cannot reach the flat region; without the small learning rate stage, it cannot converge to a local minimizer. For the two-stage PAC-Bayes training, Stage 1 (PAC-Bayer stage) guides the solution to flat regions by minimizing the generalization bound, and Stage 2 is necessary for an actual convergence to a local minimizer.
Regularizations in the PAC-Bayes training: can see that in the case of Gaussian priors and posteriors, the PAC-Bayes loss nothing but the original training loss augmented by noise injection and a weight. decay, except that strength of both of them are automatically learned during training. More discussions are available in Appendix B.3]
By plugging the KL divergence Equation 7 into we
Prediction: After training, we use the mean of the posterior as the trained model and perform deterministic
prediction on the lest dataset. In Appendix [B.]] we provide some mathematical intuition of why the deterministic predictor is expected to perform even better than the Bayesian predictor.
5 Experiments
In this section, we demonstrate the efficacy of the proposed PAC-Bays training algorithm through extensive numerical experiments. Specifically, conduct comparisons between our algorithm and existing PAC-Bayes training algorithms, well conventional training algorithus based on Empirical Risk Mininization (ERM). Our approach yields competitive test accuracy in all settings and exhibits a high degree of robustness w.r.t. the choice of hyperparameters.
Comparison with different PAC-Bayes bounds and existing PAC-Bayes training algorithms: We compared our PAC-Bayes training algorithm using the layerwise prior with baselines in Pérez- Ortiz et al. (2021): quod (Rivasplata et al. 2019), lambda {Thiemann et al.) 2017), classic (McAllester 1999 and bbb Blundell et al., 2015) in the context of deep convolutional neural networks. The bascline PAC-Bayee algorithms contain a variety of crucial hyperparameters, including variance of the prior (1e-2 to 5e-6), learning rate (1e-3 to 1e-2), momentum (0.95, 0.99), dropout rate (0 to 0.3) in the training of the prior, and the KL tradc-off cocfficient (1c-5 to 0.1) for bbb. These hyperparametors were chosen by grid search. The batch size is 250 for all methods. Our findings, as detailed in Table [ show that our algorithm outperforms the other PAC-Bayes methods regarding test. accuracy. It important to note that all four baselines employed the PAC-Bayes bound for bounded loss. Therefore, they need to convert unbounded loss into bounded loss for training purposes. Various conversion methods were evaluated by Pérez- Ortiz et al. (2021) and the most effective one was selected for producing the results presented.
Note that there are cases when the second stage not necessary, including but not limited to 1) the network shallow 2) the dataset is simple 3) good prior chosen. In these cases, the training accuracy can already reach 100% in Stage
12

---

Published in Transactions on Machine Learning Research (09/2022)
0.8
sub- Gaussian
Test Accurac
0.6
Bound
CGF
0.5
0.4
Initial test
0.3
0.2
0.1
100
200
300
400
500
100
200
300
400
500
Training Epoch
Training Epoch
Figure 3: Minimizing PAC-Bayes bounds based on sub-Gaussion, CGF and our proposed bound CIFARIO using CNN9. The lest error of randomly initialized model is shown inilial lest error. Minimizing our bound (ours) achieves better test accuracy compared with optimizing the other two (sub-Gaussin Gaussian and CGF).
Table 1: Test accuracy of convolution neural networks on CTFAR10. The test accuracy of baselines for bounded loss from Table 5 of Ortiz et al. (2021), calculated as 1-the zero-one error of the deterministic predictor. subG represents the sub-Gaussias bound. Our proposed PAC-Bayes training with layerwise prior (layer) achieves the best test accuracy across all models.
bounded
unbounded
quad lambda
classic
bbb
subG
CGF
layer
CNN9
78.63
79.39
78.33
83.49
78.53
78.35
85.46
CNN13
84.47
84.48
84.22
85.41
84.30
84.42
88.31
CNN15
85.31
85.51
85.20
85.95
84.98
85.13
87.55
To demonstrate the necessity of our newly proposed PAC-Bayes bound for unbounded loss, we compared this new bound with two existing PAC-Bayes bounds for unbounded loss. One is based on the sub- Gaussian assumption (Theorem 5.1 of Alquier (2021]), while the other (Theorem 9 of RRodiiguez Gálvez et al. (0023 assumes the loss function is a bounded cumulant generating function (CGF). It important to note that, as of now, no training algorithms specifically leverage these PAC-Bayes bounds for unbounded loss. Therefore, for fair comparison, conducted experiment by replacing our PAC-Bayes bound with the other two bounds and using the same two-stage training algorithm with the trainable layerwise prior.
We also visualized the test accuracy when minimizing different PAC-Bayes bounds for unbounded loss in Stage 1. As shown in Figure 3 minimizing our PAC-Bayes bound can achieve better generalization performance. As discussed in Remark 4.12 the K terms in the two baseline bounds for unbounded loss are much larger compared to ours. This results in smaller 7, which increases the coefficient of the KL divergence term and forces the posterior to remain close to the prior rather than fitting the data elfectively. The details of the two baseline bounds are in Appendix C.2
Comparison with ERM optimized by SGD/Adan with various regularizations: We tested our PAC-Bayes training on CIFAR1O and CTFARIO R100 datasets with no data augmentation on various popular deep neural networks, VGG13, VGG19 (Smonnynn & Zisserman 2014), ResNet18, ResNet34 (He et al. 2016), and Densel21 (Huang et al., 2017) by comparing its performance with conventional empirical risk minimization by SGD / Adam enhanced by various regularizations (which we call baselines). The training of baselines involves grid search for the best hyperparameters, including momentum for SGD (0.3 to 0.9), learning rate (1e-3 to 0.2), weight decay (1e-4 to le-2), and noise injection (5e-4 to le-2). The batch size set to be 128. We reported the highest (est accuracy obtained from this search as the baseline results. For all convolutional neural networks, our method employed Adam with a fixed learning rate of le-4.
'Result with data augmentation can found Appendix C.3
13

---

Published in Transactions on Machine Learning Research (09/2024)
Table 2: Test accuracy of CNNs on C10 (CIFARIO) and C100 (CIFAR100) with batch size 128. Our PAC-Bayes training with scalar and layerwise prior are labeled scalor and layer. The best and sccond-best test accuracies are highlighted and underlined. Our PAC-Bayes training can approximately match the best performance of the basclinc.
VGG13
VGG19
ResNet18
ResNet34
Dense121
C10
C100
C10
C100
C10
C100
C10
C100
C10
C100
SGD
90.2
66.9
90.2
64.5
89.9
64.0
90.0
70.3
91.8
74.0
Adam
88.5
63.7
89.0
58.8
87.5
61.6
87.9
59.5
91.2
70.0
AdamW
88.4
61.8
89.0
62.3
87.9
61.4
88.3
59.9
91.5
70.1
scalar
88.7
67.2
89.2
61.3
88.0
68.8
89.6
69.5
91.2
71.4
layer 89.7
67.1
90.5
62.3
89.3
68.9
90.9
69.9
91.5
72.2
Table 3: Test accuracy of GNNs trained with AdamW versus our proposed method with scalar prior scalor. The best test accuracies are highlighted. The performance of our training can almost match the best results of the baseline obtained after carefully tuning hyperparameters.
CoraML
Citeseer
PubMed
Cora
DBLP
AdamW 85.7100.7
90.3
85.010.6
60.7:00.7
80.6+1.
GCN
scalar
86.1
90.0#0.4
84.9=0.8
62.0=0.4
80.5
AdamW
85.7+11.0
90.8++0.3
84.0+0.4
63.5++044
81.8+0.6
GAT
scalar
85.9110.8
90.611..5
84.4
60.9
81.0++..5
SAGE AdamW 85.71:0.5
90.5
83.5+0.4
60.6
80.7+0.6
scalar
86.5+0.55
90.0
84.410.6
61.2+0.2
79.9
APPNP AdamW
86.6+0.7
91.010.4
85.1
62.540.4
80.642.8
scalar 87.1
90.4#0.5
85.7+0.4
63.5+0.4
81.8+0.5
Since the CIFARIO and CIFAR100 datasets do not have a published validation dataset, we used the test dataset to find the best hyperparamcters of basclines during the grid scarch, which might lead a slightly inflated performance for baselines. Nevertheless, as presented Table 2 the test accuracy of our method is still competitive. Please refer to Appendix C.4 for more details.
Evaluation graph neural networks: To demonstrate the broad applicability of the proposed PAC-Bayes training algorithm to different network architectures, evaluated it on graph neural networks (GNNs). Unlike CNNs, optimal GNN performance has been reported using the AdamW optimizer for ERM and enabling dropout. To ensure the best baseline results, we conducted a hyperparameter search over learning rate (1e-3 to 1e-2), weight decay (0 to 1e-2), noise injection (0 to 1e-2), and dropout (0 to 0.8) and reported the highest test accuracy as the baseline result. For our method, we used Adam and fixed the learning rate to be lc-2 for all graph ncural networks. We follow the convention for graph datascts by randomly assigning 20 nodes per class for training, 500 for validation, and the remaining for testing.
We tested four architectures GCN (Kipf & Welling 2016), GAT (Veličković et al. 2017 SAGE Hamilton ct al.] 2017), and APPNP (Gasteiger et al. 2018) on 5 benchmark datasets CoraML, Citeseer, PubMed, Cora and DBLP (Bojchevski & Günnemann) 2017). Since there are only two convolution layers for GNNs, applying our algorithm with the scalar prior is sensible. For our PAC-Bayes training, we retained the dropout layer in the GAT is, since it differs from the conventional dropout and essentially drops the edges of the input graph. Other architectures do not have this type of dropout; hence, our PAC-Bayes training for these architectures does not include dropout.
14

---

Published in Transactions on Machine Learning Research (09/2024)
The test accuracy for CNNs on CIFARIO (C10) and CIFAR100 (C100) using a batch size of 2048. Values in (-) indicate how much the results differ from using batch size (128). Our PAC-Bayes traiming with scalar and layerwise prior are labeled as scalar and lager. The most robust results w.r.t. the increase of batch size are highlighted, indicating the elevated robustness of our method compared to the bascline regarding batch sizes.
Table 4:
VGG13
1es1et18
C10
C100
C10
C100
SGD
87.7 (-2.5)
60.1 (-6.8)
85.4 (-4.5)
61.5 (-2.6)
Adam
90.7 (+2.2)
66.2 (+2.5)
87.7 (+0.2)
65.4 (+3.8)
AdamW
87.2 (-1.1)
61.0 (-0.8)
84.9 (-2.9)
58.9 (-2.5)
scalar
88.9 (+0.2)
66.0 (-1.2)
88.9 (+0.9)
68.7 (-0.1)
layer
89.4 (-0.3)
67.1 (0.0)
89.2 (-0.1)
69.3 (+0.3)
Table 5: Test accuracy of 18 and VGG13 trained with different learning rates CIFAR1O. The best test accuracies are highlighted. Our method is more robust to learning rate variations.
Model
Method 30-5
5e-5
lc-4
20-4
3o-4
5c-4
le-3
layer Adam
88.4
88.8
89.3
88.6
88.3
89.2
87.3
ResNet18
66.6
73.9
81.2
85.3
86.4
87.0
87.5
layer
88.6
88.9
89.7
89.6
89.6
89.5
88.7
VGG13
Adam
84.3
84.8
85.8
87.4
87.9
88.3
88.5
Table 3 demonstrates that the performance of our algorithm closely approximates the best outcome of the bascline. Appendix C.5 provides additional details and more results. Extra analysis on few-shot text. classification with transformers is in Appendix C.6
100
500: Inle- betoott, U25e:3, Tur+taa Oua Glage 11-b-le b 9-122
100
Oussi-3e5 1,1112
100
500 Iwde bettabbd L2-bid. 6-122
Danchmied) 1ux50o0n
500 t-tel, belandid 12-2d, 2 120
90
Oun odage bgai22
90
Buttiited, bur-196
test test accuracy (%)
90
son bete- R
500 Iwdo bolowo
accuracy
test accuracy
80
70
70
70
60
60
60
40
40
40
40
50
60
70
80
90
100
40
50
60
70
80
90
100
40
50
60
70
80
90
100
training accuracy (9))
training accuracy (%)
training accuracy (%)
(a) PAA-Baaes and ERM training.
(b) PAC-Bayes training.
(c) ERM training.
Figure 4: Generalization gap (the difference between the training and the testing accuracy) in PAC-Bayes training versus ERM training using ResNet18 on the CIFAR10 dataset. Each point represents an intermediate model during training, plotted according to its test accuracy versus training accuracy. The linc y = z indicates the optimal, zero generalization gap. PAC-Bayes training has a smaller generalization gap throughout the training process (Fig. 4ab and remains stable despite changes hyperparameters (Fig. 4b). In constant, ERM training (Fig. 4cp very unstable to hyper-parameter changes. When comparing ERM with our method in Fig. 4a we picked the best ERM result (the blue one) in Fig. 4c that achieved the best final test accuracy. The discontinuity it has around the testing accuracy of 87%, is due to the activation of the learning rate scheduler.
15

---

Published in Transactions on Machine Learning Research (09/2024)
Evaluation on the sensitivity of hyperparamettrs: In previous experiments, selected specific batch sizes and learning rates the only two tunable hyperparameters of our algorithm, with all other parameters remaining constant across all experiments. We further demonstrate that batch size and learning rate variations do not significantly impact our final performance. This suggests a general robustness of our method to hyperpaaaretters. reducing the necessity for extensive luning. More specifically, with a fixed learning rale 5e-4 in our method, Table f shows that changing the batch size from 128 to a very large one; 2048, for VGG13 and ResNet18 does not significantly affect the performance of the PAC-Bayes training compared to ERM with extensive tuning as before. Also, shown by Table [ our algorithm is more robust to learning rate changes than ERM, which utilizes the optimal weight decay and noise injection settings from Table
We also examine the change in the generalization gap across the training process Figure Generalization gap is defined to be the difference between the training and testing accuracy. Algorithms with better generalization ability should yield a smaller generalization gap. The smallest possible gap O, which corresponds to the line y = z in Figure We observe in Figure 4aj that our PAC-Bayes training, when compared the ERM with the optimal hyperparameter setting, has smaller generalization gap (i.e., closer to the line y = x)) over the course of training, although the final test accuracies extremely small, confirming the effectiveness of using the PAC-Bayes bound for achieving good generalization. In addition, as we vary the choice of hyperparameters. than ERM Figure 4c indicating less need for hyperparameter tuming. Furthermore, Fig. suggests future directions for improving the PAC-Bayes training. Since Stage 1 yields the best generalization gap, we should focus on developing numerically tighter PAC-Bayes bounds to prolong Stage 1. Alternatively, we can aim improve the heuristic algorithm in Stage 2 to minimize the increase in the generalization gap during this
similar. The generalization gap in Stage 1 is
FAC-Bayes training is much more stable Figure 4b
stage.
Please refer to Appendix C.7 for extra experimental studies.
6 Conclusion and Discussion
In this paper, we demonstrated the great practical potential of PAC-Bayes training by proposing a numerically tighter PAC-Bayes bound and applying to train deep neural networks. The proposed framework significantly improved the performance PAC-Baycs training, making nearly match the best results of ERM. We hope this result inspires researchers in the field to further explore the practical implications of PAC-Bayes theory and make these bounds more useful in practice.
Acknowledgments
We thank Andrew Gordon Wilson for insightful and valuable discussions. A.G. was supported by NSF grant CCF-2212065. R.W. was partially supported by NSF grant CCF-2212065 and BCS-2215155.
References
Pierre Alquier. User-friendly introduction to pac-bayes bounds. ar preprint arXiv:2110.11216 2021.
Pierre Alquier and Benjamin Guedj. Simpler pac- mac-bayesian bounds for hostile data. Machine Learning, 107 (5):887-9022 2018.
Jean- Yves Audibert and Olivier Catoni. Robust linear least squares squaress The Annals of Statistics, 39 (5):2766 2794, 2011. doi: 10.1214/11-A05918. URL
David GT Barrett and Benoit Dherin. Implicit gradient regularization. arXiv preprinl wxXiw:2009.11162 2020.
Felix Biggs and Benjamin Guedj. Differentiable pac- bayes objectives with partially aggregated neural networks. Entropy, 23(10): 2021.
Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference machine learning, pp. 1613 1622. PMLR, 2015.
16

---

Published in Transactions on Machine Learning Research (09/2024)
Aleksandar Bojchevski and Stephan Günnemann. Deep gaussian embedding of graphs: Unsupervised inductive learning via ranking. ar Xix preprint orXiv: 707.03007, 2017.
Ioar Casado, Luis A Ortega, Andrés R Masegosa, and Aritz Pérez. mae-bayes-cheroof bounds for unbounddd losses. arXiv preprint or1::2400.001448 2024.
Matias D Cattanco, Jason M Klusowski, and Boris Shigida. On the implicit bias of adam. arXiv preprint WK::2399000072 2023.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2020.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.0000 2021.
Alex Damian, Tengyn Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461, 2021.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810000803, 2018. URL http://arxiv..r org/ org/abs/ abs/ 1810.04800
Gintare Karolina. Dziugaite and Danicl M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedinggs of the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAI), 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent pac-bayes priors via differential privacy. Advances in neural information processing systems, 31, 2018.
Gintare Karolina Dziugaite, Kyle Hsu, Waseem Gharbieh, Gabriel Arpino, and Daniel Roy. On the role of data in pac-bayes bounds. In International Conference on Artificial Intelligence and Statistics, pp. 604-612. PMLR, 2021.
Johannes Gasteiger, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. arXiv preprind 0rXiv:1810.05907, 2018
Michael Gastpar, Ido Nachum, Jonathan Shafer, and Thomas Weinberger. Fantastic generalization measures are nowhere to be found, 2023.
Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavy-ball mo- mentun accelerated stochastic gradient descent. In The Eleventh International Conference on Learning Representations, 2022.
Maxime Haddouche, Benjamin Guedj, Omar Rivasplata, and John Shawe-Taylor. Pac- Pac-bayes unleashed: Generalisation bounds with unbounded losses. Entropy, 23 (31):3333,, 2021.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs Advances in neural information processing systems, 30, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In
Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Ralf Herbrich and Thore Graepel. A bayesian margin bound for linear classifiers: Why svms work. Advances neural information processing systems, 13, 2000.
Matthew Holland, Pac-bayes under potentially heavy tails. Advances in Neural Information Processing Systems, 32, 2019.
17

---

Published in Transactions on Machine Learning Research (09/2024)
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference compuler vision and pallern recognilion, pp- 4700-4708, 2017.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2019.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint ar Xi:::16009000 020,000777 2016.
Tija Kuzborskij and Csaba Szepesvári. Efron-stein pac-bayesian inequalities. arXiv preprint arXiv: 190900111,1 2019.
Sungyoon Lee and Cheongjae Jang. A new characterization of the edge of stability based on a sharpness measure aware of batch gradient distribution. In The Eleventh International Conference on Learning Representations, 2022.
Gaöl Letarte, Pascal Germain, Benjamin Guedj, and François Laviolette. Dichotomize and generalize: Pac-bayesian binary activated deep neural networks. Advances Neural Information Processing Systems, 32, 2019.
Aitor Lewkowyzz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.0220 02218, 2020.
Roi Livni and Shay Moran. A limitation of the pac-bayes framework. Advances in Neural Information Processing Systems, 33:20543-200 -20553,, 2020.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint ar Xrx:::1111.05100 2017.
Ping Luo, Xinjiang Wang, Wenqi Shao, and Zhanglin Peng. Towards understanding regularization in batch normalization. Xive preprint
2018.
Andreas Maurer. A note on the pac bayesian theorem. arXio preprint es/0411099, 2004.
David A McAllester. Some pac-bayesian theorems. In Proceedings of the eleventh annual conference on Computational learning theory, pp. 230 234, 1998.
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pp. 164-170, 1999.
Vaishnavh Nagarajan and J Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning, Advances in Neural Information Processing Systems, 32, 2019.
Arvind Neelakantan, Luke Vilnis, Quoc V Le, Ilya Sutskever, Lukasz Kaiser, Karol Kurach, and James Martens. Adding gradient noise improves learning for very deep networks. arXiv preprint ar #=Xiv:151.006077, 2015.
Antonio Orvieto, Hans Kersting, Frank Proske, Francis Bach, and Aurelien Lucchi. Anticorrelated noise injection for improved generalization. In International Conference on Machine Learning, pp. 17094-17116. PMLR, 2022.
Maria Perez-Ortiz, Omar Rivasplata, Benjamin Guedj, Matthew Gleeson, Jingyu Zhang, John Shawe-Taylon Miroslaw Bober, and Josef Kittler. Learning pac-bayes priors for probabilistic neural networks. arXiv preprint arXiv:2109. 10304, 2021.
María Pérez-Ortiz, Omar Rivasplata, John Shawe-Taylor and Csaba Szepesvári. Tighter risk certificates for neural networks. The Journal of Machine Learning Research, 22(1):10326-103886 2021.
Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvári. Pac-bayes with backprop. arXiv preprint 4rXin:1908.07339. 2019.
18

---

Published in Transactions on Machine Learning Research (09/2024)
Omar Rivasplata, Iija Kuzborskij, Csaba Szepesvári, and John Shawe-Taylor. Pac-bayes analysis beyond the usual bounds. Advances in Neural Informalion Processing Systems, 33:16833-1666 16845, 2020.
Borja Rodriguez- Rodrignez-Gálvez, Gálvez, Ragnar Thobaben, and Mikael Skoglund. More pac-bayes bounds: From bounded losses, to losses with general tail behaviors, anytime-validity. arXiv preprint arXiv:2906.122 12214, 2023.
John Shawe- Shawe-Taylor and Robert C Williamson. A pac analysis of a bayesian estimator. In Proceedings of the tenth annual conference on Computational learning theory, pp. 2 9, 1997.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 15556 2014.
Niklas Thiemann, Christian Igel, Olivier Wintenberger, and Yevgeny Seldin. A strongly quasiconvex pac- bayesian bound. In International Conference Algorithmic Learning Theory, pp. 466-492. PMLR, 2017.
Pelar Veličković, Guillem Cucurull, Aranixa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint 1rXiv:1710.111 10903, 2017.
Paul Viallard, Maxime Haddouche, Umut Simsekli, and Benjamin Guedj. Learning via wasserstein-based high probability generalisation bounds. Advances in Neural Information Processing Systems, 36, 2023.
Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International conference on machine learning, pp. 10181-10192. PMLR, 2020.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P Adams, and Peter Orbanz. Non-vacuous generalization bounds at the imagenct scale: pac-bayesian compression approach. arXiv preprint arXiv:1804.06 058662, 2018.
19

---

Published in Transactions on Machine Learning Research (09/2024)
Appendix
A Proofs
A.1 Proofs of Theorem 4.9
Theorem A.1. Given a prior Px parometrized by X A over the hypothesis set H. Fis X E A, E (0, 1) and
[2) For any choice of i..d m-sized training dataset S according to D, and all posterior distributions
Q over H, we have
1 8 _ (log KT())(P)) +7((A)
1
oh- S En-of(h;
(9)
7m
holds with probability at least 8 when ((h,)) salisfies Definition 4.7 with bound K (()).
Firstly, in the bounded interval 7 E [n, Y2], we bound the difference the expected loss over the posterior distribution evaluated on the training dataset S and with the KL divergence between the posterior distribution Q and prior distribution Px evaluated over a hypothesis space H.
Proof.
For [m, 72),
ES-D (exp (Tm (Fn D) - Rn=o((h;;)) (L((2IPx)))
=DSDD exp (ym - (h: En log dQ
(h)))
(10)
dP>
dº
EEs.DEnnggeen D) ((h: log dP> (h))]
(11)
EEn.:..Es--leeppmm((((11. D) S)))],
(12)
where dQ/dP denotes the Radon- Nidon-Nikodyym derivative.
In Equation 100, (L( En~o log de,
dQ (h). From Equation (10) to Equation 11 Jensen's 's
inequality is used over the convex exponential function. Since this argument holds for any Q, we have
sup EswD(exp (Tm: - (b=ee(h; - eEQ
( (T(2) s En Ds-peexp(((e((( D) e(h; S)))
(13)
Let X = ((h:)) - (); then X centered with E[ = 0. Then, by Definition 4.7
AK(A), E E Resplexp (fmX)) S exp (m²2k(()))
(14)
Using Markov's incquality, Equation 15 holds with probability at least 5.
exp (1mX) exp (m>²K(A)
(15)
8
Combining Equation
13 and Equation
15
the following inequality holds with probability at least 1 - ô.
sup exp (ym(En D) En ((h:)(() KL()) exp 2 K()) eec
>ym(Bn_(((h) D) En-ge(; (L(2 S log K VQ
eEn D) < (n- (log + KL( + VQ.
(16)
7m
The bound 16 exactly the statement of the Theorem.
20

---

Published in Transactions on Machine Learning Research (09/2024)
A.2 Proof of Theorem 4.15
Theorem A.2. Let == el.e) e) be the covering wambee of the set of the prior purometers. Under
Assumption 4.13 and Assumption 4.14 the following inequality holds for the minimizer (f; w. X) of upper with probabilily as least 1 -:
bound in Equation
g
1
n(e) + 12-21
Eh (n (b; & En S); +
log
KL(O PP)) + + +
Fmm
= LPAc 7, 2, +
(17)
holds for any 6,e > O, where n Be C(m(e) + m2(E)) + log(n(c)+ 22
with C
7m 72. and
B = supXEA my (KL(QW) (Px) + log }) + K().
Proof: In this proof, we extend our PAC-Bayes bound with data-independent priors to data-dependent ones that accommodate the error when the prior distribution parameterized and optimized over finite set of parameters Q (Px, À E A C RR] with a much smaller dimension than the model itself. Let T(A, ll II, be an E-cover of the set A, which states that for any À A, there exists Å T(A, II,=). such that | X|| S..
Now we select the posterior distribution as Qw, parameterized by w E Rd. Assuming the prior P is parameterized by E & (k < d).
Then the PAC-Bayes bound 9 holds already for any (w, y, A), with fixed X e A and y E 17,,72l i.c.,
1 (log + (I(Q)(IPA) (((())
e ((h:)) Exre, (h; e + ym
(18)
with probability over 8.
Now, for the collection of As in the s-net I- by the union bound, the PAC-Bayes bound uniformly holds the e-net with probability at least 1 [T 1 n(e).. For arbitrary À A, its distance to the
e-nct is most e. Then under Assumption 4.13 and Assumption 4.14 have:
min |K((2)P>) - (L m - SI S (s),
AET
and
min K(A) K(X) < 12 ÄII) S m2(s).
AET
Similarly, for 7, s-net its range 71 S 72 is the uniform grid with grid separation e, so the net contains 22-11 points. By the union bound, requiring the PAC-Bayes bound to uniformly hold for all the y within this e-net induces an extra probability of failure of 22-28. So, the total probability of failure is n(=)ō + 22
e
For an arbitrary y I, and I to m2,.)) its distance to the e-net T is most E, we have:
KI((@w1PA) + log 1 1) 7
min
[LPAc(W, ",,,,, 7, 8) LrAc(w, 7, 7,A,0) RET'
8
(((()))
(KI PA) + log
K(A)
17 1
myy
8'
1
(KT(Qw + log 5) K(A)
my?
Be,
where B = supxEA m (KL + log !) K (()), clearly, B is a constant depending on the range of the parameters.
21

---

Published in Transactions on Machine Learning Research (09/2024)
With the three inequalities above, can control the PAC-Bayes loss at the given À and y follows:
min LLAAC(W 7, 6) LPAC(W, (w,, 7, AETRET
S min [LAAC(W; 7, A, 5) L + min (PPAc(W;7:20) LPac(W, (,,5, 1EP
λET
1
$ B (1))) (m)
mm
< Be + 7/
?/1 (e) + (2))(e)
S Be C(m + 12(E))
where C= -=
71 + 72 and 71 S S 72- Since this inequality holds for any X A and y E P, it certainly holds for the optima Å and &. Combining this with Equation (88 we have
En-o. D); S (w, f,W, 2, + Be + m()))))((()))
where B: = supxeA mni
1 (2wllwll (P)) + log 31) + K(A).
Now taking e = (n(E)+ 22 )0 to be the previously calculated probability of failure, we get, with probability 6, holds that
Ex ((x;)((h; s ( 1
n(E) + 22-
log
+ KL IP3)) +AK(A) + Be C(n(() 72(=))
3m
S LPAC(W, 6) +
(19)
and the proof is completed.
A.3 KL divergence of the Gaussian prior and posterior
For k-layer network, the prior is written as PA where d4 is the random initialized model parameter and À E I&t is the vector containing the variance for each layer. The set of all such priors is denoted by B = A E A C R ulu E θ}. In the PAC-Bayes training, we select the posterior distribution to be centered around the trained model parameterized by μ,, with independent anisotropic variance. Specifically, for network with d trainable parameters, the postcrior Q,.o := N(p, diag(o)), where μ (the current model) is the mean and R2 is the vector containing the variance for each trainable parameter. The set. of all posteriors is f (Qu.o.0 E. θ}, and the KL divergence between all such prior and posterior in Q and 12 is:
b Σ
- 1 log(o,) + d.(log(A,) - 1) |lo + 11( - mo)al13)
PP = 2 1=1
(20)
A
where oi, (μ- (o); are vectors denoting the variances and weights for the i-th layer, respectively, and A is the scalar variance for the i-th layer. d, = dim(o,), and 1a, denotes an all-ones vector length d
Scalar prior special case of the layerwise prior by setting all entries of À be equal, for which the KL divergence reduces
1 -1 log(or - d(log(A) (A) 1) (1 + 11 p40 KL (T (PPo.>)
(21)
2
.Noee that with little ambignity, the here different meaning from that Equation 26 and Algorithm 2 here the ith element X, whereas Equation 26 and Algorithm 2 X, means the ith element in the discrete sel.
22

---

Published in Transactions on Machine Learning Research (09/2024)
A.4 Proof of Corollary 4.16
Recall for the training, we proposed to optimize over all four variables: μ. 7, 0, and X.
ym (,, = arg min ((fo))) (log + KK(QQ.ol1PPA +,
(22)
n.À.a. 7E(mm
Corollary A.3. Assume all paramelers for the prior and posterior are bounded, i.e., we restrict the model parameter μu, the posterior variance a and the prior variance X, and the exponential moment ({A) all to searched over bounded sets, == {u E R || I VdM}, Σ ; for E 14 ||ol11 & df), {EE 1e e, r {r E hh, Yal,, respectively, with fized M, T, a. b > Then,
Assumption 4.18 holds with mn (x) Lıx, where L1 = max e" (2V dT))
Assumption 4.14 holds with (2(x)) Lyx, where L2 = # (2dM2 + d(a+b)
With high probability, the PAC-Boyes bound for the minimizer of Equation P
has the form
E S LPAc (1,,
where 17 = 11 + log (]CL+B)Aym L I + L2, =
max + a, 2,2 -n)}, B =
1 (KL (ep > + log } + K(),, and C' = 1 + 72.
suP my?
We first prove the two assumptions are satisfied by the Gaussian family with bounded parameter spaces. To prove Assumption 4.13) is satisfied, let v; = log log1/A, i= 1,... & and perform a change of variable from X to u. The weight of prior for the ith layer now becomes M(po, e II, where d, is the number of trainable parameters in the ith layer. It is straightforward compute
Proof:
(KL P PP...) 1 [[-d, dd + t (" + 1| 14. (2)).
-
Ov:
where oú, (4i Mo,i the blocks of a, μ. no: containing the parameters associated with the ith layer, respectively. Now, given the assumptions the boundedness of the parameters, have:
( S 1P.. plo. & max e (2V + dT)) = (,,
(23)
where we used the assumption ||o||| < dl' and |oll2 ||// S VdM.
Equation 23 says L1 M, a) valid Lipschitz bound on the KL divergence and therefore Assumption 4.13 satisfied by setting 7 (I:) La (d,, a,
23

---

Published in Transactions on Machine Learning Research (09/2024)
Next, prove Assumption 4.14 is satisfied. We use Kmin (X) defined in Definition 4.7 the K(A) in the PAC-Bayes training, and verily that it makes Assumption 4.14 hold.
(min(A1) Kmin(A2)
1
sup (
722
log Ba BEND eexp (f((([==))))
sup 7E( 72
Log(BopPoo.a (Bzwolexp (>((Jo;
1
sup rel
llog(E D- AMDAA Ap Bewplexp ((fo: =)]]) log Log(EapPro.. Po. Bewploxp (x((f(====
log Dogg PMP f (ye(Jo: 2)) Puo.2:
Ppo) (0)
sup mal
logg(Bo BEND (>((fo:=)))
72
21 (0)
Auo.A (0)
< sup me(
72
sup 0E
log
AuoiA2 (0)
log Ai (0)
sup 7 heHt
Pu (0)
2dM222 + d(a d(a+00 + 6)
|| - A2112.
2
where the first inequality used the property of the supremum, the Pho.As PHo, (0) in the fourth line denote the probability density function of Gaussian with mean 140 and variance parametrized by Å X2 (i.e., X1, d are the variances for the ith layer), the second inequality use the fact that if X X(h) is a non-negative function of h and Y(h) is a bounded function of h, then
[x (b)) S (sup (Y(h)) E, (h).
heht
The last incquality used the formula of the Gaussian density
E. (2m)d/2/z111/2 exp (-2 -) (0-10--WE-0000- m)
1
and the boundedness of the parameters. Therefore, Assumption 4.14 is satisfied by setting na(x))
(2(d, M, 71, a)x, where L2(d, M, 71,4) (2dM2e20 + Ma+b))
Let L(d, a, (1dd, M,T, + L2(d, M, 71, a). Then we can apply Theorem 4.15 to gel with probability 1
Eswena D(fo:D)
< Ban (Joos)) (();s) + Fm log
1
n(E) +
22 26
+ KL(Q mo.x)) + KmKiin (2)+
(24)
(CL(d, M,, a)) + B)e.
Here, used m)(a) = Tạ and (2(2)= = L2x. Note that for the set [-b, aj8, the covering number n(e) [[[[b, a 1.E) (bta)* and the covering number 12-21 for y E [2. 2s
We introduce new variable p O, letting 2CLaMITa.4) and inserting into Equation 24 we obtain with probability -
Eg (fo:DD)
1
log 1
KTTOOpolPPP... ( 2(CT,d, M, T, a,,,,,,) + B)A
< Ea- (( fm
+ 9 minia (2) ++
log mm
24

---

Published in Transactions on Machine Learning Research (09/2024)
where A = max(b- + a, 2(%/ 7))-
Optimizing over we obtain:
Ea- (
1
KI (Q RO.D))
I Ea-Qa (ffo:S) fm
log
k.
2(CI(d, M, T, 71 B)Armm
+ Kmii (A)
+ log
mmm
k
(1+1og 2(CL(d, M, T, (,,,a) +B) m)
LPAc 0).
7m
Hence we have
Ee Qn. ((oo) S L, +
1 (1 + log(2 M, m, + B) -mmm)), mim (1 + log 2(CL(d.M,II B) yim
where v) = max
k
Remark A.4. In defining the boundedness of the domain of l Corollary 4.16. we used VMM as the bound. Here, the factor vd (where d denotes the dimension of h) used encapsulate the idca that if on average, the components of the weight bounded by M, then the lg norm would naturally be bounded by VdM. The same idea applies to the definition of E.
Remark A.5. Due to the above remark, M, T, can be treated as dimensionn-nndeeendeen constants that do not grow with the network size d. As a result, the constants L1, L2, L in Corollary 4.16 are dominated by d, and Iı; L2 (d)) This then implics the logarithm term in 1 scales as Ollog d), which grows very mildly with the size. Therefore, Corollary 4.16 can be used as the generalization guarantee for large neural networks.
B Algorithm Details
B.1 Algorithms to estimate K(A)
In this section, we explain the algorithm to compute K(A). In previous literature, the moment bound K or its analog term in the PAC-Bayes bounds was often assumed to be a constant. One of our contributions is to allow to vary with the variance À the prior, so a small prior variance is found by PAC-Bayes training, would also be small. We perform linear interpolation to approximate the function Kmin (A) defined in (2) of the main text. When À is 1D, We first compute Kmin (A) on finite grid of the domain of X, by solving Equation (25 below. With the computed function values the grid {Kmiu (A:)):
then the corresponding
can construct piecewise linear function as the approximation of Kmin(A).
Kmin (A:) arg min K
K>0
s.t. exp ())
Σ exp(:(e(fo.:S) - e(fo:;z)))
(25)
nm
1=
Yyem,
0 N Xt); Amin S S Amax
where 0₁ P P / 1, ..., n, are samples from the prior distribution and are fixed when solving Equation 25 for Kmin (2,). Equation (25) is the discrete version of the formula (2) in the main text. This optimization
problem is 1-dimensional, and the function in the constraint is monotonic in K, so it can be solved efficiently by the bisection method.
25

---

Published in Transactions on Machine Learning Research (09/2024)
When extending this procedure to high dimension, where À is k-dimension vector, need to set up grid for the domain of À in k-dimensional space and estimate min each grid point, which is time-consuming when & is large. To address this issue, we propose to use the following approximation:
K (max(A:)) arg min K K>0
n m
s.t. exp ( (T²K) 2
exp()(((Jo;:)) ((oo;
(26)
nm 1=1
A E 72).
0, ~ N(po max (x:),, Amin I max(A;) S Amax, i = 1,...,8
where X random sample from the domain A of X. Since each X is k-dimensional, max(A²) represents the maximum of the coordinates.
The idea of this formulation 26 is as follows, we use the ID function (mmxx(A:) surrogate function of the original k-dimension function (A) (i.e. Kmin (A) s K (max(A,))). Then estimating this ID surrogate function easy by using the bisection method. This procedure will certainly overestimate the true Kiin (A) but since the surrogate function is also a valid exponential moment bound, it is safe to be used as a replacement for the (A) in our PAC-Bayes bound for training. In practice, we tried to use mean(A:) replace max(A,) to mitigate the over-estimation, but the final performance stays the same. The details of the whole procedure are presented in Algorithm
Algorithm Compute (()) given set of query priors
Input: Y1 and Y2, sampling time of prior variances, the initial neural network weight Oo, the training
datasct S = {2 1m model sampling time 10
Output: the piece-wise linear interpolation ((A) for K (m)n(A)
Draw random samples for the prior variances V = (A; E C R= 1.....)
Set up a discrete grid I for the interval [71, 72] of 7.
for E do for 1: do Sampling weights from the Gaussian distribution O N X,) Use 0,, I and S compute one term in the sum in Equation 266
end for Solve (max(A:)) using Equation end for Fit a piece-wise linear function K ()) to the data ((A., (X:))
26
B.2 PAC-Bayes Training with layerwise prior
Similar to Algorithm D our PAC-Bayes training with layerwise prior is stated here in Algorithm
B.3
Regularizations in PAC-Bayes bound
Only noise injection and weight decay are essential from our derived PAC-Bayes bound. Since many factors in normal training, such as mini-batch and dropout, enhance generalization by some sort of noise injection, is unsurprising that they can be substituted by the well-calibrated noise injection in PAC-Bayes training. Like most commonly used implicit regularizations (large lr, momentum, small batch size), dropout and batch-norm are also known to penalize the loss function's sharpness indirectly. Wei et al. (2020) studies that dropout introduces an explicit regularization that penalizes sharpness and an implicit regularization that is analogous to the effect of stochasticity in small mini-batch stochastic gradient descent. Similarly, it is well-studied that batch-norm Luo et al. (2018) allows the use of a large learning rate by reducing the variance in the layer batches, and large allowable learning rates regularize sharpness through the edge of stability [Cohen] ct al. (2020). As shown in the equation below, the first term (noise-injection) in our PAC-Bayes bound explicitly penalizes the Trace of the Hessian of the loss which directly relates to sharpness and is quite similar
26

---

Published in Transactions on Machine Learning Research (09/2024)
Algorithm PAC-Bayes training (layerwise prior)
Input: imitial weight ps E RR, the number of layers k, Tí, A1 = e Tı, 11A1,,2 71 12 can be fixee in all experiments of Sec Output: trained model f, posterior noise level ā 1
0-12,,0 12 = =e² 71 = 0.5, 72 = 10, //
0 Mo: V la log(3 d Di=1 (o,s)), b 1s log Di=1 (0.s
// Initialization
Obtain the estimated K (A) with A = [11, dz) using Equation (66 and Appendix B.1
// Slage
for epoch : do
for sampling one batch s from do
// Ensure non-negative variances Construct the covariance of PooaA from X // Setting the voriance of the weights in layer-i all to the
À A exp(b), + exp(v)
scalar A(i)
Draw one 0 ~ Qu. and evaluate {(gg:))
// Stochaslic version of E {JoiS)
Compute the KL-divergence Equation 20
Compute as Equation (
Compute the loss function L as IPAC in Equation ac
b- + 7 V v V 17 Bv μ 4- 4 + Tp
1/ Update all parameters
end for end for à f- exp(v) // Stage 2
1/ Fix the noise level from
while not converge do for sampling one batch s from S do
Draw one sample ~ Qp and evaluate E(f;: as Ê ac
// Noise injection 11 Update model parameters
| μ + 1
end for end while p +
to the regularization effect of batch-norm and dropout. During training, suppose the current postcrior is 2n. =Na diag(o)), then the training loss expectation over the posterior is:
E ((fe:D) = =Ese~eo. (for
= ((ff, D) + Esa~ (Es00cos((((f:)))))) (l(f + 1 + A0 ( = =((fo:D) + [T
The second regularization term (weight decay) in the bound additionally ensures that the minimizer found is close to initialization. Although the relation of this regularizer to sharpness is not very clear, empirical results suggest that weight decay may have separate regularization effect from sharpness. In brief, we state that the effect of sharpness regularization from dropout and batch norm can also be well emulated by noise injection with the additional effect of weight decay.
B.4 Deterministic Prediction
Recall that for any μ E Rd and RR we used Qu.or to denote the multivariate normal distribition with y and covariance matrix diag(a). If rewrite the left-hand side of the PAC-Bayes bound by Taylor
27

---

Published in Transactions on Machine Learning Research (09/2024)
expansion, have:
2 (so:D) =Rsa~Q0, ".00000 ([o+00:
1 ~ ~1[Jg; D) Bao-go (V + -2200 > (e(l[g:D)A0)
(27)
=(ff: D) + 1 TT(diag(()
((fo: D).
Recall here á and & are the minimizers of the PAC-Bayes loss, obtained by solving the optimization problem Equation Equation Equation (27 states that the deterministic predictor has smaller prediction error than the Bayesian predictor. However, note that the last inequality in Equation (77 is derived under the assumption that the term {Ze D) is positive-scnne semiieeiiemmmi.iiit.... This is a reasonable assumption as fa is the local minimizer of the PAC-Bayes loss, and the PAC-Bayes loss is close to the population loss when the number of samples is large. Nevertheless, since this property only approximately holds, the presented argument can only serve as an intuition that shows the potential benefits of using the deterministie predictor.
C Extended Experimental Details
We conducted experiments using eight A5000 GPUs with four AMD EPYC 7543 32-core Processors. To speed up the training process for posterior and prior variance, we utilized a warmup method that involved updating the noise level in the posterior of each layer as a scalar for the first 50 epochs and then proceeding with normal updates after the warmup period. This method only affects the convergence speed, not the generalization, and was only used for large models in image classification.
C.1 Parameter Settings
Rocall that the exponential momentum bound ({A) is estimated over a range [n, 72l as per Definition 4.7 It means that we need the inequality
En-P, Elexp (r(E[X(()) - X(b))) - exp (?))
hold for any y in this range. One ncods to be little cautious when choosing the upper bound 72, because if it is too large, then the empirical estimale of En-PAE Elexp (>(E[X(h) X (h)))] would have too large of a variance. Therefore,
recommended "/2 to be set to no more than 10 or 20. The choice of % also does not seem to be very crucial, so we have fixed it to 0.5 throughout.
For large datasets (like in MNIST or CIFAR10), m is large. Then, according to Theorem 4.15 we can set the range M, a, b of the trainable parameters to be very large with only a little increase of the bound (as M,T, a, b are inside the logarithm), and then during training, the parameters would not exceed these bounds even if we don't clip them. Hence, no clipping is needed for very large networks or with small networks with proper initializations. But when the dataset size m is small, or the initialization not good enough, then the correction term could be large, and clipping will be needed.
The clipping is also needed from the usual numerical stability point of view. As X is in the denominator of the KI.-divergence, it camnot be too close to O. Because of this, in the numerical experiments on GNN and CNN13/ CNN13/CNNI5 CNN15, we clip the domain of À at a lower bound of 0.1 and 5e - 3, respectively. For the VGG and Resnet experiments, the clipping À is optional.
C.2 Baseline PAC-Bayes bounds for unbounded loss functions
We compared two baseline PAC-Bayes bounds when training CNNs with our layerwise PAC-Bayes bound. The bounds are expressed in our notation. Consider neural network model denoted as Se. where represents the network's architecture, and 0 is the weight.
sub-Gaussian (Theorem 5.1 of Alquier 2021
1 log 5
+ +KT(Q11))) + Kaudy
& Eawel(J:: my
(28)
28

---

Published in Transactions on Machine Learning Research (09/2024)
where Keub is the variance factor by assuming the loss function E is sub-Gaussian defined below: pS exp (x(e(fo:)) - E(fo: D))1 - exp Ks2 vy E R.
In the experiment to generate Figure 1, we restricted y to [-1,1]. Otherwise, if we use E R, the bound would be too large to calculate, resulting in NaN value. Therefore, the reported Kud in Figure 1 may be underestimated.
CGF (Theorem 9 of Gálvez el al. (0233
1 1
1 + KT(2IPP) + + 5
E of S Eo +
(log ? m
(29)
where (1(1) is a convex and continuousy difforentiable function defined on [0, b) for some b E R+ such that (0) v0()) = 0 and (e--pEss_leep((((((((( D) - E(So [))]] S exp (1(7)) for all y E [0. b).
There is no specific form of (T)) provided in the original paper, but in order to achieve the normal v'm convergence rate of the PAC-Bayos bound, we need at least set v(()) = KcaFy" (a > 2). Among these, using = 2 gives the smallest K, so used = 2 for the comparison in Figure I and 3
C.3 Compatibility with Data Augmentation
We didn't include data augmentation in the experiments in the main text. Because with data augmentation, there is no rigorous way of choosing the sample size m that appears in the PAC-Bayes bound. More specifically, for the PAC-Bayes bound to be valid, the training data has to be i.i.d. samples from some underlying distribution. However, most data augmentation Lechniques would break the i.i.d. assumption. As result, if we have 10 times more samples after augmentation, the new information they bring in would be much less than those from times i.i.d. samples. In this case, how to determine the effective sample size m to be usod in the PAC-Bayes bound is a problem.
Since knowing whether training method can work well with data augmentation is important, carried out the PAC-Bayes training with an ad-hoc choice of m, that is, we set m to be the size of the augmented data. We compared the grid-search result of SGD and Adam versus PAC-Bayes training on CIFAR10 with ResNet18. The augmentation is achieved by random flipping and random cropping. The data augmentation increased the size of the training sample by 128 times. The test accuracy for SGD 95.2%, it is 94.3% for Adam, it 94.4% for AdamW, and is 94.3% for PAC-Bayes training with the layerwise prior. In contrast, the test. accuracy without data augmentation is lower than 90% for all methods. It suggests that. data augmentation does not conflict with the PAC-Bayes training in practice.
C.4 Model analysis
We examined the learning process of PAC-Bayes training by analyzing the posterior variance for different layers in models trained by Algorithm 3 Typically, batch norm layers have smaller values than convolution layers. Additionally, shadow convolution and the last few layers have smaller or values than the middle layers. We also found that skip-connections in ResNet18 have smaller a values than nearby layers, suggesting that important layers with a greater impact on the output have smaller o values.
In Stage 1, the training loss higher than the testing loss, which means the adopted PAC-Bayes bound is able to bound the generalization error throughout the PAC-Bayes training stage. Additionally, we observed that the final value of K is usually very close to the minimun of the sampled function values. The average value of experienced rapid update during the initial 50 warmup epochs but later progressed slowly until Stage 2. The details can be found in Figure g and TO Based on the figures, shadow convolution, and the last few layers have smaller or values than the middle layers for all models. We also found that skip-connections in ResNet18 and ResNet34 have smaller a values than nearby layers on both datasets, suggesting that important layers with a greater impact on the output have smaller o values.
Computational cost: In PAC-Bayes training, we have four parameters μ. X. 0.7% y. Among these variables, Y can be computed the fly whenever needed, there no need to store them. We need to store μr, X.
29

---

Published in Transactions on Machine Learning Research (09/2024)
where a has the same size as i and the size of X is the same the number of layers which is much smaller. Hence the total storage is approxinately doubled. Likewise, when computing the gradient for μ. X, a, the cost of automatic differentiation in each iteration is also approximately doubled. In the inference stage, the complexity is the same as in conventional training.
Effect of two stages: We have tested the effect of the two stages. Without the first stage, the algorithm cannot automatically learn the noise level and weight decay to be used in the second stage. If the first stage is there but too short (10 epochs for example), then the final performance of VGG13 CIFAR100 will reduce to 64.0%. Without Stage 2, the final performance is not good reported either. The test accuracy of models like VGG13 and RosNet18 on CIFAR10 would be 10% lower as in Figure and 10
C.5 Node classification by GNNs
We test the PAC-Bayes training algorithm on the following popular GNN modcls, tuning the learning rate (1e-3, 5, le-2), weight decay (0, 1e-4, 1e-3, 1e-2), noise injection (0. 3, 5e-3, le-2), and dropout (0, 0.4, 0.8). The number of filters per layer is 32 in GCN (Kipf & Welling 2016) and SAGE (Hamilton et al. 2017). For GAT (Veličković et al. 2017), the number of filters is 8 per layer, the number of heads is 8, and the dropout rate of the attention coefficient is 0.6. Fpr APPNP (Gasteiger et al. 2018), the number of filters is 32, K 10 and a = 0.1. We set the number of layers to 2, achieving the best baseline performance. A ReLU activation and a dropout layer are added between the convolution layers for baseline training only. Since GNNs are faster to train than convolutional neural networks, we tested all possible combinations of the above parameters for the baseline, conducting 144 searches per model on one dataset. We use Adam as the optimizer with the learning rate as 1e-2 for all models using both training and validation nodes for PAC-Bayes training.
We also did a separate experiment using both training and validation nodes for training. For baselines, we nced first to train the model to detect the best hyperparameters as before and then train the model again on the combined data. Our PAC-Bayes training can also match the best generalization of baselines in this setting.
All results are visualized in Figure g The AdamW+val and + record the performances of the baseline and the PAC-Bayes training, respectively, with both training and validation datasets for training. We can see that test accuracy after adding validation nodes increased significantly for both methods but still, the results of our algorithm match the best test accuracy of baselines. Our proposed PAC-Bayes training with the scalar prior is better than most of the settings during searching and achieved comparable test accuracy when adding validation nodes to training.
C.6 Few-shot text classification with transformers
The proposed method is also observed to work on transformer networks. We conducted experiments on two text classification tasks of the GLUE benchmark as shown in Table ( SST is the sentiment analysis task, whose performance is evaluated as the classification accuracy. Sentiment analysis is the process of analyzing the sentiment of given text to determine if the emotional tone of the text is positive, negative, or neutral. QNLI (Question-answering Natural Language Inference) focuses on determining the logical relationship between a given question and a corresponding sentence. The objective of QNLI is to determine whether the sentence contradicts, entails, or is neutral with respect to the question.
We use classification accuracy as the evaluation metric. The bascline method uses grid scarch over the hyper-parameter choices of the learning rate (1e-1, 1e-2, le-3), batch size (2,8, 8, 16, 32, 80), dropout ratio (0, 0.5), optimization algorithms (SGD, AdamW), noise injection (0, le-5, 1e-4, 1e-3, 1e-2, 1e-1), and weight decay (0, 1ellee,eel-1 1e-3, 1e-4). The learning rate and batch size of our method are set to le-3 and 100 (i.e., full-batch), respectively. In this task, the number of training samples is small (80). As a result, the preset 72- = 10 is a bit large and thus prevents the model from achieving the best performance with PAC-Bayes training.
We adopt BERT (Devlin et al. 2018 as our backbone and added one fully connected layer as the classification layer. Only the added classification layer is trainable, and the pre-trained model is frozen without gradient.
30

---

89.2 Published in Transactions on Machine Learning Research (09/2024)
update. To simulate few-shot learning scenario, we randomly sample 100 instances from the original training set and take the whole development set to evaluate the classification performance. We split the training set into 5 splits, taking one split as the validation data and the rest as the training set. Each experiment was conducted five times, and we report the average performance. We used the PAC-Bayes training with the scalar prior in this experiment. According to Table [ our method is compelitive to the baseline method on the SST task, and the performance gap only 0.4 points. On the QNLI task, our method outperforms the baseline by a large margin, and the variance of our proposed method is less than that of the baseline method.
Table 6: Test accuracy on the development sets of 2 GLUE benchmarks.
SST
QNLI
baseline
0..9
62.6+0011 72.5::0.99
scalar
64.2:10.02
C.7 Additional experiments stability
We conducted extra experiments to showcase the robustness of the proposed PAC-Bayes training algorithm. Specifically, we tested the effect of different learning rales on ResNet18 and VGG13 models trained with layerwise prior. Learning rate has long been known as an important impact factor of the generalization for baseline training. Within the stability range of gradient descent, the larger the learning rate is, the better the generalization has been observed (Lewkowyzz al. 2020). In contrast, the generalization of the PAC-Bayes trained model is less sensitive to the learning rate. We do observe that due to the newly introduced noise parameters, the stability of the optimization gets worse, which in turn requires lower learning rate achieve stable training. But as long the stability is guaranteed by sctting the learning rate low enough, our results, as Table 7 indicated that the test accuracy remained stable across various learning rates for VGG13 and Resnet The dash in the table means that the learning rate for that particular setting is too large to maintain the training stability. For learning rates below le-4, we trained the model in Stage 1 for more epochs (700) to fully update the prior and posterior variance.
We also demonstrate that the warmup iterations (as discussed at the beginning of this section) do not affect generalization. As shown in Table g the test accuracy is insensitive to different numbers of warmup iterations. Furthermore, additional evaluations of the effects of batch size (Table 10), optimizer (Tables 11 and 21 and 22 (Table 12
Table 7: Test accuracy of ResNet18 trained with different learning rates.
3e-5
56-5
le-4
2
3e-4
5e 4
CIFAR10
88.4
88.8
89.3
88.6
88.3
CIFAR100
69.2
69.0
68.9
69.1
69.1
69.6
Table Test accuracy of VGG13 trained with different learning rates.
3e-5
5e-5
le le-4
2e-4
3e
4
5e 5e-4
CIFAR10
88.6
88.9
89.7
89.6
89.6
89.5
CIFAR100
67.7
68.0
67.1
31

---

Published in Transactions on Machine Learning Research (09/2024)
Table 9: Test accuracy of ResNet18 trained with warmup epochs of
10
20
80
100
150
CIFAR10
88.5
88.5
89.3
89.5
89.5
88.9
CIFAR100
69.4
69.6
68.9
69.1
69.0
68.1
Table 10:
Test accuracy
VGG13 with different batch sizes.
Batch Size
128
256
1024
2048
2500
Test Acc
89.7
89.7
88.7
89.4
88.3
Table 11: Test accuracy of ResNet18 using SGD: Effects of different momentum values (with learning rate x 10-3) and different learning rates (with momentum 0.9).
Momentum
T Learning Rate
0.3
0.6
0.9
10-4
10-4
1 10-3
Test Acc 88.6
88.8
89.2
88.3
88.8
89.2
Table 12: Test accuracy ResNet18 with different settings for 7 (with 72 = 20) and 72 (with 71 = 0.1).
?/2
0.1
0.5
1.0
10
15
20
Test Acc 88.8
89.3
88.8
89.3
89.4
89.4
32

---

Published in Transactions on Machine Learning Research (09/2024)
0.87
0.62
0.90
0.86
0.85
0.61
accuracy
0.84
AdamW
accuracy
AdamW
accuracy
AdamW
scalar
scalar
0.60
scalar
0.83
scalar+val
scalar+val
scalar+val
AdamW+val
AdamW+val
AdamW+val
0.82
0.84
0.59
0.81
0.82
0.58
0.80
50
100
150
0
100
150
index
index
(a) CoraML
(b) Citeseer
(c) CoraFull
0.86
0.82
0.80
0.84
0.78
accuracy
0.82
AdamW
scalar
accuraci
0.76
AdamW
scalar
scalar+val
0.74
scalar
0.80
AdamW+vval
AdamW+val
0.72
0.78
0.70
0.68
0
50
100
150
50
100
150
index
index
(d) DBLP
(e) DBLP
Figure 5: Test accuracy of GCN. The first and third quartiles construct the interval over the ten random splits. {+val} denotes the performance with both training and validation datasets for training.
33

---

Published in Transactions on Machine Learning Research (09/2024)
0.64
0.92
AdamW
AdamW
0.86
0.90
scalar
0.62
scalar
accuracy
0.82
0.84
scalar
AdamW
accuracy
0.86
0.88
scalar+val
AdamW+val
accuracy
0.58
0.60
scalar+vvll
AvamW+val
vcalar+val
0.80
AdamW+val
0.84
0.54
0.82
0.78
0.52
0.80
0.50
0.76
50
150
100
150
100
150
index
index
index
(a) CoraML
(b) Citeseer
(c) CoraFull
0.86
0.825
AdamW
0.84
0.800
scalar
0.775
scalar+val
0.82
accuracy
AdamW
accuracy
0.750
AdamW+va
0.80
scalar
0.725
vaa
0.78
AvamW+val
0.700
0.76
0.675
0.74
0.650
0.625
0.72
0
50
100
150
0
50
100
150
index
index
(d) DBLP
(e) DBLP
Figure 6: Test accuracy of SAGE. The first and third quartiles construct the interval over the ten random splits. {+val} denotes the performance with both training and validation datasets for training.
34

---

Published in Transactions on Machine Learning Research (09/2024)
0.88
AdamW
0.92
0.66
AdamW
scalar
scalar
0.86
scalar+val
0.90
scalar+val
accuracy
AdamW+val
accuracy
AdamW
AdamW+val
accuracy
scalar
0.84
0.88
scalar+val
0.62
AdamW+val
0.82
0.86
0.60
0.80
0.84
0.58
50
100
150
100
150
index
index
index
(a) CoraML
(b) Citeseer
(c) CoraFull
0.86
AdamW
0.82
0.84
scalar
scalar+vvl
0.80
accuracy
0.82
AdamW+vvl
accuracy
AdamW scalar
0.78
0.80
scalar+val
0.76
0.78
AdamW+val
0.74
0.76
0.74
0
50
100
150
50
100
150
index
index
(d) DBLP
(e) DBLP
Figure 7: Test accuracy of GAT. The first and third quartiles construct the interval over the ten random splits. {+val} denotes the performance with both training and validation datasets for training.
35

---

Published in Transactions on Machine Learning Research (09/2024)
0.64
0.90
AdamW
0.85
0.85
scalar+val
scalar
0.62
0.60
accuracy
0.80
accuracy
AdamW+val
accuracy
0.58
0.75
0.56
0.75
AdamW
0.70
AdamW
scalar
0.54
0.65
0.70
scalar+val
scalar+val
AdamW+val
0.60
AdamW+val
0.50
50
100
150
100
150
50
100
150
index
index
index
CoraML
(b) Citeseer
(c) CoraFull
0.86
AdamW
AdamW
0.84
scalar
0.80
scalar
scalar+val
scalar+val
accuracy
0.82
AdamW+val
accuracy
0.75
Ava
0.80
0.78
0.76
0.74
0.60
0.72
0.70
50
100
150
100
150
Index
index
(d) DBLP
(e) DBLP
Figure 8: Test accuracy APPNP. The first and third quartiles construct the interval over the ten random splits. {+val} denotes the performance with both training and validation datasets for training.
1e-6
Ie-3
336
0.8
-3.8
0.4
4.4
0.0
4.6
0
10
20
30
40
5
10
15
20
200
400
600
weight/blas index of batchnorm
weight index of conv layer
Epoch
mean(or of batch-norm layers. (b) mean(ar of convolution layers.
(c) mean in training.
2.0
10a
2.0
Accuracy
Loss
0.7
train loss
10a
test loss
0.6
1.0
train acc
test acc
10-2
10-1
10a
query 1/2
200
400
Epoch
(d) Function K (2).
(e) Training and testing process.
Figure 9: Training details of ResNet18 on CIFAR10. The red star denotes the final K.
36

---

Published in Transactions on Machine Learning Research (09/2024)
1e-3
-34
10-3
-3.6
-3.8
4.0
10-4
4.2
10-5
weight/bias index batchnorm
weight index conv layers
Epoch
(a) mean(or) of batch-norm layers.
(b) mean(or of convolution layers.
(c) mean(o)
in training.
2.5
train loss
test loss
0.9
2.0
train acc
0.8
test acc
1.5
0.6
10a
1.0
0.5
10-2
10-1
109
0.5
query 1/2
200
400
600
Epoch
(d) Function ((X).
(e) Training and testing process.
Figure 10: Training details of VGG13 on CIFAR10. The red star denotes the final K.
37


Reviews:
[{"review_id": "jB9N5nzBrV", "reviewer": "Reviewer_Gcra", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n1. The authors provide a new PAC-Bayes algorithm with a relaxed assumption.\n2. The new algorithm outperforms other selected PAC-Bayes baselines with mean predictors and achieves results comparable to the best in test accuracy.\n\nWeaknesses: \nIt is interesting to explore the PAC-Bayes training objectives. If only the test accuracy is evaluated, it is unclear why this new algorithm, which considers both prior and posterior training and involves more hyperparameters than ERM, is preferable. It would be better to show how accuracy changes over time or to provide additional evaluation metrics.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a new PAC-Bayes algorithm with theoretical analysis. They consider unbounded loss functions with a relaxed exponential moment assumption for hypotheses. Additionally, they propose an optional second stage of Bayesian training to further increase training accuracy. They conduct experiments on several datasets using CNNs and GNNs and compare the results with other baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The new algorithm is compared with other selected PAC-Bayes baselines. It would be better to provide some additional comparisons using the evaluation metrics considered by these baselines.", "broader_impact_concerns": "N. A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3cI65tYgiq", "reviewer": "Reviewer_4s6g", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "__Strengths:__\n\nI would say the strongest point of this paper is that the authors have developed a practical procedure that seems to work quite well on a variety of datasets and neural network architectures. Their basic motivations for designing this procedure are explained clearly, and rooted (partially) in solid theory. In addition, overall, the paper is quite well written, easy to read, and does a pretty good job of describing previous research and its relation to the present work.\n\n\n__Weaknesses:__\n\nOn the other hand, the presentation of some technical points is rather sloppy, and most importantly (in terms of review standards at TMLR) I feel like the main claims and relevant evidence in this work are not sufficiently clear, leaving readers wondering \"what have I learned?\" after reading this paper. Yes, empirical tests suggest that the authors proposed procedure is \"good\" when our definition of \"good\" is \"behaves like a well-tuned ERM algorithm\", but how this relates to the novel content in section 4 is not in my opinion clear at all, and the main claims stated explicitly by the authors are not well-organized, and scattered throughout the paper. I will address these points in a bit more detail in the next section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "To start, I will briefly summarize by understanding of the PAC-Bayes literature which acts as a broad context for this paper. Typical PAC-Bayes bounds for machine learning tasks are formulated for \"randomized\" learning algorithms, i.e., we have a hypothesis class (a set of model candidates, e.g., a set of classifiers) over which we consider so-called \"prior\" and \"posterior\" distributions; the candidate selected by the posterior distribution is the \"output\" of a randomized learning algorithm. As in standard statistical learning theory, we have a loss function depending on a (model, data) pair, and the ultimate objective of interest is the expected loss incurred by the aforementioned learning algorithm output. Here \"expected loss\" includes taking expectation with respect to a new test point as well as expectation over the posterior; this is sometimes referred to as the \"Gibbs risk\" if I recall correctly. One important quality of PAC-Bayes bounds is that they hold with high probability (over training data), uniform in the choice of posterior distribution, so the posterior is allowed to depend on training data.\n\nThe notion of \"PAC-Bayes training\" plays a central role in the paper being reviewed; as just mentioned, the choice of posterior can depend on training data, and thus in principle it can be \"optimized\" based on various empirical criteria. Indeed, much of the past PAC-Bayes literature provides (bound, algorithm) pairs as their main results. Under various assumptions new bounds are derived that are completely (or partially) determined by data available at training time, and algorithms are derived as procedures that minmize those bounds in terms of the posterior (either exactly or approximately).\n\nFrom what I gather, the main interest of the authors, in terms of models, is neural networks (for classification). Their main motivation overall is to design PAC-Bayes training algorithms which (without much or any tuning) perform just as well as traditional empirical risk minimization (ERM) with lots of tuning. Of course, plenty of past research has looked at neural network PAC-Bayes training with the same goals. The main technical points which the authors emphasize differentiate this work from existing literature are as follows: losses are allowed to be unbounded, bounds to be optimized are completely empirical, and the *prior* distribution can also be data-dependent without having to split the data into two disjoint subsets. In addition, they provide a practical training algorithm assuming the distribution over the hypothesis class is Gaussian, and empirically evaluate their procedure, which they find to be competitive.\n\nTheir main result in Thm 4.9 is obtained by assuming a Lipschitz property on the KL divergence and the prior-dependent exponential moment bound.\n\nTheir proposed algorithm is a \"two-stage\" procedure which is not purely PAC-Bayes training, but rather does typical PAC-Bayes training on their objective, and then fixes everything but the mean of the posterior and does essentially ERM on the Gibbs empirical risk. My overall impression is that this practical modification is what leads their proposal to get closer to well-tuned ERM performance.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Here I will highlight some points I found confusing or troublesome. I'm not worried about novelty, but the exposition and main claims have some rather significant issues of clarity in my opinion.\n\n\n__Exposition of the exponential moment approach__\n\nA lot of space is dedicated to the exponential moment bound approach starting in Defn 4.1 (and going through Lemma 4.2 to Defn 4.4), but I found the overall logic a bit hard to parse, and I feel like other readers would as well. The authors claim that their exponential moment condition is \"weaker than the 2nd-order moment condition\" (page 4). I get that *if* we only need one-sided bounds (i.e., for $\\\\mathbb{E}[X]-X$ but not for $X-\\\\mathbb{E}[X]$) and $X \\\\geq 0$, then the exponential bound is weaker than the second moment bound, assuming $\\\\gamma$ is small enough. That said, for newcomers to PAC-Bayes who read \"Unbounded Loss\" in the title, one expects they would assume that control of $X-\\\\mathbf{E}[X]$ will be important, i.e., the right-hand tails of the loss distribution will cause problems for analysis. If this is not the case, and a one-sided approach is sufficient (even if say the right tails are heavy), some exposition of *why* would be meaningful.\n\nNext, the authors also state that the main motivation for Defn 4.1 is to get \"a bound with a smaller numerical value.\" This is unclear. A smaller numerical value when compared with *what?* In the proof of Lemma 4.2, the authors mention how setting $K$ to the larger of $\\\\text{Var}(X)$ or $(\\\\mathbb{E}X)^{2}$ is sufficient for the condition in Defn 4.1 to hold; are these values supposed to be compared with $\\\\mathbb{E}X^{2}$, which is no smaller than either of the previous two quantities?\n\n\n__Claims regarding empirical bounds__\n\nIn Remark 4.6, the authors emphasize that their bound *\"seems to be the first purely empirical bound (i.e., computable from data) that can be easily used for training.\"* I take this is to be a key claim for this paper, as it corresponds as a direct answer to the limitation highlighted in previous literature in the last sentence of the first paragraph of section 3 (*\"... all suffer from this issue.\"*). The latter part of this claim (\"can easily be used for training\") is unclear, but more importantly, upon what do the authors base their claim that the bound in Thm 4.5 is \"purely empirical\"? As I read it from Defn 4.4, the key new quantity $K(\\\\boldsymbol{\\\\lambda})$ depends critically on the population distribution. Maybe I'm missing something, but this is a key point that needs clarification.\n\n\n__Unclear main claims__\n\nConsidering the [acceptance criteria](https://jmlr.org/tmlr/acceptance-criteria.html) for TMLR, the most important point is that submissions need to have claims that are supported by accurate, convincing, and clear evidence. For evidence to be meaningful, the claims that the evidence supports must of course be clear, but I feel this point is lacking in the current submission. I'll try and be concise by placing my focus on the summary of contributions at the end of section 1. Some of these points are really vague and not well-rooted in the substantive content of sections 4-6. A bit more detail:\n\n- Point 1 (*\"We introduce a new...\"*): Yes, the authors do introduce a new bound, but how does it compare to previous bounds? In sections 3-4 one gets the impression that smaller bounds that are purely empirical is a key point, but this isn't touched on in the contribution summary. Of course, the evidence is also unclear, as I pointed out earlier. The second sentence in point 1 (*\"This algorithm simultaneously optimizes...\"*) is very misleading. If the authors are referring to the algorithm in section 5, the second stage is in my opinion a critical point, yet it is glossed over, making it seem like their practical algorithm is just a pure PAC-Bayes bound minimizer.\n\n- Point 2 (*\"The test performance of the proposed algorithm is theoretically justified.\"*\"): this is a totally ethereal statement, really impossible to prove or disprove, and should not be positioned as a main claim. Furthermore, even for the casual reader, intuitively, since stage 2 of Algorithm 1 is a big heuristic modification, it doesn't really jive with this vague claim.\n\n- Points 3, 4, 5: these points are quite lengthy, and really should be summarized into a single point that discusses the actual proposed algorithm and its strengths/weaknesses found in the empirical tests.\n\n\n__Technical points can be quite sloppy__\n\nThere are numerous points I found troublesome related to notation and technical exposition. Here is a handful of representative examples.\n\n- Page 5: parameterizing the posterior distribution, the authors immediately dive into parameterization with mean and variance vectors $\\\\mathbf{h}$ and $\\\\boldsymbol{\\\\sigma}$, but this is quite unnatural. At this stage in the paper, the set of distributions from which the posterior is taken need not be Gaussian, yet the formulation already essentially assumes the Gaussian special case that comes up in section 5. I personally feel this kind of inconsistency in generality can spoil the exposition. The content in section 4 in no way depends on the distribution family being Gaussian.\n\n- Page 5 (and all over the place): the poor usage of $\\\\mathbf{h}$ and $\\\\theta$ really clutters the notation and makes things inconsistent. For example, in section 2, $\\\\mathbf{h}$ is a classifier which returns a *\"predicted label\"*. This is consistent with the text before Defn 4.4, where $\\\\mathbf{h} = f\\_{\\\\theta}$, but does not match up with the last paragraph of page 5, where $\\\\mathbf{h} \\\\in \\\\mathbb{R}^{d}$ and $\\\\mathbf{h}$ has been re-appropriated to do what really should be done by $\\\\theta$ parameterizing the family $\\\\{f\\_{\\\\theta}: \\\\theta \\\\in \\\\mathbb{R}^{d}\\\\}$. The authors make this switch in section 5, but this really should be cleaned up a bit.\n\n- Thm 4.9: the bound holds *\"with probability at least $1-\\\\epsilon$\"*, and yet $\\\\epsilon$ appears nowhere. Presumably it is buried within $L\\_{PAC}$, whose original definition uses $\\\\delta$ rather than $\\\\epsilon$.\n\n- Thm 4.9: $C$ is said to be a \"constant\" but it can also depend on $\\\\eta\\_{2}$; this hides the true order of $\\\\eta_{2}(\\\\varepsilon)$ in the bound (3).\n\n\n__Miscellaneous points__\n\n- $\\\\mathrm{P}$ that appears in the probability space notation in Defn 4.1 looks identical to equation (P) on page 5. The former isn't really critical, is it?\n- Double spaces (two consecutive spaces between words): e.g., Corollary 5.1 third bullet.\n- Most of the ideas underlying this paper are totally model agnostic, i.e., they hold regardless of what $\\\\mathcal{H}$ is, and yet for some reason the authors start in the *abstract* by talking about updating *\"the network weights\"* without any mention of neural networks before that. This is another example of inconsistency in terms of generality, which I mentioned earlier regarding the posterior parameterization.", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8XnmfwcPIt", "reviewer": "Reviewer_g1cP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- new theoretical results\n- good empirical results\nWeaknesses\n- authors should stress more the novelty of the proposal\n- authors should better clarify the selection of the experimental results", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Authors introduce a new PAC-Bayes training algorithm with improved performance and reduced reliance on prior tuning by proposing a new PAC-Bayes bound for unbounded loss and a theoretically grounded approach that involves jointly training the prior and posterior using the same dataset. Results agains both PAC and ERM based model support the proposal.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Authors should stress more the novelty of the proposal.\nAuthors should better clarify the selection of the experimental results.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors have carefully revised their work following the reviews and discussion with reviewers, and this submission is now worthy of publication at TMLR. I kindly invite the authors to polish the last details and remove the colours used to signal updates from the initial submission. Congratulations on a fine piece of work!
Decision: Accept (poster)