Paper: Unlocking Tuning-free Generalization: Minimizing the PAC-Bayes Bound with Trainable Priors (ICLR-2024)

Reviewer comment (Theory):
**Poor compatibility with data augmentation**: a major downside of the method is its poor compatibility with data augmentation techniques, which are important and beneficial to boost a model's performance. The reason is that the algorithm needs the exact number of training data sizes and data augmentation breaks this assumption. This is discussed by the authors in the appendix. Nevertheless, this is an important limitation of the approach.

Author response:
`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.

Reviewer comment (Experiments):
**Lack experiments**: On CIFAR10 and CIFAR100 experiments, the proposed method is on par with existing optimization methods, which need extensive hyper-parameter tuning, especially for regularization heuristics. Given the theory-oriented nature of this paper, this is reasonable. However, it would be great to scale up experiments to medium/large datasets. More importantly, given the claim that the proposed framework automatically learns regularization, it is *necessary* to test its performance under a low-data regime for example using only 10% of CIFAR10 and CIFAR100. This will make the claim of being *nearly tuning-free* much more convincing.

Author response:
`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. 
||CoLA |SST |MNLI-m| QNLI| RTE|
|---|---|---|---|---|--|
|metrics|MCC|Accuracy|Accuracy|Accuracy|Accuracy| 
| | | | | |
|Vanilla-tuning| .235| .773 |.369 |.702 |.589 |
|Data Augmentation[1]|.171 |.817 |**.395**| .705 |.594| 
|Noise Injection[2]|.233 |.783 |.371 |.706 |.588| 
|LoRA[3]| .298 |.792| .385 |.669 |.592 |
|Prefix-training[4]| .191 |.704 |.375 |.649 |.565| 
|BitFit[5]| .267 |.768 |.376| .647 |.588| 
|Ours| **.335** |**.834** |.387 |**.709** |**.601**|

[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."

Reviewer comment (Theory):
**Need more clarity on $K(\lambda)$**: The loss function derived from the PAC-Bayes bound has three important terms, the empirical data-driven loss, the KL divergence between the posterior and the prior, and $K(\lambda)$. It's not clear what role $K(\lambda)$ plays in the overall algorithm. Since it is estimated before training, it is also not clear how it affects the learning in terms of gradient.

Author response:
`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.

Reviewer comment (Experiments):
My major concern is over the experiments. While it is reasonable to have small-scale experiments, I would really like to see how the proposed method works for low-label situations since a major claim is its tunning-free capability. This means that it can adjust its built-in regularization strength to different dataset sizes.

Author response:
`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\%.

Reviewer comment (Experiments):
It would make the claim much more convincing if the algorithm showed promise in adapting to varying numbers of training data. I would suggest the authors run a study on varying the number of training data, e.g., 10%, 50%, and report the performance and the learned variances for the prior and posterior. What do you expect the prior and posterior variances to be when the dataset size is small and how do the results support this?

Author response:
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.

Reviewer comment (Experiments):
The authors develop a novel and arguably tighter PAC-Bayes bound but never put it to test (at least from what I could see). 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 in terms of tightness, compared to the previous works cited in this paper. I understand that this is not the main focus of the paper but it would definitely strengthen its technical contribution, i.e. the bound itself. I would have also liked to see a discussion regarding the trade-offs of having a learnable prior. What if the set of priors is so large that it includes the posterior, i.e. the KL term could be set to zero perfectly. How large would the resulting penalisation term be? Is there ever a scenario where this could lead to a non-vacuous bound? Also, how does this learnable prior compare to methods that perform a (discrete) grid search over the prior and then perform a union bound, resulting in an additional penalty term (e.g. [1])?

Author response:
`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.

Reviewer comment (Experiments):
There seems to be a certain arbitrariness as to what hyper-parameters are simply chosen ad-hoc and turn out to work well. 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. I also could not find a stability analysis for the default choices, were those just the first values tried or was an initial grid search actually needed to find those values? There is also a warmup period only detailed in the Appendix to ensure faster convergence, how was the duration of this set? I also would like to see a comparison to SGD/Adam/AdamW, where the default hyper-parameter settings are used. For instance, I very rarely see anyone changing the default momentum value for simple tasks such as CIFAR10/CIFAR100. I would also like to see a plot depicting how sensitive the baselines are to individual parameters such as the learning rate, momentum etc, I think the results are already in Figure 1 but need to be re-grouped accordingly.

Author response:
`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.

Reviewer comment (Experiments):
My biggest concern however is due to the training time required for the proposed method. A single run on CIFAR10, according to Table 8 in the Appendix takes roughly 7000 seconds, while for SGD the same takes roughly 600 seconds. This means that one can roughly perform 12 SGD runs in the time it takes to perform a single run with the PAC-Bayes technique. I would argue first that (1) it does not take 12 runs of SGD to find an acceptable hyper-parameter setting, especially not on simple tasks such as CIFAR10/CIFAR100. (2) Even if 12 runs are required, I would argue that one has obtained a better understanding of the task in the sense of what techniques work etc, and subsequent re-runs of the same method will be very cheap. Moreover, one could even make use of those 12 runs and build an ensemble, perform uncertainty quantification etc. I also believe that this should be discussed more transparently in the main text, instead of the Appendix.

Author response:
`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.

||CIFAR10	| CIFAR100|
|---|---|--|
|SGD	|11501	|17095|
|Adam	|28508	|28335|
|AdamW	|24761	|9560|
|scalar	|15912	|15216|
|layer	|15960	|15168|

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.

Reviewer comment (Experiments):
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. It would be beneficial to investigate and compare the impact of other regularization techniques commonly used in deep learning, such as dropout or batch normalization, within the proposed framework. This would provide a more comprehensive understanding of the interplay between different regularization methods and their contribution to generalization performance.

Author response:
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."

Reviewer comment (Theory):
The method proposed in this paper may require i.i.d. data and may not be able to deal with out-of-distribution tasks.

Author response:
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.

Reviewer comment (Theory):
Could the authors elaborate more on why it is so important to give a generalization bound on unbounded loss? Existing bounded one can not work well in practice? An empirical compare with them when apply those bounds for training? From my understanding, the bounded part has now been shifted to the bounded $\gamma$ with $\gamma_1$ and $\gamma_2$ in the unbounded bounds.

Author response:
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.

Reviewer comment (Theory):
Could the authors elaborate more on why optimizing prior on training dataset will be helpful for the generalization performance. Normally, we should fix prior or choose a good one using validation dataset.

Author response:
Thank you for highlighting this important central question on PAC-Bayes bound for neural networks. In applications of PAC-Bayes bounds to generalization error, the contribution of the KL divergence often dominates the bound: To have a small KL, the posterior must be close to the prior. This means it is critical to choose a good prior. However, this is difficult without knowledge of (or access to) the data distribution. Indeed, using data-independent priors (choosing the prior that is independent of data) often leads to vacuous bounds due to ignoring important, favorable properties of the data distribution. Alternatively, one can search for good priors based on validation data, but as the searching grid gets denser and denser, this is almost equivalent to optimizing over the priors, except that the latter is more efficient.

Moreover, some tasks lack training data samples, e.g., node classification based on graph neural networks only has around 100 training data points, and splitting more nodes to select the prior will compromise the test performance in this scenario. At last, hyper-parameter tuning could be time-consuming, which we intend to avoid in this paper.