Unlocking Tuning-free Generalization: Minimizing the PAC-Bayes Bound with Trainable Priors

1. Poor compatibility with data augmentation.

From our experiment in the appendix, the method works well with data augmentation, even though we don't have the theoretical explanation. Since data augmentation violates the i.i.d. assumption, having $10$ times augmented data does not mean having $10$ times new independent data. So, in the future, it would be good to explore if we can compute some sort of "effective" dimension of the data augmentation, which characterizes how much new information the data augmentation introduces, or what is the equivalent number of i.i.d. data that brings the same amount of information as the given data augmentation.

2. Lack experiments for the low-data regime.

We strongly agree that the small data and large model regimes are the hardest for PAC-Bayes training because the KL term would be large in this case. Running new experiments with our method is relatively fast, but performing the baseline search is quite time-consuming. We tried fine-tuning the pre-trained large language models to validate our method in the low-data regime. Unlike the experiments we did in the appendix, where the large pre-trained part of the model is fixed, and only the small classification layer is to be trained on the small fine-tuning dataset, we now allow both the pre-trained model and the classification layer to be trained using the small fine-tuning dataset. This is what people do in practice because its performance is much better than freezing the pre-trained model. However, this will make PAC-Bayes training challenging because the number of training samples is around 100 while the total number of training parameters is as large as $110$ million. Table 1 contains the result. It shows that PAC-Bayes training works better than all baselines for fine-tuning.

Table1. Experimental Results for BERT-base-uncased Backbone Model.

[1]. Sennrich, Rico, et al. "Improving neural machine translation models with monolingual data."

[2]. Orvieto, Antonio, et al. "Explicit regularization in overparametrized models via noise injection."

[3]. Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models."

[4]. Liu, Xiao, et al. "P-tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks."

[5]. Zaken, Elad Ben, et al. "Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models."

3. Need more clarity on $K(\lambda)$

$K(\lambda)$ is a function of $\lambda$, and $\lambda$ is the trainable parameters of the prior representing the variance. So $K(\lambda)$ affects what $\lambda$ can be achieved in training, and the prior variance $\lambda$, in turn, affects the posterior variance and the model parameters through their interaction in the bound.

4. Experiments with medium/large datasets.

We have obtained additional numerical results for image classification tasks. Our comparison of PAC-Bayes training with SGD, using a batch size of 128 on the TinyImageNet dataset, involved 100,000 images across 200 classes. Using ResNet18, the test accuracy of PAC-Bayes training is 53.2%, while the best test accuracy achieved with SGD/Adam/AdamW after hyperparameter searching is 46.4%/46.7%/46.8%. This further demonstrates the advantages of our proposed PAC-Bayes training.

5. adapting to varying numbers of training data

We compared PAC-Bayes training with SGD/Adam/AdamW on CIFAR10, using a batch size of 128, and allocated 10% of the training data for training and the remaining 90% for hyper-parameter searching in SGD/Adam/AdamW. With ResNet18, the test accuracy of PAC-Bayes is 67.8%, while the best test accuracies for SGD, Adam, and AdamW, after hyper-parameter searching, are 64.00%, 64.96%, and 65.59%, respectively. When training ResNet18 with all the training data using a batch size of 128, SGD typically achieves the best test accuracy among the baselines. However, AdamW outperforms SGD when using only 10% of the training data. This demonstrates the necessity of both hyper-parameter searching and choosing the appropriate optimizer for baselines. With the same setting and model on CIFAR100, the accuracies are as follows: SGD - 22.16%, Adam - 24.95%, AdamW - 25.68%, and PAC-Bayes - 26.22%.

Thank you for the intriguing question! We feel it makes sense based on the following explanation. In order to fit more training data, the (final) model needs to be able to go further away from the random initial model, which means smaller confidence should be given to the prior (i.e., larger prior variance) to allow this to happen (otherwise if we're very certain about the prior, then the model will just stay near the initial model and won't be able to fit the data). Since the prior and posterior variances are positively correlated, the posterior variance will also be larger. This is a possible explanation of why larger data leads to larger variance.

However, we feel that there is another scenario, that is, when we have a very large amount of data compared to the model (i.e., the under-parametrized regime). Then, adding more data wouldn't pull the model much further from the initial model than it already is. In this case, more data should reduce the uncertainty and, therefore, lead to a smaller variance. We're currently doing more experiments on various amounts of training data of CIFAR10 and will report the result soon. However, we feel since CIFAR10 on ResNet/VGG is in the over-parametrized regime, we probably are not going to see this latter scenario.

Thank you very much for raising the score!

As suggested, we conducted additional experiments and observed that the variances of the learned posterior and prior change due to the different numbers of training samples.

With batch size 128, by training with full/10% of CIFAR10 on VGG13, the final standard deviations of posteriors are 0.032/0.010, and the mean of the standard deviations of priors are 0.042/0.0158. Training with full/10% of CIFAR10 on ResNet18, the mean of the final variances of posteriors are 0.031/0.010, and the final variances of priors are 0.044/0.008.

1. The overall contribution of this work seems minimal, since there is previous literature loosening the bounded loss assumption, and there is little additional information provided in the current work.

Thank you very much for the question. We want to emphasize the gap between the theoretical bounds and applying them to training deep neural networks in practice.

Even though there are many existing PAC-Bayes bounds in the literature, to the best of our knowledge and according to our experiment, none of them works in this difficult setting (training deep CNN networks on Cifar10 with cross-entropy loss). More explicitly, when putting these bounds into practice, they are either non-applicable or too vacuous that the model does not even train (i.e., the training accuracy stays at $10\\%$ and does not increase). Therefore, we did not find a method that gives OK performance to compare with our method.

For example, in the paper https://arxiv.org/pdf/1605.08636.pdf, the reviewer mentioned that Corollary 4 has a PAC-Bayes bound for unbounded loss. However, the $s^2$ is too large in practice. If we train the model with this bound, then the training accuracy stays at $10\\%$ and does not increase.

Another example mentioned in the paper, "Efron-stein pac-bayesian inequalities," has a PAC-Bayes bound for unbounded loss. $E_{h\sim Q}\ell(h;D)\leq E_{h\sim Q}\ell(h;S) + \sqrt{\frac{1}{m} E_{h\sim Q} \left[\ell_1(h;S) + E_{z'\sim D} \ell^2(h;z') \right]\mathrm{KL}(Q||P)} +\frac{1}{m}$.

$\ell_1(h;S):=\frac{1}{m} \sum_{i=1}^m \ell^2(h;z_i)$, and $z'\sim D$ is a test sample drawn from the data distribution. This bound holds for any unbounded loss with a finite second-order moment. However, this bound is semi-empirical, meaning that the bound contains a term, $E_{h\sim Q}E_{z'\sim D} \ell^2(h;z')$, that is hard to estimate from data.

Another example is in the paper [1], where the bound requires the loss to be bounded, and works for two-class classification, rendering it unsuitable for direct use with CIFAR10/CIFAR100.

[1] Dziugaite, Karolina, et al.. "Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data."

To the best of our knowledge, we did not find any other available bounds that are suitable to train deep neural networks directly with unbounded loss and trainable priors.

2. The posterior distribution given by $Q_\sigma(h)$ does not appear to be a valid probability distribution, as defined in the paper.

Does the reviewer refer to the definition here? $Q_\sigma(h):=h+Q_\sigma(0)$. We apologize for the confusion and have changed the notation in the revision (please see the highlighted part in red on page 4). Here, by this notation, we meant $Q_\sigma(h)$ is the distribution of $Y = X+h$, with $X\sim Q_\sigma(0)$, and $Q_\sigma(0)$ is an arbitrary 0 mean distribution parametrized by $\sigma$. So $Q(h)$ is a valid distribution. In the case of Gaussian posteriors in later sections of the paper, this definition leads to $Q_\sigma(h):=N(h, diag(\sigma))$, where $N(h, diag(\sigma))$ refers to the multivariate Gaussian with mean 0 and covariance $diag(\sigma)$

3. Although the terminology surrounding prior and posterior distributions is consistent with Maurer 2004 and parts of the PAC-Bayes literature, the more recent prominence of Bayesian methods in the machine learning literature and their more specific use of these terms leaves room for confusion. Distinguishing between the distributions appearing in the bounds and the specific choice of the Gibbs posterior, for example, would be helpful (see e.g. https://arxiv.org/pdf/1605.08636.pdf, https://arxiv.org/pdf/2110.11216.pdf).

We're not sure if we have understood your question correctly. We believe our notation is consistent with the literature you mentioned https://arxiv.org/pdf/2110.11216.pdf. More explicitly, our Theorem 4.1 is of a similar type to Theorem 2.1 in the referenced paper, which deals with general posteriors. The referenced paper also includes several theorems for Gibbs posterior, but our work does not. This is because, according to our method, the final posterior we used for prediction is computed as the minimizer of the PAC-Bayes bound, which does not admit a closed-form expression like Gibbs posterior.

In addition, because we also optimize during the optimization over the prior, there is no existing bound in the literature that can be used to bound the population error of this posterior. Therefore, we build Theorem 4.2 to characterize the algorithm's performance. We believe this is the first bound of this type which allows data-dependent prior (previous work all require fixed prior). Please let us know if this response addresses your question, or if we have completely missed the point you were making.

1. I don’t see any generalisation bounds reported in the paper, so it is hard to gauge how much of a contribution the novel bound is.

Thank you very much for the question. The short answer is that despite many existing PAC-Bayes bounds in the literature, none of them works in our settings (training deep CNN networks on Cifar10 with cross-entropy loss). More explicitly, when putting these bounds into practice, they are too vacuous that the model does not even train (i.e., the training accuracy stays at $10\%$ and does not increase). Therefore, we did not find a method that gives OK performance to compare with.

We want to emphasize that there is a gap between the theoretical bounds and applying them in training deep neural networks in practice.

For example, in the paper "Efron-stein pac-bayesian inequalities," the bound is provided as:

$E_{h\sim Q}\ell(h;D)\leq E_{h\sim Q}\ell(h;S) + \sqrt{\frac{1}{m} E_{h\sim Q} \left[\ell_1(h;S) + E_{z'\sim D} \ell(h;z')^2 \right]\mathrm{KL}(Q||P)} +\frac{1}{m}$.

$\ell_1(h;S):=\frac{1}{m} \sum_{i=1}^m \ell(h;z_i)^2$, and $z'\sim D$ is a test sample drawn from the data distribution. This bound holds for any unbounded loss with a finite second-order moment. However, the term $E_{h \sim Q} E_{z' \sim D} \ell(h;z')^2$ is almost as difficult to estimate as the generalization error itself.

To the best of our knowledge, we did not find any other available bounds that are suitable to train deep neural networks directly with unbounded loss and trainable priors.

In the paper (Dziugaite and Roy) mentioned by the reviewer, the bound requires that the loss be bounded for two-class classification, which can not be used for CIFAR10/CIFAR100 directly.

2. I would have also liked to see a discussion regarding the trade-offs of having a learnable prior.

Thanks for bringing up this critical point. In the theoretical guarantee (Theorem 4.2), it says our PAC-Bayes training is valid if the correction terms $\eta +\frac{\log(n(\epsilon)}{\hat \gamma m}$ is small. Note that the $\log(n(\epsilon))$ in the correction term represents how rich the set of priors is. So, if the set of priors we're searching over is too rich/large, then this $\log(n(\epsilon)$ would make the correction term large, and the PAC-Bayes training will no longer be reliable. In practice, we observe that for deep network on cifar10, setting the prior to having 1 degree of freedom per layer is fine, but setting the prior to having 1 degree of freedom per weight is too much. This is the motivation of the proposed layerwise prior.

3.The $\gamma_1$ and $\gamma_2$ values are, for instance set to values that work well for the vision tasks, but they actually need adaptation for the graph and text tasks, suggesting that some tuning is actually needed.

We used the same range of $\gamma$ $[0.5,10]$ across different tasks. We can add a sensitivity analysis of the upper and lower bounds of this range to the paper. In short, the result is quite insensitive to the lower bound $\gamma_1$ and is relatively more sensitive to the upper bound $\gamma_2$. So, we ran the CIFAR10 experiment several times to obtain a good value of $\gamma_2$ and fixed it to this value in all other settings.

4. There is also a warmup period only detailed in the Appendix to ensure faster convergence, how was the duration of this set?

The warmup is not necessary and is only used to accelerate convergence. Please refer to the answer to the next question for more explanation of the warmup.

5. My biggest concern however is due to the training time required for the proposed method.

The long wall time of our method is due to the coding issue. Pytorch is not optimized for the PAC-Bayes training. Currently, we use for-loops in Python to inject the posterior noise in each iteration, and that makes the algorithm slow. The conclusion will change if we look at the actual complexity of the algorithms. For both the normal SGD/Adam training and our method, the computation bottleneck in each iteration is the backpropagation, so we may use the total number of backpropagations to represent the complexity. Although our method needs the gradients for both the model and the variance, they can be computed from the same quantity obtained by a common backpropagation. Therefore, our method also only requires 1 backpropagation per iteration. The total number of backpropagations in a single run is thus equal to the total number of iterations, that is:

the number of epoch $* \frac{\textrm{training data size}}{\textrm{batch size}}$

The table below (Table 9 in the appendix) provides this number for CIFAR10/CIFAR100 with ResNet18.

It can be seen that the complexity of the methods is similar. We can expect with code optimization, the wall-time issue can be fixed.

The "trick" like the warm-up iteration mentioned in the supplementary, is used to accelerate the convergence without affecting the final result. Once the code optimization is done, this trick can be removed.

8. Sensitivity analysis of baselines. (Part 2)

We also compared PAC-Bayes training with SGD/Adam/AdamW on CIFAR10, using a batch size of 128, and allocated 10% of the training data for training and the remaining 90% for hyper-parameter searching in SGD/Adam/AdamW. With ResNet18, the test accuracy of PAC-Bayes is 67.8%, while the best test accuracies for SGD, Adam, and AdamW, after hyper-parameter searching, are 64.00%, 64.96%, and 65.59%, respectively. When training ResNet18 with all the training data using a batch size of 128, SGD typically achieves the best test accuracy among the baselines. However, AdamW outperforms SGD when using only 10% of the training data. This demonstrates the necessity of both hyper-parameter searching and choosing the appropriate optimizer for baselines.

7. Performance of default settings of baselines.

Since for SGD, Pytorch has 0 as the default momentum, we guess by default, you mean the setting people commonly use, i.e., momentum $= 0.9$ (please correct us if we're wrong). We are not sure about the common choice of weight decay, so we found the default configuration for ResNet from [1], where it sets the learning rate to 0.1, weight decay to 1e-4, and momentum to 0.9, and uses a learning rate scheduler (reducing by a factor of 10 when the error plateaus). But with this parameter setting, the test accuracy of ResNet18 on CIFAR10 and CIFAR100 is only 85.21% and 60.66%, respectively, which is approximately 4% to 5% lower than that achieved with the best-tuned settings. This suggests when changing network type or dataset, in order to achieve the best performance, hyper-parameter tuning is still needed.

The default learning rate for Adam is 1e-3, and the default weight decay is 0. Using Adam's default settings, the test accuracies for VGG13 and ResNet18 on CIFAR10 and CIFAR100 are displayed in Tables 1 and 2. Although the best test accuracy is similar to that obtained with Adam's default settings at a batch size of 128, the discrepancy is more pronounced with a batch size of 2048. Optimal test accuracy for VGG13 is attainable only with larger batch sizes, where the default settings are ineffective. Our PAC-Bayes training achieves the best test accuracy across various settings.

Table1. The test accuracy when using the default setting of Adam with a batch size of 128. The accuracy inside the $()$ denotes the best test accuracy with a batch size of 128.

Table2. The test accuracy when using the default setting of Adam with batch size as 2048. The accuracy inside the $()$ denotes the best test accuracy with a batch size of 2048.

The default learning rate for AdamW is 1e-3, and the default weight decay is 1e-2. Under AdamW's default settings, the test accuracies for VGG13 and ResNet18 on CIFAR10 and CIFAR100 are presented in Tables 3 and 4. Although these default settings generally yield the best test accuracy for AdamW when using small batch size, they are not as effective as Adam/SGD, making optimizer selection essential. Our PAC-Bayes training achieves better test accuracy than AdamW.

Table3. The test accuracy when using the default setting of AdamW with a batch size of 128. The accuracy inside the $()$ denotes the best test accuracy with a batch size of 128.

Table4. The test accuracy when using the default setting of AdamW with batch size as 2048. The accuracy inside the $()$ denotes the best test accuracy with a batch size of 2048.

[1] He, Kaiming, et al. "Deep residual learning for image recognition."

Thank you very much for thinking our theoretical contribution is nice!

Here, we want to further address the reviewer's questions.

1. Tightness of the generalization bound.

(1) As shown in the first 500 epochs (Stage 1) of Figures 10e-15e in the Appendix, the test loss decreases along with the training loss (our PAC-Bayes bound), and the training loss (our PAC-Bayes bound) is always higher than the test loss. Therefore, our PAC-Bayes bound is tight because: 1. it is a valid bound, and 2. the network can be optimized effectively with this bound.

(2) In Dziugaite and Roy, authors used 1 and -1 to represent classes $\{0,...,4\}$ and $\{5,...,9\}$. So, it is a binary classification problem with bounded loss. We target multi-class classification tasks with unbounded loss. So, our PAC-Bayes bound is essentially different from the one in the paper (Dziugaite and Roy).

2. $\gamma_1$ and $\gamma_2$

We only tried $10$, $20$ for $\gamma_2$ on ResNet18 with CIFAR10. After checking the numerical results, we selected $10$ as the default and fixed it for all experiments and all models (including NLP models such as GPT, GNN models for citation graphs, and CNN models for image classification). So, there is no searching on $\gamma$.

3. Training time

There are currently two challenges for us to achieve the theoretically predicted speed and make the PAC-Bayes training as fast as SGD: software and hardware limitations.

a. Limited access to the source code of PyTorch: Implementing noise injection in PyTorch can be challenging and time-consuming. Adding noise to algorithms like SGD or Adam slows each iteration down to the same speed as PAC-Bayes training, which is around three times slower than a standard SGD iteration. The training time for one epoch needed for noise injection and our PAC-Bayes training are both 27s for ResNet18 on CIFAR10 with one RTX3090 GPU. For pure Adam without noise injection, it only takes 10s for one epoch. This is despite the fact that the complexity of these methods is the same. Noise injection in Julia does not suffer from this issue, but we're less familiar with Julia.

b. GPU memory limitations impact our training speed. While our PAC training benefits from using large batch sizes without performance loss, and larger batch sizes could speed up the algorithm by processing more data simultaneously, our GPU memory limits us. For example, we use a batch size of 2048 for PAC-Bayes training and 128 for SGD. Ideally, training with a batch size of 2048 should be 16 times faster than with 128, but due to GPU memory constraints, it's only twice as fast.

These are soft constraints to the community but are hard constraints to us.

`4. Stability of hyper-parameters

The current tuning seems efficient thanks to existing knowledge, which allows people to follow established hyper-parameters. However, this is not the case if we lack such knowledge about which hyper-parameters are effective. In the next two sets of experiments, we used the same set of hyper-parameters for searching, as mentioned in the paper for CIFAR10/CIFAR100, showing sensitivity when lacking the knowledge for searching.

On CIFAR10 with a batch size of 128, we allocated 10% of the training data for actual training and the remaining 90% for hyper-parameter searching in baselines. With ResNet18, the test accuracy of PAC-Bayes is 67.8%, while the best test accuracies for SGD/Adam/AdamW are 64.00%/64.96%/65.59%. When training ResNet18 with all the training data and a batch size of 128, SGD typically achieves the best test accuracy among the baselines. However, with only 10% of the training data, AdamW outperforms SGD.

We also trained ResNet18 with a batch size of 128 on the TinyImageNet. The test accuracy of PAC-Bayes training is 53.2%, while the best test accuracy achieved with SGD/Adam/AdamW after hyperparameter searching is 46.4%/46.7%/46.8%. This further demonstrates the advantages of our proposed PAC-Bayes training.

While SGD is faster on image tasks, it doesn't consistently deliver the best performance. Even with some knowledge about hyper-parameter searching, it might still be inadequate. Moreover, there are 4 parameters to tune for SGD and 5 for Adam, making it relatively easy to conduct over 12 tuning iterations in total, even if we ignore Adam needs longer training time than SGD.

Tightness of the generalization bound. (numerical comparison)

It is difficult to compare our bound with those for bounded loss (e.g., Dziugaite and Roy) because if the loss is bounded, then our bound has similar tightness as theirs, and if the loss is unbounded, their bound (for bounded loss) cannot be applied. But to partially address your concern, we have numerically compared our bounds with some recent bounds in the literature for unbounded loss, as mentioned by Reviewer YdhE. Those bounds are crude and do not enable PAC-Bayes training because the model does not see an update by minimizing those bounds.

Nevertheless, as you correctly mentioned, we can still compare them as pure bounds instead of training algorithms. Figure (https://tinyurl.com/ycksnjvb) displays the result. More explicitly, we numerically compared our new PAC-Bayes bound with the existing one in the reference (Corollary 4, https://arxiv.org/pdf/1605.08636.pdf) mentioned by the reviewer, which is also for unbounded loss. Even though the forms of the bounds look similar, their actual numerical values are significantly different. As shown in this figure, our proposed bound is far tighter than this baseline one. The sub-Gaussian norm term ($s^2$) in this baseline bound alone is larger than $100$, making it not appealing in practice.

Thank you very much for the constructive suggestions. We will take your comments into consideration in the revision.

1. How does the proposed bound compare to Dziugaite et al. or Haddouche et al. on the same experiment setup?

Thanks for the question, which also highlights our contribution: we bridged the gap between theory and practice while using PAC-Bayes bound. For example, Haddouche et al. (PAC-Bayes Generalisation Bounds for Heavy-Tailed Losses through Supermartingales) have data-free prior to searching the hyper-parameter, which is time-consuming and needs extra validation data to enable the selection. Dziugaite et al. (On the role of data in pac-bayes bounds) need extra data to train the prior, which could also be expensive in practice, e.g., collecting medical data. We have also tried to use existing PAC-Bayes bounds derived for unbounded losses, e.g., [1][2]. However, the cross entropy loss does not satisfy the bound in [1] without putting extra assumptions on the data, while [2] is too difficult to estimate in practice since it requires estimating the expectation of the second-order moment of the loss by drawing a test sample.

In summary, we did not find an existing PAC-Bayes bound that can be directly applied to the classification task without hyper-parameter searching or extra data for the prior. The primary purpose of this paper is to bridge the gap between theory and practice. For our method, we proposed a data-dependent prior that can be trained with the training dataset without requiring extra data to learn the prior or extra tuning to find a good prior.

[1]. Haddouche, Maxime, et al. "PAC-Bayes unleashed: Generalisation bounds with unbounded losses."

[2]. Kuzborskij, Ilja, et al. "Efron-stein pac-bayesian inequalities."

2. What does “index” stand for in this context?

We sorted the test accuracy of different searched hyper-parameters. The x-axis denotes the sorted index of the experiment.

3. The conclusion section makes the claim that the proposed method also mitigates the curse of dimensionality. How do we conclude this from the rest of the paper?

The curse of dimensionality refers to the fact that when the model gets large, the KL term (and therefore the PAC-Bayes bound) could potentially get very large. Experiments show our method can achieve state-of-the-art results, indicating this curse of dimensionality is mitigated. We also tried other bounds, which cannot handle models as large as resnet18 (on cifar10), as the bound gets too large to be practical.

1. The paper suggests that only weight decay and noise injections are essential for PAC-Bayes training. However, this conclusion seems premature and lacks comprehensive analysis.

We agree that studying the effect of the interplay between different kinds of regularization is still an important open problem, and we have mentioned it in the introduction. Here, we state that only noise injection and weight decay are essential from our derived PAC-Bayes bound.

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. [1] 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 [2] 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 [3].

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 to the regularization effect of batch-norm and dropout. During training, suppose the current posterior is $\mathcal{Q}_{\hat{\sigma}}(\hat h) = \mathcal{N}(\hat h,\textrm{diag}(\hat \sigma))$,

Let $h \sim \mathcal{Q}_{\hat{\sigma}}(\hat h)$,

and $\Delta{h} \sim \mathcal{Q}_{\hat{\sigma}}(0)$.

The training loss expectation over the posterior is:

$\mathbb{E}_{h} \ell(h;\mathcal{D})$

$= \mathbb{E}_{ \Delta{h} } \ell(\hat{h}+\Delta {h};\mathcal{D})$

$\approx \ell(\hat h;\mathcal{D})+\frac{1}{2} \mathrm{Tr} (\textrm{diag}(\hat\sigma) \nabla^2 \ell(\hat{h};\mathcal{D}))$

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 a separate regularization effect from sharpness. So, 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.

1. Wei, Colin, et al. "The implicit and explicit regularization effects of dropout."
2. Luo, Ping, et al. "Towards Understanding Regularization in Batch Normalization."
3. Cohen, Jeremy, et al. "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability."

2. The method proposed in this paper may require i.i.d. data.

Yes, that is correct. The PAC-Bayes analysis to study generalization is limited to in-distribution tasks. We do not think of it as a drawback but rather out of the scope of this paper. Analyzing out-of-distribution generalization is still a developing field and will need additional statistical tools to derive generalization bounds for OOD tasks. However, there are papers discussing using the PAC-Bayes bound to detect the out-of-distribution samples [4], which could be the future work.

3. Why is it so important to give a generalization bound on unbounded loss? Existing bounded one can not work well in practice? Any empirical comparison with them when applying those bounds for training?

Since the prevalent cross-entropy loss (CE) is unbounded, it is important to derive a PAC-bayes bound to deal with it. Existing ones for bounded loss (such as Theorem 2.1) do not work well on the CE loss because the $C$ in Theorem 2.1 has to be set to infinity for CE loss, which makes the bound meaningless. We also tried to approximate the CE loss by a bounded loss first and then used the PAC-Bayes bound for bounded loss, which again failed to work when we explicitly tried to truncate/clip the cross-entropy loss to make the loss bounded. In the experiment, we observe that this will cause the training and test accuracy to plateau at very low levels around $10\\%$ to $20\\%$.

4. The bounded part has now been shifted to the bounded $\gamma$ .

The bounded part has been shifted to the $K(\lambda)$ in Theorem 4.1. More explicitly, the $K(\lambda)$ in Theorem 4.1 (our theorem) corresponds to the $C$ (the upper bound of the loss) in Theorem 2.1 (the one for bounded loss). So, we can understand $K(\lambda)$ as an effective bound of the unbounded loss. Limiting the range of $\gamma$ to $[\gamma_1,$ $\gamma_2]$ reduces the value of $K(\lambda)$, which makes the bound tighter. For instance, if we don't limit this range, the $K(\lambda)$ we found for cifar10 on resnet18 would be $4$ versus $0.3$ when we limit the range.