Compute-Optimal LLMs Provably Generalize Better with Scale

We thank you for your detailed review. We address your questions one by one.

For the most closely related work see Lotfi et. al (b, 2024) which we cite and compare our generalization bounds to. As stated at the end of the background section and additionally in the related work section, this work uses Azuma’s inequality producing bounds that scale $\sqrt{C} \log V$ whereas our empirical Freedman inequality based bounds scale as $C\log V + \sqrt{C}\Sigma$. $C\log V$ dominates the other contributions to the bounds and as $C\approx 1/9$, our bounds are roughly 3x tighter than those derived using Azuma’s inequality for the same sized model and dataset.

You mention an ‘extreme’ dependence on the number of parameters of the model. One of the crucial insights that we present here not previously explored in other work is that despite the conventional wisdom, this linear dependence (in the case of using quantization) is not that extreme due to the way that large models are scaled up in both model size and data size. The ratio of parameters to data points is a constant along the compute optimal frontier, and the decreasing information content per parameter means that this ratio becomes even smaller with scale. We do not see how this forms substantive critique of the paper when competing alternatives do not exist within the research literature. If other approaches had been demonstrated on similarly sized LLMs with a superior scaling that would be a different story, but this is not the case.

As mentioned to another reviewer, the evaluation on just the pythia family of models is by necessity. In order to perform the analysis and evaluation of generalization bounds that we do in the paper, we need to have access to both the model checkpoints and unbiased random samples from the data that the model was trained on. To the best of our knowledge Pythia is the only set of models spanning multiple orders of magnitude in parameter count where this is the case. Almost all open source LLMs do not provide their training data, and without that training data we cannot conduct the analysis we do here. While it would be great if additional model families were available that is not currently the case. This is a theory paper and we do not consider it within the scope of such a project to undertake the training of multi billion parameter LLMs. Respectfully, we do not consider this to be a fair reason to advocate for rejecting the paper.

In section 5 we make the argument based on the functional form of the scaling law and prequential coding that the information content in the model grows sublinearly in the number of parameters (or the number of data points) in the compute optimal regime. For an overview of section 5, the argument can be summarized as follows. Using prequential coding, the model and data can be coded jointly in a number of nats equal to the cumulative sum of the NLL values over the course of training (for a modern reference, see e.g. Blier and Ollivier 2018). Given the final model, the data alone has a code length equal to the sum of loss values of the final model, and by assumption this is well described by the scaling law. Thus we can use this value to estimate the complexity of the data given the final model. Through the symmetry of information property of Kolmogorov complexity, the difference between the two (up to logarithmic factors) is equal to $K(M)$, the complexity of the model.

Getting from equation 5 to equation 6, we simply evaluate the functional form of the scaling law $R(N,D) = E+AN^{-\alpha}+BD^{-\beta}$. $\sum_{k=1}^D R(N,k) = D(E+AN^{-\alpha}) + \sum_{k=1}^d B k^{-\beta}$. Subtracting the two terms, $\sum_{k=1}^D R(N,k) - DR(N,D) = B(\sum_{k=1}^D k^{-\beta}-D^{1-\beta})$. Asymptotically, $\sum_{k=1}^D k^{-\beta} \sim \frac{1}{1-\beta}D^{-\beta}$ therefore the difference is $O(D^{1-\beta})$. Does this clear up your understanding of the section? And if not, please let us know which parts you don’t follow so we can clarify them, particularly as you listed this as a principal concern in your review.

“From what I understand, the authors claim that their bound improves as the number of parameters increases only if the size of the dataset increases as well. “ Yes. The information content in the model grows sublinearly while the dataset size grows linearly, leading to a shrinking complexity term and thus improving bounds. We believe this is a highly significant result as such observations have not been made in the literature.

Thank you for addressing my questions in great detail.

• Concerning the dependence in parameters, I respectfully stand by my opinion: I do believe that the dependence is extreme, and hence question the quality of this bound. Its main strength in my opinion lies in the fact that it is empirical.
• I do agree with your point on the evaluation, given that your explanation is satisfactory and that this has not raised major concerns in other reviews.
• The discussion concerning section 5 is of little help and seems to require theoretical elements with which I am not familiar. Other reviewers have also stressed a lack of clarity: this should raise your attention on this point. Please improve the clarity of this section.
• Concerning the scaling of the bound, I stand by my point: I still believe that the title of the paper is clearly claiming something stronger than what you show. The fact that under certain assumptions, the complexity-to-datapoints ratio tends to $0$ is a well studied fact in learning theory and is not surprising - even though your proof techniques are novel. If you consider a linear model with sparsity growing at a sufficiently slow speed, such a result would hold, but I doubt that this would be consider as a proof of "larger models generalising better". From the title, I would have expected that for a fixed datasize, increasing the number of parameters leads to better performance - see for instance http://proceedings.mlr.press/v97/brutzkus19b/brutzkus19b.pdf or https://arxiv.org/abs/1810.02032 for the analysis of such a case.
• Concerning my concerns with related work of Lotfi et al. (2024a), I thank you for the clarity of your response. I had missed the second submission of this article and have reviewed the discussion. This fully answers my question.

To sum up, my main concerns lie in (1) the strength of your generalisation bounds, (2) the clarity of your work and (3) the choice of title. To make my comment actionable, I suggest the following:

• To solve problems (1) and (3), please reconsider the strength of your results by changing the title and maybe amending your article to include a discussion of what I consider to be weaknesses.
• To solve problem (2), I would suggest a rewriting of section 5 which would include more intuitions and explanations, and the addition of a supplementary that gives some background.

I am not able to raise my score yet but am in no way hostile to raising it during the rest of the discussion period. I will remain active and discuss these issues, since I believe this work has to be improved but builds on solid ground.

We thank you for your detailed and thoughtful review. We’d like to address your questions in turn and hopefully clear up any questions you may have.

Regarding the main mathematical results of the paper we would like to clarify a few things that may not be immediately apparent. The random variables $Y_k$ serve as a proxy for the conditional mean $\mathbb{E}[X_k|\mathcal F_{k-1}]$. We don’t know the mean (because we want to be able to empirically evaluate the bounds) but in its place we can substitute any deterministic function of quantities that are already known (at time k-1), or in other words $Y_k$ is $\mathcal F_{k-1}$ measurable. As $Y_k$ can be any $\mathcal F_{k-1}$ measurable function, this is why it has no additional definition in theorem 3.1.

As an example, one choice of $Y_k$ and the one we use in evaluating our bounds is choosing $Y_k$ as the mean NLL for token k according to the model taken over samples from the model distribution (conditioned on tokens <k), whereas $X_k$ is the the model NLL for the token drawn from the actual data distribution. Crucially, this model mean can be computed without knowing what token will come next. This choice of $Y_k$ is implicit in theorem 3.2 and 3.4, though we recognize that split between the main text and the appendix this detail may be less forthcoming, and perhaps especially so since in contrast to theorem 3.1, in 3.2 and 3.4 $X_k,Y_k$ are the token k and the model sampled token k. To address this we will add a small amount of additional exposition for this point to appendix A.5 and to theorem 3.2 in the main text.

Thus the with probability $1-\delta$ statement refers to the sampling of the token sequence ($X_k$ in theorem 3.4), from which both the per token loss values and the model mean loss values can be derived.

On the finite set $K$, your observation is appropriate. The discrete set $K$ is chosen to enable optimization over the scale variable $t$ in Theorem A.5. In the proof of A.5 (“Inequality Sketch”) we sketch out how the proof would look if one could freely optimize this variable without regards to making an additional union bound, and why the union bound is necessary. The optimal scale for $t$ depends on how closely $X_k$ concentrates to the mean, however theorem A.3 holds for a single value of $t$ rather than $t$ itself being a random variable. To circumvent this one can quantize the values of $t$ and optimize over them. At $|K|=1000$ the contribution to the complexity term the $(1/D)\log |K|$ factor (see line 264) is completely inconsequential in comparison to the term depending on the model parameters. $|K|=1000$ was chosen not to be larger so as not to slow down the evaluation of the bounds by too much and at that size $\Sigma$ has already converged to sufficient precision. We will add these details to the appendix on your suggestion.

On $Y_k$ and $\Delta$. As the NLL is strictly positive on the finite token distribution, $Y_k-X_k > -X_k >-\Delta$ where $\Delta$ is an upper bound on the per token NLL. The per token NLL is bounded in theorem 3.3 by introducing label smoothing, and $\Delta$ is optimized to be $\log V+\log((1+C)/C)$, slightly larger than random guess NLL of $\log V$. The reason we cannot take $Y_k=X_k-\Delta+\epsilon$ is that then $Y_k$ would not be a $\mathcal F_{k-1}$ measurable random variable and instead would be a $\mathcal F_k$ measurable random variable (it would be a function of unknown information and thus random even given context <k).

We evaluate our generalization bounds on the compute optimal frontier (shown in figure 1). Though large models trained by companies can and do sometimes differ in this choice due to other reasons, the largest and most expensive runs tend to be close to the compute optimal curve (see e.g. LLAMA 3 405B). Abstractly one can consider what would happen if models were trained on far more data or far less data than they do currently, but practically speaking model size and dataset size are scaled together. It is our view that the choice of how compute is expended between these two is a critical component of the algorithm, and thus evaluating the algorithm that produces LLMs requires respecting this choice. In this regard we break with prior work on generalization bounds and attending to this aspect provides motivation for the current work. For checkpoints trained with more data than optimal, their bounds will be tighter in a predictable way and likewise looser for checkpoints trained with less data.

-- continued --

We thank you for your engagement and thoughtful suggestions.

Returning to theorem 3.1, the randomness is over $(X_k,Y_{k})_{k=1}^n$ and so the inequality is not random. We believe that we have not stated in any of the relevant theorems that the randomness is only over $X_k$ though we recognize the ambiguity.

To clarify the statement of the theorem, we have reworded it as: (forgive the spacing as OpenReview's latex rendering is somewhat idiosyncratic)

"Let $(X_k)_{k=1}^n$

and $(Y_k)_{k=1}^n$

be sequences of random variables adapted to a filtration $(F_k)_{k = 0}^{n-1}$,

where each $X_k$ is $F_k$-measurable and each $Y_k$ is $F_{k-1}$-measurable. Assume the scaled differences $A_k = (Y_k - X_k)/\Delta$ are bounded below by $-1$ for some $\Delta > 0$. Let $K$ be a finite subset of $(0,1)$. Then, with probability at least $1-\delta$ over the probability space filtered by $(F_k)_{k = 0}^{n-1}$ …"

We believe this clears up any impreciseness in the statement of the theorem and we hope this is to your satisfaction.

In terms of the title, while we have a different perspective we can understand how it could be interpreted as stating something different from what we had intended. If you're willing to raise your score to support acceptance, we would be happy to commit to changing the title to something along the lines of "Compute Optimal LLMs Provably Generalize Better with Scale ” which less ambiguously is in line with the results of our paper.

For further empirical validation, we have added a table comparing the results of our bounds vs using bounds derived from Azuma’s inequality, employed in Lotfi et al (b, 2024). Leveraging our empirical Freedman concentration inequality, we achieve a substantially lower bounded loss value over equivalent Azuma based bounds, nearly halving the gap from the training loss across the different model sizes.

We really appreciate your support and suggestions. It would be a pity to have to go through a whole new review cycle for changes that, while important, can be straightforwardly resolved.

We really appreciate your supportive and thoughtful review! We have taken your feedback and will incorporate additionally summary for the bounds we derived in the revised version of the main text. For the placement of figure 2 and section 4 between section 3 (deriving the bounds) and section 5 (compressibility and information growth), this was a conscious choice. We felt that if we did not unpack the relevant components of theorem 3.4 with their relative contributions that it would detract readability and add additional load on the reader. In section 5 we explore a different way of bounding the complexity term and thus present those comparative results there. We appreciate the feedback however and will consider ways of improving the flow of the paper. With regards to “Beyond Chinchilla-Optimal…” as the inference costs shift the compute optimal curves to smaller ratios of parameter counts and tokens, the generalization bounds would improve. In the paper we stuck with the original and simpler problem framing, but incorporating recent works looking at modifications to the laws such as the paper you described would be a good direction to incorporate into future work.

We thank you for your thoughtful and constructive feedback. Below we address your specific questions in detail.

You question the practical merit of decomposing the generalization bound, particularly in relation to the demonstrations in Figures 2 (right) and 3 (center). We respectfully submit that this decomposition provides several essential insights. By separating the bound into its constituent components, we can precisely analyze how each element scales with increasing model and data size. The decomposition also reveals which components of the bound (e.g., the smoothing cost) present the most promising opportunities for improvement, enabling more targeted future research efforts. Additionally, this structure allows for meaningful comparisons with existing generalization bounds. While one might initially expect our Freedman-like inequality to be dominated by the variance term (similar to the O(√C log V) bounds from Azuma's inequality), our decomposition reveals a notably different behavior.

The exponent $\beta$ from the scaling law tends to be around $\beta \approx 0.37$ as observed in prior work. This value is not constrained by our work, however the limited information absorption by the model implied by the scaling law produces a generalization gap of $\tilde{O}(D^{-\beta})$ (neglecting the suboptimal smoothing term).

Regarding the use of the Pythia family of models and the potential for extending this analysis, we would like to clarify several points. Our choice of the Pythia family was indeed motivated by the unique availability of their complete training data. To our knowledge, they represent the only publicly available model family that spans orders of magnitude of compute while providing access to the training data. We emphasize that most of our analysis, particularly in Section 5.3, remains valid independent of the specific empirical behavior of the model, provided the Chinchilla scaling laws hold. The primary model-specific component is the behavior of the loss variation Σ, which is not directly derived from the scaling laws. While the loss variation behavior was previously unknown, we have strong reason to believe that it will exhibit similar scaling properties across large language models trained on comparable text corpora.

We appreciate your attention to detail regarding formatting issues. We will correct references to the appendix and update the equation citation format.

Thank you again for your thoughtful review. We would appreciate it if you could consider raising your score in light of our response. We are also happy to answer any additional questions.

We thank you for your review. We provide some additional context and clarifications in response to your questions.

Regarding the evaluation on the pythia family of models we would like to make a few points. In order to perform the analysis and evaluation of generalization bounds that we do in the paper, we need to have access to both the model checkpoints and unbiased random samples from the data that the model was trained on. To the best of our knowledge Pythia is the only set of models spanning multiple orders of magnitude in parameter count where this is the case. Almost all open source LLMs do not provide their training data, and without that training data we cannot conduct the analysis we do here. While it would be great if additional model families were available that is not currently the case. This is a theory paper and we do not consider it within the scope of such a project to undertake the training of multi billion parameter LLMs. Respectfully, we do not consider this to be a fair reason to advocate for rejecting the paper. We also want to highlight that for a paper on generalization bounds our empirical evaluation is actually quite strong: it is highly novel to be evaluating bounds on models with billions of parameters. In fact, to the best of our knowledge, there are only a couple papers about generalization bounds on large language models at all (e.g., Lotfi et. al 2024), and our bounds are significantly tighter and for larger models. There are no papers, to our knowledge, that evaluate bounds that match scaling law behavior.

Regarding the loss variation $\Sigma$, on line 189 (in the broader context of the concentration inequality) we state that it can be upper bounded in terms of $\sqrt{(1/n)\sum_k (X_k-Y_k)^2}$ and in the more narrow context of generalization bounds for autoregressive language models on line 269. With this upper bound, it can be interpreted as an estimate of the average variance of the loss at each timestep, but with the mean replaced by the mean over the model predicted tokens, therefore increasing it slightly. Intuitively as the variance of this distribution decreases, the random variables will be more concentrated around the mean.

While we agree that it would be highly interesting to understand complexity term changes for other datasets and modalities (it seems likely to be similar for different architectures on the same dataset based on similarities in the scaling laws), this isn’t possible at this time due to the paucity of publicly available LLMs and their training data. If we were to speculate, it seems plausible that image and video data (and perhaps to a lesser extent code) could have a substantially different model complexity $\mathcal C$ on the frontier of best performing models (for the given compute budget).

Thanks again for your questions. We would appreciate it if you could consider raising your score in light of these clarifications and additional context.