Paper: Non-Vacuous Generalization Bounds for Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
The paper seems to lack a thorough comparison to previous theoretical work that highlights the novel theoretical contributions. For example, the cited **Aghajanyan et. al (2020)** proves a generalization bound for classifiers that have been obtained by fine-tuning a pre-trained language model. Their bound already relies on intrinsic dimension (Ingredient 2) to reduce the hypothesis space and what seems to me a version of Ingredient 1 for compression based on **Arora et al**: *Stronger Generalization Bounds for Deep Nets via a Compression Approach*. Ingredients (1, 2) and a version of Ingredient (4) seem to be present in **Lofti et. al**.

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

Reviewer comment (Theory):
While the authors emphasize the importance of non-vacuous bounds during the pretraining phase, a more detailed justification of its significance *within the LLM context* would enhance the paper's impact. At least I am not sure that for LLM, generalization in the pretraining phase is such important.

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

Reviewer comment (Novelty):
Given the point 1, I wish to see some novel techniques. However, it seems that the techniques in this paper are not very novel. The methods presented in this paper largely combine existing techniques (e.g., SubLoRA).

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

Reviewer comment (Theory):
The paper claims that tokens exhibit non-iid behavior but resolves this issue by considering entire sequences, which might be considered a somewhat coarse approach.

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

Reviewer comment (Theory):
The technical contribution is moderate on the conceptual part is fair, however, engineering a working non-vacuous bound can be challenging.

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

Reviewer comment (Theory):
In Equation (2) it seems that LoRA is applied to $Pw$, it is not clear to me how that can be implemented, in particular, how the weights can have the forms $Pw$ and $UV$ simultaneously. It might, however, be just a typo and the equation should be $P \cdot LoRA(w)$.

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

Reviewer comment (Theory):
In section 4.4, it is not clear why we can assume that $\hat{R}_{\sigma_{i}}(h)$ are independent. In general, taking a random sample of a sequence does not make such a random sub-sample independent.

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

Reviewer comment (Writing):
A minor typo: $Q_1, Q_2 \sim \mathcal{N}(0,1)^{\sqrt{D}\times d}$ ----> $Q_1, Q_2 \sim \mathcal{N}(0,1)^{\sqrt{D}\times \sqrt{d}}$

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

Reviewer comment (Evaluation):
Also perhaps more clarity about errors. When it is said top-1 error, do you mean the worst sequence, or the worst token?

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

Reviewer comment (Theory):
i.i.d. assumption is a very strong assumption.

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

Reviewer comment (Experiments):
A workaround is found which considers sequences not tokens, but then it's unclear how the parameter L is chosen.

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

Reviewer comment (Reproducibility):
It seems the provided code is running the glue suite, and the code involves LoRA, but not subLoRA. And the other folder seems to be from nanoGPT.

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

Reviewer comment (Evaluation):
Glue suite seems to be unrelated to the task in paper: token prediction.

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

Reviewer comment (Reproducibility):
Can you please check if the provided code is what you intended to submit? And let me know if I am missing something? And also some rough system specifications on which it can be run?

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

Reviewer comment (Evaluation):
When comparing LoRA and Subspace (e.g. in figure 1) we see a comparison of train error, but not test error. Is there a reason for that?

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

| Metric | SubLoRA | LoRA Only  | Subspace Only  | 
| ---       | ---          | ---        |  ---         | 
|  Training Loss (NLL)    |   7.20    |   6.42    |  8.85   |    
|  Validation Loss (NLL)    |   7.19     |    6.48      |  8.85    |

Reviewer comment (Theory):
Can you please comment on why, despite the i.i.d. assumption, these results are still significant?

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

Reviewer comment (Writing):
Page 6: 'payed'.

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

Reviewer comment (Writing):
Abbreviation NLL used without definition.

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

Reviewer comment (Writing):
Page 4, second last line, shouldn't it be "u: = flatten( ...)" instead of "LoRA(u) := flatten ( ... )".

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

Reviewer comment (Writing):
Page 6, last paragraph says we use "several" values for r, giving the impression they are more than two, but the appendix mentions two values {1,4}.

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

Reviewer comment (Writing):
Page 14, second last paragraph: "Choosing an overall ... ". Please recheck this for clarity.

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