Predicting Task Performance with Context-aware Scaling Laws

TL;DR

This work introduces a context-aware scaling law that extends conventional neural scaling by modeling downstream task performance as a product of two saturating power laws in training compute $C$ and context length $n_{\text{pmt}}$, with a sigmoid penalty for exceeding the model's context limit $n_{\text{ctx}}$. The authors validate the framework by fitting it to 65{,}500 instances from extended-context Llama-2 models across arithmetic reasoning, common-sense reasoning, and machine translation, and demonstrate robust generalization across three orders of magnitude in compute and four orders of magnitude in context. The empirical results show accurate in-distribution fits and reliable extrapolation to longer contexts and different context-extension techniques, supporting practical guidance for designing efficient long-context LLMs. The approach provides an interpretable link between compute, context utilization, and downstream task performance, with publicly available code for replication.

Abstract

Scaling laws have transformed our understanding of large language models by linking upstream metrics like cross-entropy loss to design factors such as model size, training data, and compute. However, these conventional laws fail to capture downstream task performance, where context plays a critical role. In this work, we propose a straightforward, interpretable framework that jointly models downstream performance as a function of the training compute and the provided context. We empirically validate our framework by fitting it on the observed downstream performance of extended-context variants of Llama-2-7B and Llama-2-13B across 65,500 unique instances spanning three tasks: arithmetic reasoning, common sense reasoning, and machine translation. Our results demonstrate that our framework accurately models in-distribution downstream performance, generalizes across three orders of magnitude in training compute, and reliably extrapolates performance as the amount of context increases. These findings offer valuable insights into the interplay between training compute and context utilization, providing guidance for designing more efficient long-context LLMs for diverse downstream tasks. Our code is available at https://github.com/wang-research-lab/context-scaling.

Figures (6)

• Figure 1: Existing approaches ignore the impact of context length and predict an average performance level regardless of the number of in-context demonstrations. In comparison, our context-aware fit closely tracks the observed performance as additional context is provided.
• Figure 2: Contours of fits at $C=7.8 \times 10^{22}$ (red) and $C=1.5 \times 10^{23}$ (blue) for $n_\text{ctx}=8\text{k}$ on three tasks: arithmetic reasoning (left), common sense reasoning (middle) and machine translation (right).
• Figure 3: Contours of fits at $C=7.8 \times 10^{22}$ (red) and $C=1.5 \times 10^{23}$ (blue) for $n_\text{ctx}=128\text{k}$ on three tasks: arithmetic reasoning (left), common sense reasoning (middle) and machine translation (right). Held-out observations are colored in purple and green for Llama-2-7b and Llama-2-13b, respectively.
• Figure 4: Contours of our fit at $C=7.8 \times 10^{22}$ (left) and $C=1.5 \times 10^{23}$ (right) for the arithmetic reasoning task.
• Figure 5: Contours of our fit at $C=7.8 \times 10^{22}$ (left) and $C=1.5 \times 10^{23}$ (right) for the common sense reasoning task.