Compute-Optimal LLMs Provably Generalize Better With Scale (2504.15208v1)

Abstract: Why do larger LLMs generalize better? To investigate this question, we develop generalization bounds on the pretraining objective of LLMs in the compute-optimal regime, as described by the Chinchilla scaling laws. We introduce a novel, fully empirical Freedman-type martingale concentration inequality that tightens existing bounds by accounting for the variance of the loss function. This generalization bound can be decomposed into three interpretable components: the number of parameters per token, the loss variance, and the quantization error at a fixed bitrate. As compute-optimal LLMs are scaled up, the number of parameters per data point remains constant; however, both the loss variance and the quantization error decrease, implying that larger models should have smaller generalization gaps. We examine why larger models tend to be more quantizable from an information theoretic perspective, showing that the rate at which they can integrate new information grows more slowly than their capacity on the compute-optimal frontier. From these findings we produce a scaling law for the generalization gap, with bounds that become predictably stronger with scale.

• The paper introduces a novel empirical Freedman inequality that tightens the generalization bound for compute-optimal LLMs.
• It demonstrates that increased model scale reduces per-token loss variance, thereby narrowing the generalization gap.
• By linking quantization techniques with scaling laws, the study offers practical insights into enhancing LLM compressibility and efficiency.

This paper investigates why larger LLMs tend to generalize better, focusing specifically on models trained in the "compute-optimal" regime described by the Chinchilla scaling laws (2203.15556). The core idea is to develop and analyze generalization bounds for the pretraining objective (next-token prediction negative log likelihood - NLL) that can explain the empirically observed small generalization gap (difference between training and test loss) and its behavior with increasing model scale.

The analysis breaks down the expected test error into three components: irreducible error (inherent randomness in data), approximation gap (training error), and generalization gap (difference between test and train error). The paper focuses on bounding the token-wise generalization gap, building upon prior work (2402.06073) that used martingale theory for non-IID sequence data.

A key contribution is a novel, fully empirical Freedman-type martingale concentration inequality (Theorem 3.1). Unlike standard Azuma-Hoeffding inequalities which depend only on the range of the loss, or Freedman's inequality which depends on the (often unknown) conditional variance, this new bound incorporates a fully empirical variance proxy term, denoted $\Sigma$.

• Empirical Freedman Inequality (Theorem 3.1): For a martingale difference sequence $X_k - E[X_k | \mathcal{F}_{k-1}]$, this theorem bounds the average difference:

  $$\frac{1}{n}\sum_{k=1}^n\big(E[X_k\mid\mathcal{F}_{k-1}]-X_k)\big) \le \Delta C + \Sigma\sqrt{C}$$

  where $C$ relates to complexity (like $\log(1/\delta)/n$), $\Delta$ is a bound on the loss difference, and $\Sigma$ depends on the empirical variation of the sequence relative to a predictable quantity $Y_k$. Crucially, $\Sigma$ can be estimated directly from training data. When $\Sigma$ is small, the bound tightens significantly, moving from a $\sqrt{C}$ dependence towards a linear $C$ dependence.

This inequality is then incorporated into a generalization bound for LLMs (Theorem 3.4). The bound requires quantizing the model to $b$ bits per parameter to make the hypothesis class finite and relate complexity to model size. It also uses prediction smoothing (mixing model output with a uniform distribution) to bound the maximum possible NLL loss ($\Delta$).

• Main Generalization Bound (Theorem 3.4): With probability $1-\delta$, the token-wise population risk of a smoothed, quantized model ($R_{sq}$) is bounded by:

  $$R_{sq} \le \hat{R}_h + \mathcal C \log V + \Sigma \sqrt{\mathcal C} + \sqrt{2\mathcal C} + (\hat{R}_q-\hat{R}_h)$$

  • $\hat{R}_h$: Empirical training loss of the original model $h$.
  • $$\mathcal C = (\frac{N}{D})b\log2 + \frac{1}{D}\log\frac{|K|}{\delta}$$: Complexity term. $N$ is parameters, $D$ is tokens. In the Chinchilla regime, $N/D$ is constant ($\approx 1/20$). $b$ is bits per parameter for quantization.
  • $\log V$: NLL of random guessing (Vocabulary size $V$). $\mathcal C \log V$ is often the largest term.
  • $\Sigma$: Empirical loss variation term (from Theorem 3.1), computed on training data using the quantized model $q$.
  • $\sqrt{2\mathcal C}$: Cost associated with prediction smoothing.
  • $(\hat{R}_q-\hat{R}_h)$: Quantization gap (difference in training loss between original and quantized model).

The paper evaluates this bound empirically using the Pythia model suite (2304.01373) trained on the Pile dataset (2009.07311), selecting checkpoints along the compute-optimal frontier ($N/D \approx 1/20$).

• Empirical Findings:
  • The empirical loss variation term $\Sigma$ decreases predictably with model size $N$, roughly as $c + k N^{-0.5}$. Larger models have lower loss variance per token.
  • The quantization gap $(\hat{R}_q-\hat{R}_h)$ for a fixed bitrate (e.g., $b=4$ using GPTQ (2210.17323)) also tends to decrease as models get larger.
  • As $N$ and $D$ increase along the compute-optimal path, the overall bound $R_{sq}$ decreases, mirroring the decrease in the empirical training loss $\hat{R}_h$. This supports the idea that the theory captures the trend of better generalization with scale.
  • The $\mathcal C \log V$ term dominates the bound value. Since $N/D$ is constant and $b$ is fixed, this term doesn't inherently decrease with scale in this formulation.

The paper further argues that the complexity term $\mathcal C$ might actually decrease with scale, making the bound even tighter for larger models. Two arguments are presented:

1. Improved Quantizability (Appendix B): Drawing on theory from QuIP quantization (2307.13304, 2310.16669), it's argued that if the Hessian spectrum decays rapidly (which empirical estimates suggest), larger models become inherently more quantizable. This means a smaller $b$ could achieve the same quantization gap, reducing $\mathcal C = (N/D) b \log 2$. The analysis involves Hessian incoherence and relies on theoretical quantization schemes (LDLQ with random rotations) that are computationally expensive but illustrate the principle. Empirical estimates of $\mathrm{Tr}(H^{1/2})$ suggest it grows sublinearly with $N$, supporting the idea that $b$ can decrease with scale.
2. Information-Theoretic Complexity (Section 5.2): Using prequential coding arguments [Rissanen 1984, Dawid 1984], the information content $K(h)$ of the model $h$ (related to its description length) is upper-bounded by the cumulative reduction in loss during training (area between the initial and final loss curves). For models following scaling laws $R(N,D) \approx E + AN^{-\alpha} + BD^{-\beta}$, this analysis yields $K(h) \propto D^{1-\beta}$. Since $\beta < 1$, the information stored grows sublinearly with the data $D$. If complexity $L(h) \propto K(h)$, then $\mathcal C \propto L(h)/D \propto D^{-\beta}$. This implies complexity decreases with scale along the compute-optimal path ($D \propto C^{0.5}$). Evaluating the bound with this information-theoretic complexity shows it decreases predictably with scale ($O(D^{-\beta/2})$, dominated by the smoothing term), although the absolute values are looser than the quantization-based bounds for current model sizes.

Conclusion: The paper provides theoretical and empirical evidence that LLMs trained compute-optimally generalize better with scale because:

1. The ratio of parameters to data points ($N/D$) remains constant.
2. The per-token loss variance ($\Sigma$) decreases.
3. The models become more compressible/quantizable (effectively smaller $b$ or sublinear $K(h)/D$), reducing the effective complexity $\mathcal C$.

These factors combine within the derived generalization bound to predict a shrinking generalization gap as models scale. The work introduces a novel empirical Freedman inequality applicable to martingales and connects scaling laws, generalization theory, quantization, and information theory to explain a key phenomenon in LLMs. Limitations include the pessimism of the smoothing term and the need for deeper explanations of why loss variance and compressibility behave as observed.