Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models

Many thanks for your encouraging and thoughtful feedback! We respond to your questions below.

An intuitive application of our bounds to downstream tasks:

Inspired by your comments, we provide another intuitive application of our bounds. In this case, to a downstream scientific task. Our token-level generalization bounds are particularly descriptive of antibody design in biology. An antibody sequence is usually composed of 20 different amino acid tokens to bind to a target of interest. In therapeutic antibody design, biologists propose mutations to existing antibody sequences by changing the amino acid tokens at specific positions in the sequence. Recent works have shown that LLMs pretrained on large antibody datasets can be used to propose mutations conditioned on starting antibody sequences. Our token-level generalization bounds match the antibody design setting by bounding the expected next amino acid token negative log likelihood averaged over training contexts that serve as starting sequences for iterative mutations. In the table below, we show that language models based on the Mistral 7B architecture pretrained on a processed subset of the Observed Antibody Sequences (OAS) from scratch achieves non-vacuous token-level generalization bounds:

Using generalization bounds for model optimization:

It is indeed possible to integrate the elements of our bounds directly into the training regimen to enhance model generalization directly. One example is quantization-aware training, which we in fact already use for post-training quantization of all pretrained models except the LLaMA models. In this scenario, we wish to map the pretrained weights of the neural networks into a significantly smaller number of quantization clusters. The quantized vector $\hat{w} = [\hat{w}_1,\dots,\hat{w}_d]$ can be constructed from the original weights vector $w = [w_1,\dots,w_d]$ by assigning these weights to different clusters $c = [c_1,\dots c_L]$, where $\hat{w}_i =c_q$ such that ${q= \operatorname{argmin}_k |w_i-c_k|}$. The quantization clusters $c$ are learned alongside $w$, such that we optimize the empirical risk and the compressed size of the model as well. We will add more details about our quantization-aware training procedure in the appendix. The focus of our work however is to compute non-vacuous generalization bounds that we can use to understand when and why LLMs generalize and provide a prescription for how to design better models in practice.

We provide additional experiments in the general response showing that not only is the quantity we bound predictive of downstream performance, but the bounds themselves positively correlate with downstream performance. We also demonstrate that the trade-off between the empirical risk and the compressed size of the model highly depends on the compression scheme and the size of the training data. The additional figures can be found in the Rebuttal PDF.

Inspired by your feedback, we also made sure to include additional intuitive explanations and details throughout the revised manuscript to make sure that the paper's content is accessible to readers.

Thank you again for your supportive and positive review. We made a significant effort to revise our paper and run additional experiments in light of your feedback. 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 or comments.

We really value your thoughtful and supportive feedback! We provide several additional results inspired by your comments.

An intuitive application of our bounds to downstream tasks:

Inspired by your comments, we provide another intuitive application of our bounds. In this case, to a downstream scientific task. Our token-level generalization bounds are particularly descriptive of antibody design in biology. An antibody sequence is usually composed of 20 different amino acid tokens to bind to a target of interest. In therapeutic antibody design, biologists propose mutations to existing antibody sequences by changing the amino acid tokens at specific positions in the sequence. Recent works have shown that LLMs pretrained on large antibody datasets can be used to propose mutations conditioned on starting antibody sequences. Our token-level generalization bounds match the antibody design setting by bounding the expected next amino acid token negative log likelihood averaged over training contexts that serve as starting sequences for iterative mutations. In the table below, we show that language models based on the Mistral 7B architecture pretrained on a processed subset of the Observed Antibody Sequences (OAS) from scratch achieves non-vacuous token-level generalization bounds:

Correlation between our bounds and the performance on downstream tasks:

Following your suggestion, we compute the direct correlation between our bounds and downstream performance on the tasks reported in Table 6 of the paper. In Figure 1(left y-axis) reported in the attached Rebuttal PDF, we plot the average zero-shot error (Error), defined as 1 - the accuracy, and the perplexity (PPL) achieved by GPT2 small, medium and large on downstream tasks, as reported in Radford et al. [1]. On the right y-axis, we plot the token-level bounds achieved by the GPT2 models with different sizes on the OpenWebText dataset that they were partially trained on. Our token-level BPD bounds achieve 98.9% and 99.4% correlation with the downstream perplexity and error, respectively, and are indeed predictive of generalization on downstream tasks. Given the significance of these results, we included them in our revised manuscript.

The distribution over next tokens being diffuse:

As you mentioned and as shown in Figure 2 (middle), the average entropy of the next token distribution conditioned on fixed contexts decreases for later token positions compared to the beginning of text token. Therefore, if we have a fixed prompt, this distribution might be less diffuse for the first token that we predict given the prompt. However, the main takeaway from this experiment actually pertains: even for later token positions, the distribution does not collapse entirely and the entropy is non-zero. For instance, if the prompt is composed of 127 tokens, then the average entropy for the next token is equal to 3 bits after transformation to the logarithm base 2, which corresponds to 8 choices for the next tokens. The number of choices is far greater than a single predetermined token, which implies that it is infeasible to make predictions on new resampled tokens from the empirical distribution alone. Therefore, non-vacuous token-level bounds are indicative of generalization beyond the training data.

Thank you again for your detailed review. We put a significant effort into our response and would appreciate it if you could consider raising your score.

References:

[1] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are 454 unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.

We thank you for your thoughtful and supportive feedback! We respond to your questions below.

Effect of finetuning on downstream task performance:

As per your suggestion, we run additional experiments where we finetune the GPT2 large model (774M parameters) – pretrained on the WebText dataset – on multiple downstream tasks: grade school math, stack exchange and planning. These datasets are publicly available on HuggingFace. We report the results in the table below. While finetuning visibly improves the model’s performance on the downstream task, it negatively affects the performance on the pretraining dataset and hence leads to a worse upstream bits-per-dimension (BPD) bound since the compressed size of the model remains approximately constant.

Trade-off between the empirical risk and the compressed size of the model:

Our generalization bounds, as described by Equation 2, can be conceptually written as:

$\text{Expected Risk} \leq \text{Empirical Risk} + \sqrt{\text{Compressed Model Size} / \text{Train Data Size}}.$

Therefore, the trade-off between the empirical risk and the compressed model size heavily depends on the compression approach and the size of the dataset. Having a dataset that contains a higher number of tokens puts a bigger emphasis on the empirical risk compared to the compressed model size, and vice versa. Likewise, if the compression approach is very aggressive, it will significantly reduce the compressed model size while potentially causing a deterioration of the empirical performance. The rate at which each element of the bound changes dictates whether the bound will increase or decrease for larger models. We demonstrate this effect in Figure 2 in the attached Rebuttal PDF, where we see that an aggressive compression scheme for the GPT2 models consisting of pretraining them in restricted SubLoRA spaces with 25k parameters only then quantizing them leads to an improvement in the token-level bounds as we increase the size of the original model. In contrast, quantizing the LLaMA2 models using QuIP# maintains a good empirical performance but does not reduce the compressed size significantly compared to an approach like SubLoRA, therefore the bounds deteriorate as the models become larger. If the Amber dataset contained more tokens, one would expect that the bounds would improve for larger models given the improvement in the empirical risk.

An intuitive application of our bounds:

In addition to the above experiments, we ran additional experiments for the antibody design downstream task to provide more intuition on the quantity that we bound .In fact, biologists propose mutations to existing antibody sequences by changing the amino acid tokens at specific positions in the sequence. Recent works have shown that LLMs pretrained on large antibody datasets can be used to propose mutations conditioned on starting antibody sequences. Our token-level generalization bounds match the antibody design setting by bounding the expected next amino acid token negative log likelihood averaged over training contexts that serve as starting sequences for iterative mutations. We provide results in the general response showing that these bounds are non-vacuous.

Background and proof sketch of the main theorem:

We provide a sketch of the proof in Appendix B.1, and it consists of three components: set up the empirical loss as the average of a martingale difference sequence, apply Azumas inequality for each hypothesis assigning failure probability proportional to the prior p(h), and then apply a union bound to relate the concentration around the many for individual hypotheses to any given hypothesis that can depend on the training data. We take your point that it would be beneficial to have this summary and some additional background in the main text, to provide additional context for the bounds we use. We will update the paper accordingly.

We provide additional experiments in the general response showing that not only is the quantity we bound predictive of downstream performance, but the bounds themselves positively correlate with downstream performance.

Thank you again for your supportive feedback. We made a significant effort to address your comments and run several new experiments inspired by your feedback, which we believe have improved our paper. We would appreciate it if you would consider raising your score in light of our response. We would be happy to engage in additional discussion if there are further questions.