PaperlayerView exact revision on arXiv

Source-linked AI summary

Predicting Task Performance with Context-aware Scaling Laws

Kyle Montgomery, David Park, Jianhong Tu, Michael Bendersky, Beliz Gunel, Dawn Song, Chenguang Wang

arXiv:2510.14919v1cs.CLcs.AIcs.LG

TL;DR

Conventional scaling laws link upstream metrics to model resources but do not capture downstream performance where context matters. This paper jointly models downstream performance using compute and context, and reports accurate fits that generalize across compute and context ranges and extension techniques.

  • Problem

    Conventional scaling laws focus on upstream metrics, while downstream task performance can diverge from those metrics and depends on context.

    Source: §1 · p. 1

  • Method

    The framework directly models downstream performance as a function of training compute and provided context, combining saturating power-law terms with a penalty for exceeding the model’s context limit.

    Source: §1 · pp. 1–2

  • Results

    The framework accurately predicts Llama-2 downstream performance and generalizes across 3 orders of magnitude in training compute, 4 orders in context length, and different context-extension techniques.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Downstream performance benefits from increased compute and longer relevant context only up to saturation, offering guidance for long-context model design.

    Source: §6 · p. 9

    The formulation assumes performance scales with compute and context, which may not hold under extreme scaling or adversarial attacks, and omits data mix, alignment, and architecture factors.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses the gap between upstream scaling metrics and downstream task performance, which can diverge in real-world applications. It proposes a context-aware scaling framework and reports accurate fits and generalization across compute, context lengths, and context-extension techniques.

1 Introduction

  • Existing downstream-performance prediction methods are described as overly complicated or less interpretable, while upstream metrics can diverge from downstream performance.The paper motivates upfront estimates as a way to guide development and identify task emergence or saturation with fewer costly experiments.Source: §1 · p. 1
  • The framework directly models downstream performance using training compute and provided context, with saturating power-law terms and a penalty when context exceeds the model limit.Its design reflects the assumption that performance gains from more compute and longer relevant context eventually saturate.Source: §1 · pp. 1–2
  • The framework generalizes across 3 orders of magnitude in training compute, 4 orders of magnitude in context length, and different context-extension techniques.It also accurately predicts downstream performance for both Llama-2-7B and Llama-2-13B, including at longer contexts beyond a model’s context limit.Source: §1 · p. 2
  • The framework is fitted to extended-context Llama-2-7B and Llama-2-13B performance across 65,500 instances in arithmetic reasoning, common-sense reasoning, and machine translation.The contributions describe the evaluation as spanning three tasks.Source: §1 · pp. 1–2
  • The framework is presented as an interpretable tool for understanding compute, context, and downstream performance to guide long-context model design.Figure 1 contrasts context-aware fits that track observed performance as demonstrations increase with existing approaches that predict an average level regardless of context length.Source: §1 · p. 2, Figure 1 · p. 2

2 Background

This background introduces notation and context extension for Llama-2. Because long sequences are difficult to handle and the original context limit is 4k tokens, the paper uses YaRN-extended models for its study.

2.1 Notation

  • The paper denotes downstream task performance as P, non-embedding training compute as C, context limit as nctx, and input query or context length as npmt.The context limit is the maximum number of positional embeddings computed for a training sequence; npmt excludes generated tokens.Source: §2.1 · p. 2

2.2 Extending Llama-2’s Context Limit

  • Long-sequence attention has quadratic complexity, and models commonly pre-train on short sequences, so they can struggle to generalize beyond those lengths.The paper gives 4k and 128k tokens as examples and notes that Llama-2’s original 4k context limit is insufficient for studying performance across context lengths.Source: §2.2 · pp. 2–3

2 Background

  • The 12 scaling-curve checkpoints include official 4k Llama-2-7B and Llama-2-13B variants, with YaRN extension training tokens and compute included in D and C.The background introduces notation and the process of extending Llama-2’s context window.Source: §2 · p. 2, Table 1 · p. 3

2.2 Extending Llama-2’s Context Limit

  • YaRN extends Llama-2’s context limit by fine-tuning on sequences longer than the model’s pre-training context limit.The paper selects YaRN for its compute efficiency and empirical results compared with other techniques.Source: §2.2 · p. 3
  • On RULER’s 13 synthetic needle-in-a-haystack tasks, results suggest YaRN context extension is somewhat effective, and models can underperform at their extended context limit.The passage notes that training with a context limit well beyond the target evaluation range can improve performance within that desired range.Source: §2.2 · p. 3

3 Method

The method models downstream performance jointly over training compute and context length, penalizing prompts beyond the model’s context limit, and fits this form to three-task Llama-2 evaluations.

3 Method

  • Performance is modeled as the product of saturating power laws in compute and prompt length, with a sigmoid penalty when prompt length exceeds the context limit.The product encodes compute and context as complementary: a shortfall in either limits the benefit of the other.Source: §3 · pp. 3–4

3.1 Datasets

  • Evaluation covers 65,500 varied-length instances across arithmetic reasoning, common sense reasoning, and machine translation.The arithmetic and common-sense tasks use up to 511 demonstrations; translation uses same-source-language demonstrations and BLEU-4 scoring.Source: §3.1 · p. 4

3.2 Fitting Procedure

  • The fitting data are grouped by context length into records of compute, prompt length, context limit, and average metric value, then fit per task.Optimization uses a global search followed by local curve_fit refinement from the global estimate.Source: §3.2 · pp. 4–5

4 Empirical Results

The joint scaling framework fits observed task performance closely and generalizes across compute, longer contexts, and context-extension methods, while the penalty term matters near the context limit.

4 Empirical Results

  • Average prediction errors |P − ˆP| are 0.010 for arithmetic reasoning, 0.037 for common sense reasoning, and 0.007 for machine translation.The fit also models behavior around the prompt-length/context-limit boundary.Source: §4 · p. 5

4.1 Generalization along C

  • Across five test models spanning three orders of magnitude in compute, many errors fall near or below 5 points; generalization is stronger on arithmetic reasoning and translation than common sense reasoning.Generalization is weakest on Gemma-2-2B despite good generalization to Gemma-2-9B and Gemma-2-27B; errors tend to underestimate at small compute and slightly overestimate at large compute.Source: Table 5 · p. 6, §4.1 · p. 6

4.2 Generalization along npmt

  • Held-out-context prediction errors are 0.017 for arithmetic reasoning, 0.067 for common sense reasoning, and 0.006 for machine translation.The fits were refit while holding out observations above 10,000 tokens; Figure 3 shows contours at nctx = 128k and held-out model observations.Source: §4.2 · p. 6, Figure 3 · p. 6
  • On common-sense and translation tasks, performance can decline with longer context; the authors hypothesize misalignment with the extension-training mix, which is predominantly English PG-19 text.Source: §4.2 · pp. 6–7

4.3 Does the choice of context extension technique matter?

  • Prediction errors for positional-interpolation and YaRN extensions are similar, suggesting the choice of extension technique has little impact on downstream scaling properties.The models used different training mixes and token quantities, though their estimated compute was similar.Source: §4.3 · p. 7

4.4 Ablation over the penalty term

  • Without the sigmoid penalty, the fit underestimates performance when prompt length is within the context limit and overestimates it when prompt length exceeds that limit.Source: §4.4 · p. 7

5 Related Work

Related work develops upstream scaling laws and several approaches to downstream prediction, while prior context-length analyses focus largely on upstream loss. This paper positions context and compute jointly as predictors of downstream performance.

5 Related Work

  • Neural scaling laws relate upstream performance to model parameters, dataset size, or training compute, and have guided compute-optimal model design.Related work also extends scaling analyses to autoregressive image and video models and smoothly broken laws.Source: §5 · p. 7
  • Prior downstream predictors address emergence, transfer learning, multiple-choice tasks, benchmark aggregates, or intermediary metrics such as upstream loss.This work instead adds context length to a direct downstream-performance model.Source: §5 · p. 7
  • Earlier context-length scaling analyses examined upstream validation loss or in-context demonstrations, not the relationship between context length and downstream performance.The authors identify their work as the first to explicitly focus on that relationship and jointly consider context and compute.Source: §5 · pp. 7–8
  • Positional encodings affect extrapolation to longer sequences, motivating techniques that extend model context limits.Some positional-encoding methods offer limited extrapolation, while others do not.Source: §5 · p. 8
  • Context-extension approaches include attention modifications and positional-encoding rescaling; this work uses YaRN for Llama-2 extensions.The passage motivates YaRN by its compute efficiency and empirical results relative to other techniques.Source: §5 · p. 8

6 Conclusion

The paper models downstream performance jointly with training compute and context, then evaluates the framework across three task types. It reports accurate in-distribution fits and generalization across compute, context lengths, and context-extension techniques, with benefits from compute and relevant context saturating.

6 Conclusion

  • The framework fits in-distribution performance and generalizes across 3 orders of magnitude in non-embedding training compute, 4 orders in input context length, and other context-extension techniques.Experiments cover arithmetic reasoning, common-sense reasoning, and machine translation.Source: §6 · p. 8
  • Downstream performance benefits from increased compute and longer, relevant context, but only up to a saturation point.The authors say the framework offers interpretable insights that can guide long-context LLM design.Source: §6 · p. 9

Limitations

The framework’s assumptions and omitted factors constrain its scope, and its fits cover a narrow range of training compute that may not generalize far beyond that range.

Limitations

  • The assumed scaling of performance with compute and context may not hold under extreme scaling or adversarial settings such as many-shot jailbreaking.The authors identify these as limitations of the formulation.Source: Limitations · p. 9
  • Pre-training data mix, post-training and alignment, and architecture are not explicitly modeled, though they may affect fitted parameters.The paper suggests these factors may alter parameters without changing the structure of its equation.Source: Limitations · p. 9
  • The scaling curves cover a narrow training-compute range and may fail to generalize to models trained far beyond it.This is an explicit scope boundary for extrapolating the fitted curves.Source: Limitations · p. 9

A Dataset Details

The appendix specifies filtering, sampling, generation, and scoring procedures for arithmetic, common-sense, and translation datasets, generally using randomly selected training and test subsets.

A Dataset Details

  • GSM8K and MATH each use 511 training and 250 testing instances, scoring extracted final answers by numerical equivalence.Instances over 256 tokens are filtered out; inference allows up to 400 new tokens.Source: A Dataset Details · p. 17
  • AQUA-RAT uses 511 training and 250 testing instances, allows up to 5 generated tokens, and checks whether the returned choice matches the reference.Instances over 256 tokens are filtered out.Source: A Dataset Details · p. 17
  • DeepMind Math samples 511 training and 50 testing instances from each of 56 subsets, allowing up to 400 generated tokens and checking numerical equivalence.Instances over 256 tokens are filtered out.Source: A Dataset Details · p. 17
  • WinoGrande, ARC Easy, and ARC Challenge each use 511 training and 250 testing instances; choice outputs are checked against reference answers.Instances over 256 tokens are filtered out, and generation allows up to 5 new tokens.Source: A Dataset Details · pp. 17–18
  • CommonSenseQA, PIQA, OpenBookQA, SIQA, and HellaSwag use 511 training and 250 testing instances; generation allows up to 5 tokens except SIQA’s 1.Instances over 256 tokens are filtered out, and returned choices are compared with references.Source: A Dataset Details · pp. 17–18
  • The five WMT14 pairs—CS-EN, DE-EN, FR-EN, HI-EN, and RU-EN—use 511 training and 250 testing instances and are scored with BLEU-4.Instances over 256 tokens are filtered out; generation allows up to 256 new tokens.Source: A Dataset Details · p. 18

B Full Results

The section presents full aggregate accuracy results for arithmetic reasoning, common sense reasoning, and machine translation, alongside fitted-performance contours for each task at two compute levels.

Related papers

Loading 2510.14919v1…