Non-Vacuous Generalization Bounds for Large Language Models

Model scale and few shot capabilities:

In this work, we focus on constructing the first non-vacuous bounds for LLMs. We have added additional results for larger variants of GPT-2 of sizes of size 354M (GPT-2 medium), 458M, 773M (GPT-2 large), and 849M parameters; we report the results below. The biggest model we trained previously was nanoGPT, also known as GPT small, of size 124M parameters trained on the OpenWebText dataset that contains about 9B tokens (9,035,582,198). However, these are still far smaller than the most capable LLMs. Training LLMs of 10B or larger requires significant computational resources, not to mention the tremendous engineering and infrastructure effort required. While we do believe that constructing bounds for such large models is valuable, our work is merely the first effort at constructing non-vacuous bounds for LLMs and it may be prudent to await several advancements in bound construction that are likely soon to come before expanding such an investment. We believe that our paper can lay the groundwork for such later improvements. Most of the zero, few-shot, and instruction-following capabilities of LLMs don’t emerge until model sizes significantly larger than GPT-2. Thus while the questions of how generalization bounds relate to these capabilities are definitely of interest, they cannot straightforwardly be answered for models this size, and we are optimistic about addressing these questions in future work.

Bounds for larger models:

Note that for the new experiments involving larger models, our approach yields non-vacuous bounds, even for models with nearly a billion parameters. Moreover, we see that the smallest model, where we previously performed experiments and tuned our hyperparmeters, actually achieves the worst bound on bits per dimension.

It is important to note that due to time and computational constraints, we pretrain these models with SubLoRA only for a single hyperparameter setting in contrast to the 124M model for which we did a thorough hyperparameter sweep; we also did fewer optimization steps for larger models given the time constraints. Therefore, it is likely that the tightest empirically achievable bounds are much stronger for the new large models than what we report in the table below. The main conclusions that we draw from these results are: (1) our approach extends naturally to much larger language models; (2) it is possible to achieve tighter bounds as we increase the size of the model.

Subspace only is enough to achieve non-vacuous bounds in table 1:

That is correct. When combined with our ingredients (3,4,5), subspace-only bounds (aka 1,2) are sufficient to achieve non-vacuous bounds in many cases. However, adding SubLoRA (6) substantially improves these bounds.

Tradeoff between SubLoRA and standard pre-training:

SubLoRA converges much faster in practice than standard pretraining. While there are no memory gains, we effectively have a more compressed delta between the randomly initialized weights and the final model, leading to tighter bounds.

Thank you again for your review and suggestions. We believe your feedback has improved our paper. We made a significant effort to run additional experiments and address your questions and would appreciate it if you would consider raising your score in light of our response. Please let us know if you have any additional questions we can address.

References:

[1] Lotfi, S., Finzi, M., Kapoor, S., Potapczynski, A., Goldblum, M. and Wilson, A.G., 2022. PAC-bayes compression bounds so tight that they can explain generalization. Advances in Neural Information Processing Systems, 35, pp.31459-31473.

[2] Zhou, W., Veitch, V., Austern, M., Adams, R.P. and Orbanz, P., 2018. Non-vacuous generalization bounds at the imagenet scale: a PAC-bayesian compression approach. arXiv preprint arXiv:1804.05862.

[3] Aghajanyan, A., Zettlemoyer, L. and Gupta, S., 2020. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. arXiv preprint arXiv:2012.13255.

[4] Arora, S., Ge, R., Neyshabur, B. and Zhang, Y., 2018, July. Stronger generalization bounds for deep nets via a compression approach. In International Conference on Machine Learning (pp. 254-263). PMLR.

[5] Haddouche, M., Guedj, B., Rivasplata, O. and Shawe-Taylor, J., 2021. PAC-Bayes unleashed: Generalisation bounds with unbounded losses. Entropy, 23(10), p.1330.

We thank you for your detailed and thoughtful feedback. First, we address the comparison with prior work and clarify which elements of our bounds are novel (including how they relate to other methods) and which are not.

We would first like to point out that ingredients (1) and (2), i.e., using a non-uniform hypothesis prior and the intrinsic dimensionality to compress neural networks, are not novel to our paper, and we build on these techniques as most closely done in Lotfi et al. (2022) [1], but these techniques have been used to differing extents in other papers such as in [2, 3, 4]. We note that Aghajanyan et. al (2020) [3] does not actually prove any new bounds, or evaluate any existing bounds on LLMs, but instead highlights how intrinsic dimensionality could be used for such a purpose. We have made an effort to clarify the relation with these works in the updated draft.

In terms of the novel ingredients of our bounds, we would like to emphasize (3,4) that you identify as well as (5) the overall approach to handling the non-IID token level structure, and (6) the use of LoRA in the compressed parametrization. (3) is a novel approach, and other existing approaches for handling unbounded objectives are not well suited to LLMs. For example, in [5] the authors introduce a hypothesis dependent range $\sup_x R(h,x) \le K(h)$ for some function $K(h)$, then $K(h)$ can essentially be used in place for the bounded range. However for neural networks, $K(h)$ grows exponentially in the number of layers of the network and relates to the Lipschitz constant for a fixed network, which is far too large to be making bounds on large networks like LLMs. We have added discussion of this and other theoretical techniques for handling unbounded objectives, explaining why they are not well suited to LLMs.

You state that ingredient (4) exists in some form in Lotfi et al. (2022), however this is not the case. Perhaps some confusion arises in the case of the data dependent bounds which reserve a fraction of the training data for computing a data dependent prior, however this is of a very different nature than the subsampling bounds which we employ here. For the full training dataset of size $m$, our subsampling bounds allow us to compute the empirical risk on a subsample of size $n<<m$ while maintaining the dominating complexity term $\frac{\log 1/P}{2m}$ using $m$ from the full dataset for a data independent bound. In contrast, for the data dependent priors a small fraction $(m-n)$ is reserved for training the prior, but the bound complexity scales as $\frac{\log 1/P}{2n}$ because only $n$ are considered as the training data points for adapting the prior to posterior. For this reason in data dependent bounds, $n$ should be chosen not much smaller than $m$ so as to preserve the tightness of the bound, and these subsampling bounds serve an entirely different purpose to our subsampling bounds (which are data independent). The subsampling bounds that we introduce reduce the bound evaluation time dramatically from 3 days on a 8 GPUs in parallel to 45 minutes on a single GPU.

Below, we address your other questions and comments.

Thank you for your supportive feedback; we appreciate it! Inspired by all the reviews, we have improved the clarity of our paper and run additional experiments to extend our generalization bounds from the nanoGPT model of size 124M parameters (GPT-2 small) to much larger GPT-2 variants. We also obtained classification error bounds for fine-tuning tasks on the GLUE datasets [2], highlighting the generalization implications of pretraining LLMs.

Generalization bounds for much larger models: We use SubLoRA to obtain generalization bounds for much larger variants of GPT-2 of sizes 354M (GPT-2 medium), 458M, 773M (GPT-2 large), and 849M parameters. Note that our approach yields non-vacuous bounds, even for models with nearly a billion parameters. Moreover, we see that the smallest model, where we previously performed experiments and tuned our hyperparmeters, actually achieves the worst bound on bits per dimension.

It is important to note that due to time and computational constraints, we pretrain these models with SubLoRA only for a single hyperparameter setting in contrast to the 124M model for which we did a thorough hyperparameter sweep; we also did fewer optimization steps for larger models given the time constraints. Therefore, it is likely that the tightest empirically achievable bounds are much stronger for the new large models than what we report in the table below. The main conclusions that we draw from these results are: i. our approach extends naturally to much larger language models; ii. it is possible to achieve tighter bounds as we increase the size of the model.

Generalization bounds for downstream tasks:

For fine-tuning experiments, we consider two binary classification tasks on the QQP and Cola datasets from GLUE [2]. We contrast the bounds that we obtain for fine-tuning a pretrained GPT-2 small model [3], to training from scratch the same model on the two datasets. For both the pretrained and randomly initialized GPT-2 large models, we fine-tune using SubLoRA with rank = 8 and intrinsic_dim=30000 and perform fine-tuning. We run the fine-tuning for 5 epochs with a learning rate of 2e-5. We obtain the following non-vacuous classification accuracy bounds for both QQP and Cola:

As we can see, using pretrained LLMs leads to tighter, non-vacuous bounds in contrast with training from scratch. Our bounds thereby provide a theoretical certification on the value of pretraining LLMs.

Thank you for your supportive, constructive, and detailed feedback. We address your comments and questions below.

The value of non-vacuous bounds for LLMs:

As we mentioned in the general comment, obtaining non-vacuous bounds for LLMs is of paramount significance given the debate about their ability to generalize to unseen data. While empirical non-vacuous bounds in the literature and in our work are looser than the validation error, their existence tells us that the model’s performance is not trivial in expectation for any data coming from the same distribution as the training data. This theoretical guarantee settles the debate about LLMs’ ability to generalize. Another fundamental question for the ML community is whether bigger LLMs are more likely to merely memorize their training samples and not perform any meaningful generalization. We show using our non-vacuous bounds that larger models are in fact more compressible and lead to better bounds. Finally, our work is merely the first step towards a deeper investigation of the generalization capabilities of LLMs through the lens of generalization bounds, and we expect future work to uncover more insights, in the same fashion Lotfi et al. (2022) did for vision.

Fundamentally novel contributions:

We want to clarify that our work does provide contributions that are distinctly tailored to LLMs. While we were able to bound the top-1 error for next token prediction using existing PAC-Bayes bounds, we strived to design generalization bounds that extend to the BPD continuous and unbounded loss since it is the metric of interest in LLMs. Our subsampling bounds are also specifically designed to accommodate for the massive datasets that LLMs are trained on. In fact, we reduce the bound evaluation time from 3 days using 8 GPUs in parallel to 45 minutes using a single GPU. While SubLoRA is a combination of existing methods, it comes from our novel observation that pretraining using LoRA on the last fully connected layer in addition to the attention layers – used exclusively for LoRA fine-tuning – leads to a non-trivial performance. Another major challenge in applying existing bounds to LLMs lies in the autoregressive nature of LLMs, and how the tokens themselves are not i.i.d. To address this challenge, we first observed that unlike some models where sequences draw on the previous history that lies outside the context window, such as with Transformer-XL or Mistral-7B, the GPT2 model takes in these sequences without considering any previous history. Then, we made the choice to divide the training data into non-overlapping sequences of size equal to the context length and randomly select from those sequences. The bounds that we obtained for LLMs in this work would not have been possible to achieve without addressing each of these challenges.

Significance of the results with the I.I.D assumption:

We believe that our bounds are still of practical significance and are informative about generalization despite considering entire sequences to satisfy the I.I.D assumption.In fact, the bits-per-dimension for a given sequence can be computed as the average error for each token in the sequence given previous tokens, where the token error here refers to the negative log-likelihood $\mathrm{BPD}(h, X):= -\log_2 p_h(X)/L = - \sum_i^L \log_2 p_h(x_i|x_{<i}) /L$. Therefore, an upper bound on the expected BPD error still reflects a guarantee on the average performance of the model at the token level, conditioned on previous tokens within independent sequences, and is a common quantity of interest in language modeling.

Thank you again for your review and suggestions. We hope that we were able to address all of your questions. We also ran new experiments that we detail in the general comment, including non-vacuous generalization bounds for very large models with up to 849 million parameters and non-vacuous generalization bounds for LLMs fine-tuned on downstream benchmark tasks. Please let us know if you have any additional questions we can address.

Thank you for your feedback. According to your feedback and feedback from the other reviewers, we have made updates to the PDF to help clarify the construction of the SubLoRA parametrization, the IID assumption requirements in both the full and subsampling bounds, and small typos such as the one you point out for the matrices $Q_1$ and $Q_2$.

SubLoRA implementation:

For equation (2) and SubLoRA, the expression is not a typo. In the forward pass of the model, first the matrix P is applied to the reduced dimension weights w to project to a space with dimension equal to the number of LoRA parameters. These LoRA parameters are in the form of biases and matrices to be multiplied together to form the low edits to the random initialization. If $v=Pw$, $\mathrm{LoRA}(v)$ reshapes $v$ into a list of these biases and low rank factors, multiplies the factors and then reshapes the result back into a vector (with ordering so that when used by the network the biases and weight matrices will be in the right places).

Subsampling bounds:

In Section 4.4 with the subsampling bounds, the essential point is that we only require that $\hat{R}\_{\sigma(i)}(h)$ for different $i$ are independent when conditioned on $X$. In this setting, the only randomness in $\hat{R}\_{\sigma(i)}(h)$ is over $\sigma(i)$. With this in mind, we can rewrite the quantity $\hat{\hat{R}}$ as merely the sum of $n$ randomly chosen elements from the set of deterministic elements $\{ c_i \}\_{i=1}^m$ where $c_i:=\hat{R}\_{\sigma(i)}(h)$. This random choice can be with replacement (to strictly be independent) or without replacement (Hoeffding inequalities also apply to this setting).

Technical contributions of our paper:

As we mentioned in the general responses, we make several technical contributions to address challenges that are specific to large language models. First, the loss function of interest in LLMs is the number of bits-per-dimension, which is neither bounded nor discrete. Therefore, the PAC-Bayes bounds used by Lotfi et al. (2022) are no longer valid in this setting. In contrast with other generic approaches for handling unbounded objectives, our bounds can be practically applied to the BPD objective of LLMs. Moreover, the bound evaluation for the massive datasets that LLMs are trained on is very computationally expensive. For instance, it takes about 3 days to evaluate our bounds on the OpenWebText dataset using 8 GPUs in parallel. The subsampling bounds that we propose in this work bring down that evaluation time to 45 minutes on a single GPU, while paying a very small penalty on the tightness of the bound. Besides these theoretical contributions, we also provide a practical method to achieve non-linear subspace compression for LLMs, while noting from Figure 1 and Table 1 that LoRA alone is not sufficient to achieve non-vacuous top-1 error bounds.

As you highlighted, engineering a working non-vacuous bound can be challenging. In particular, we had to choose our instances to correspond to entire non-overlapping sequences in order to preserve the I.I.D assumption and to select an optimal subset of layers to apply SubLoRA to, which involved both the attention layers and the final fully connected layer. We believe that our work is highly significant to the machine learning community, as it provides the first theoretical certification that large language models are capable of generalization beyond their training data, and that as their size increases, they become more compressible and lead to better bounds, in line with observed performance benefits of bigger models.

Thank you again for your review and suggestions. We made a significant effort to address your questions and also have run new experiments that we detail in the general comment, including non-vacuous generalization bounds for very large models with up to 849 million parameters and non-vacuous generalization bounds for LLMs fine-tuned on downstream benchmark tasks. We would appreciate it if you would consider raising your score in light of our response. Please let us know if you have any additional questions we can address.

We thank you for your detailed and supportive review. We respond to your questions below.

Task and Top1 error:

Indeed, we consider the token level error averaged over the sequence chunk as the empirical risk which we bound. In the case of top1 error, this is the top1 error per token averaged over the chunk $R(h,X_k) = \frac{1}{L} \sum_{i=1}^L 1[{\mathrm{argmax} \ p(x_{i}|x_{<i})=x_{i}^k}]$, where the upper index $k$ denotes the chunk index and the lower index denotes the position within the chunk.

I.I.D assumption and choice of $L$:

We satisfy the I.I.D assumption by dividing the data into non-overlapping chunks/sequences of size L, and then randomly selecting from those sequences. Unlike some models where the input sequences draw on the previous history that lies outside the context window, such as with Transformer-XL or Mistral-7B, the GPT2 model takes in these sequences without considering any previous history. In this setup, it makes sense to set the size of the sequence L to be equal to the context length (which is 1024 by default for nanoGPT/gpt-2 small). This value of L ensures both the computation of the complete joint distribution $p_h(X) := \Pi_i^L p_h(x_i|x_{<i})$ for each sequence and the preservation of the I.I.D assumption.

Clarification about the code:

We apologize for the confusion about the code. The Glue code was used to run some fine-tuning experiments that we did not report in the paper initially, but that we report in the rebuttal as shown in the general comment. The nanoGPT folder contains all the code required to reproduce our experiments. To pretrain using SubLoRA, you can run the file train.py with the configuration file config/train_gpt2_LoRA_better_hparams.py using the following command, where the provided hyperparameters lead to the best SubLoRA bound: python train.py config/train_gpt2_LoRA_better_hparams.py --intrinsic_dim=25000 --learning_rate=5e-3 --attention_linear_lora_r=4 --linear_head_lora_r=4.

To compute the bounds for the trained checkpoint, you can run the file bound_pipeline.py with the configuration file train_gpt2_LoRA_bounds.py using the following command, where again, the provided hyperparameters lead to the best SubLoRA bound: python bound_pipeline.py config/train_gpt2_LoRA_bounds.py --intrinsic_dim=25000 --learning_rate=5e-3 --levels=17 --best_checkpoint_path=$PATH_TO_CHECKPOINT

We provided a ReadMe file in the zipped code to facilitate navigating it. Your point about the Hugging Face Repository is correct, since we copied it without removing the copyright but removed any other information that can reveal our identity.

Validation loss for the different methods:

When comparing LoRA, linear subspace training, and SubLoRA, we use the training error because it contributes directly to the bound, as it constitutes the first term in the bound. To give you an idea about how the test error compares between the 3 methods, we report below the BPD test values for the LoRA, linear subspace, and SubLoRA for the models that yield the best bounds. Note however that we are not optimizing for the validation loss, but rather for the error bound.

Significance of the results with the I.I.D assumption:

We satisfy the I.I.D assumption by considering instances to be non-overlapping sequences of size equal to the context length instead of single tokens. This procedure is precisely how training data is sampled when training language models in practice. Moreover, the bits-per-dimension for a given sequence can be computed as the average error for each token in the sequence given previous tokens, where the token error here refers to the negative log-likelihood $\mathrm{BPD}(h, X):= -\log_2 p_h(X)/L = - \sum_i^L \log_2 p_h(x_i|x_{<i})/L$. Therefore, an upper bound on the expected BPD error still reflects a guarantee on the average performance of the model at the token level, conditioned on previous tokens within independent sequences, and is a common quantity of interest in language modeling.

Thank you for pointing out the minor typos and unclear sentences in the text, we have updated the paper to reflect these corrections and also add clarifications about the task, the top-k errors, and the I.I.D assumption. We also run new experiments that we detail in the general comment, including non-vacuous generalization bounds for very large models with up to 849 million parameters and non-vacuous generalization bounds for LLMs fine-tuned on downstream benchmark tasks.

Thank you again for your feedback, we believe it has positively impacted our paper. We hope that we have been able to address your points, We would also be happy to further engage to answer other questions you may have.

Hi, Any updates about the supplementary material?

Significance: Aside from the novelty of our approach, we believe that constructing nontrivial generalization bounds for LLMs is also significant. We emphasize the purpose of generalization bounds in deep learning more broadly, as well as why constructing them for LLMs is particularly significant and timely.

1. Broader significance of generalization bounds

The generalization of neural networks is not well understood theoretically. Large neural networks have been shown empirically time after time to generalize well on common problems despite their high complexity, but the familiarity of this behavior does not constitute an explanation. Different possible explanations have emerged, but in the absence of a strong theoretical argument, it is difficult to say which, if any, of these explanations are correct. One way to determine which explanation holds explanatory power is to craft a generalization bound for each and see which explanation generates bounds which can actually explain generalization. While bounds give a certificate on future performance, this certificate is not our aim and would be better achieved by creating a statistical bound on held out data in our opinion. Rather, the purpose of generalization bounds is to determine why models work and to provide concrete evidence for compression as an explanation for the behavior of LLMs.

2. The significance of non-vacuous bounds for LLMs

The ability of large language models to generalize beyond their training data is still an active topic of debate within the machine learning community, referring to them sometimes as “stochastic parrots”. By obtaining non-vacuous bounds for LLMs, our work provides a mathematical proof that they are, in fact, powerful knowledge compressors and are capable of generalization beyond their training samples.

3. Generalization and compression improve with model size

Lotfi et al. (2022) note that their compression bounds prefer models with a smaller number of parameters, instead of larger models which tend to generalize better in practice. They hypothesize that nonlinear parameter compression schemes might lead to a different trend where larger models are more compressible than smaller models. In this work, we demonstrate using our nonlinear compression method that larger models are actually more compressible and have better bounds. In other words, as we increase the size of the model, its ability to effectively generalize beyond the training data instead of merely copying it improves as well. To the best of our knowledge, our work is the first to show that generalization bounds improve with more parameters on models of practical sizes, in line with empirical benefits of large models.

4. The importance of pretraining LLMs

Our additional transfer learning experiments show that we obtain tighter generalization bounds when we used a pretrained LLM instead of training from scratch, providing a theoretical quantification of the benefits of pretraining in LLMs.