Long-Context Benchmarks and Context Extension Methods

A model that claims to read a novel but forgets the first chapter by the last page has a context window in name only.

Eval, Forgetfully Long-Context AI Agent

1. The Gap Between Claimed and Effective Context Length

When a model provider advertises a 128K or 1M token context window, that number represents the maximum sequence length the model can accept as input. It does not guarantee that the model will accurately attend to, reason about, or recall information at every position within that window. Several factors contribute to degraded performance at long context lengths:

• Attention dilution: As sequence length grows, attention weights spread across more tokens. Information at positions that receive low attention scores is effectively invisible to the model, even though it is technically within the context window.
• Position bias: Models trained predominantly on shorter sequences develop a bias toward information near the beginning and end of the context, with reduced recall for content in the middle. This is the "lost-in-the-middle" phenomenon documented by Liu et al. (2023).
• Training distribution mismatch: If a model was pre-trained on 4K token sequences and then extended to 128K through positional extrapolation, its ability to reason over the extended range may be significantly weaker than at the original training length.
• Task complexity interaction: Simple retrieval (finding a specific fact) degrades more gracefully than complex reasoning (synthesizing information from multiple positions) as context grows.

This gap between claimed and effective context length has practical consequences. A developer who stuffs 100K tokens of context into a prompt, assuming the model will use it all, may get worse results than one who carefully selects 10K tokens of highly relevant content. Benchmarks in the following sections provide the tools to measure this gap empirically.

2. Needle-in-a-Haystack (NIAH) and Multi-Needle Variants

The Needle-in-a-Haystack (NIAH) test, popularized by Greg Kamradt in late 2023, is the simplest and most intuitive long-context benchmark. A synthetic "needle" (a specific fact, such as "The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day") is inserted at a controlled position within a large "haystack" of irrelevant text. The model is then asked to retrieve the needle.

By varying both the total context length and the needle position (from 0% to 100% depth), NIAH produces a two-dimensional heatmap showing retrieval accuracy across the entire context window. This reveals position-dependent failure patterns that a single aggregate accuracy number would hide.

2.1 Multi-Needle and Reasoning Variants

The basic single-needle test has been extended in several ways to increase difficulty:

• Multi-needle: Multiple facts are scattered throughout the context, and the model must retrieve all of them. This tests whether the model can attend to multiple positions simultaneously.
• Reasoning needles: Instead of retrieving a verbatim fact, the model must combine information from two or more needles placed at different positions to produce the correct answer.
• Adversarial needles: Distractor facts similar to the target needle are inserted at other positions, testing whether the model can distinguish the correct needle from plausible alternatives.

# Code Fragment 29.11.2: Needle-in-a-Haystack evaluation
# Key operations: needle insertion, position sweep, heatmap generation
import numpy as np
from openai import OpenAI

client = OpenAI()

def run_niah_test(
 model: str,
 context_lengths: list[int],
 depth_percents: list[float],
 needle: str = "The secret project code name is AURORA-7.",
 question: str = "What is the secret project code name?",
 expected: str = "AURORA-7",
 filler_text: str = None,
) -> np.ndarray:
 """Run NIAH across context lengths and needle positions."""
 results = np.zeros((len(context_lengths), len(depth_percents)))

 for i, ctx_len in enumerate(context_lengths):
 # Generate filler text to fill the context
 haystack = generate_filler(ctx_len, filler_text)

 for j, depth in enumerate(depth_percents):
 # Insert needle at the specified depth
 insert_pos = int(len(haystack) * depth)
 prompt = (
 haystack[:insert_pos]
 + f"\n{needle}\n"
 + haystack[insert_pos:]
 )

 response = client.chat.completions.create(
 model=model,
 messages=[
 {"role": "system",
 "content": "Answer based on the provided context."},
 {"role": "user",
 "content": f"Context:\n{prompt}\n\n{question}"},
 ],
 max_tokens=50,
 temperature=0.0,
 )

 answer = response.choices[0].message.content
 results[i, j] = 1.0 if expected.lower() in answer.lower() else 0.0

 return results

# Run across multiple context lengths and positions
context_lengths = [1_000, 4_000, 16_000, 32_000, 64_000, 128_000]
depth_percents = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]

scores = run_niah_test(
 model="gpt-4o",
 context_lengths=context_lengths,
 depth_percents=depth_percents,
)
print(f"Overall retrieval accuracy: {scores.mean():.1%}")

# pip install inspect-ai
from inspect_ai import eval
from inspect_ai.dataset import Sample
from inspect_evals.niah import niah

# Run NIAH across context lengths with one call
results = eval(
 niah(min_context=1_000, max_context=128_000, n_contexts=6, n_positions=7),
 model="openai/gpt-4o",
 log_dir="./niah_logs",
)
# Results include a 2D accuracy matrix and interactive log viewer

3. RULER: Parametric Context Utilization Benchmark

RULER (Hsieh et al., 2024) addresses the limitations of NIAH by providing a parametric benchmark with configurable task complexity. While NIAH tests simple retrieval of a single fact, RULER includes tasks that require reasoning, aggregation, and multi-step operations across the context. RULER defines four task categories:

1. Retrieval (NIAH variants): Single-key, multi-key, multi-value, and multi-query needle retrieval at controlled positions.
2. Multi-hop tracing: Variable-length chains where the model must follow a sequence of references (e.g., "X is stored in Y, Y is stored in Z; where is X?").
3. Aggregation: The model must count, sum, or find the most frequent item among values scattered throughout the context.
4. Question answering: Complex questions requiring information from multiple context positions.

The key insight of RULER is that model performance degrades at very different rates depending on task complexity. A model that achieves 95% on single-needle retrieval at 128K tokens may drop to 40% on multi-hop tracing at the same length. RULER's parametric design allows benchmarking at any target sequence length, producing "RULER curves" that show how each task type degrades with context.

# Code Fragment 29.11.5: RULER-style multi-hop tracing task
# Key operations: chain generation, context construction, evaluation
import random
import string

def generate_ruler_tracing_task(
 context_length: int,
 chain_length: int = 3,
 num_distractors: int = 20,
) -> dict:
 """Generate a RULER-style variable-length tracing task.

 Creates a chain: X is in box A, box A is in box B, ...
 The model must determine the final location of X.
 """
 # Generate unique variable names
 names = [
 "".join(random.choices(string.ascii_uppercase, k=4))
 for _ in range(chain_length + 1 + num_distractors)
 ]

 target = names[0]
 chain_locations = names[1:chain_length + 1]
 distractor_names = names[chain_length + 1:]

 # Build the chain statements
 chain_facts = []
 chain_facts.append(
 f"The item '{target}' is stored in container '{chain_locations[0]}'."
 )
 for k in range(len(chain_locations) - 1):
 chain_facts.append(
 f"Container '{chain_locations[k]}' is inside "
 f"container '{chain_locations[k+1]}'."
 )

 # Build distractor statements
 distractor_facts = []
 for name in distractor_names:
 loc = random.choice(distractor_names)
 distractor_facts.append(
 f"The item '{name}' is stored in container '{loc}'."
 )

 # Interleave chain facts at random positions among distractors
 all_facts = distractor_facts.copy()
 for fact in chain_facts:
 pos = random.randint(0, len(all_facts))
 all_facts.insert(pos, fact)

 # Pad to target context length with filler
 context = "\n".join(all_facts)
 context = pad_to_length(context, context_length)

 return {
 "context": context,
 "question": f"Where is '{target}' ultimately stored? "
 f"Follow the chain of containers.",
 "expected": chain_locations[-1],
 "chain_length": chain_length,
 }

# Evaluate across context lengths and chain depths
for ctx_len in [4_000, 32_000, 128_000]:
 for chain_len in [2, 3, 5, 7]:
 task = generate_ruler_tracing_task(ctx_len, chain_len)
 result = evaluate_model(task)
 print(f"Context: {ctx_len:>7,} | Chain: {chain_len} | "
 f"Correct: {result}")

4. LongBench v2: Multi-Task Long-Context Evaluation

LongBench v2 (Bai et al., 2024) provides a comprehensive, multi-task evaluation suite for long-context language models. Unlike synthetic benchmarks (NIAH, RULER), LongBench v2 uses real-world documents and tasks that reflect actual long-context use cases. The benchmark spans six task categories:

• Single-document QA: Questions answerable from a single long document (academic papers, legal contracts, technical manuals).
• Multi-document QA: Questions requiring synthesis across multiple documents concatenated into a single context.
• Summarization: Producing summaries of long documents that capture key information without hallucination.
• Few-shot learning: In-context learning with many demonstration examples, testing whether models can leverage large example sets effectively.
• Code completion: Completing code given a large repository context with cross-file dependencies.
• Synthetic tasks: Controlled tasks (including NIAH variants) for isolating specific capabilities.

LongBench v2 addresses several shortcomings of earlier benchmarks. It uses a "length-instruction-enhanced" evaluation protocol that controls for the confound between task difficulty and context length. Tasks are designed so that the relevant information truly requires the full context, preventing models from "cheating" by ignoring the long context and relying on parametric knowledge.

5. Context Extension Methods

Models trained with a fixed maximum sequence length face a hard boundary: they cannot process inputs longer than their training context. Positional encoding extension methods address this by modifying the position representation so that the model can generalize to longer sequences without full retraining. The key methods are:

5.1 Position Interpolation

Position interpolation (Chen et al., 2023) is the simplest extension method. Instead of extrapolating RoPE (Rotary Position Embedding) frequencies beyond the training range, it rescales all positions to fit within the original range. For a model trained on $L$ tokens extended to $L' > L$ tokens, position $i$ is mapped to $i \cdot \frac{L}{L'}$.

The intuition is that interpolation between known positions is more stable than extrapolation beyond them. A model trained on positions 0 through 4,096 can reasonably estimate what position 2,048.5 "looks like," but has no reliable basis for estimating position 8,192. After interpolation, a small amount of fine-tuning (typically 1,000 steps) adapts the model to the compressed position representation.

5.2 NTK-Aware Scaling

NTK-aware scaling (bloc97, 2023) observes that position interpolation compresses high-frequency RoPE components too aggressively, destroying local position information. Instead of applying a uniform scaling factor, NTK-aware scaling adjusts the RoPE base frequency:

Here, $\alpha = L'/L$ is the extension ratio and $d$ is the embedding dimension. This preserves high-frequency (local) positional information while stretching low-frequency (global) components to accommodate longer sequences. NTK-aware scaling often works without any fine-tuning at moderate extension ratios (2x to 4x).

5.3 YaRN (Yet another RoPE extensioN)

YaRN (Peng et al., 2023) combines the best ideas from position interpolation and NTK-aware scaling into a unified framework. It introduces a per-dimension scaling strategy that treats RoPE frequency bands differently based on their wavelength relative to the original context length:

• High-frequency dimensions (wavelength much shorter than training length): No scaling needed; these encode local position information that transfers directly.
• Low-frequency dimensions (wavelength comparable to or longer than training length): Apply NTK-aware scaling to stretch these for the extended context.
• Intermediate dimensions: Blend between no scaling and full scaling using a smooth ramp function.

YaRN also applies a temperature scaling factor to the attention logits to compensate for the increased entropy of attention distributions at longer sequence lengths. This correction prevents the "attention dilution" that would otherwise degrade performance.

# Code Fragment 29.11.4: Implementing YaRN-style RoPE extension
# Key operations: frequency scaling, dimension-wise interpolation
import torch
import math

def yarn_rope_scaling(
 dim: int,
 max_position: int,
 original_max: int = 4096,
 base: float = 10000.0,
 beta_fast: int = 32,
 beta_slow: int = 1,
) -> torch.Tensor:
 """Compute YaRN-scaled RoPE frequencies.

 Args:
 dim: RoPE embedding dimension
 max_position: target context length
 original_max: original training context length
 base: RoPE base frequency
 beta_fast: fast (high-freq) interpolation boundary
 beta_slow: slow (low-freq) interpolation boundary
 """
 scale = max_position / original_max

 # Compute original frequencies
 freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

 # Compute wavelengths for each frequency dimension
 wavelengths = 2 * math.pi / freqs

 # Determine per-dimension scaling using ramp function
 low_bound = original_max / beta_fast
 high_bound = original_max / beta_slow

 # Ramp: 0 for high-freq (no scaling), 1 for low-freq (full scaling)
 ramp = ((wavelengths - low_bound) / (high_bound - low_bound)).clamp(0, 1)

 # Interpolate between original freqs and NTK-scaled freqs
 ntk_factor = base * (scale ** (dim / (dim - 2)))
 ntk_freqs = 1.0 / (ntk_factor ** (torch.arange(0, dim, 2).float() / dim))

 # Blend: high-freq keeps original, low-freq uses NTK scaling
 scaled_freqs = (1 - ramp) * freqs + ramp * ntk_freqs

 return scaled_freqs

# Example: extend a 4K model to 128K
freqs_4k = yarn_rope_scaling(dim=128, max_position=4096,
 original_max=4096)
freqs_128k = yarn_rope_scaling(dim=128, max_position=131072,
 original_max=4096)

print(f"Scaling ratio: {131072 / 4096}x")
print(f"Low-freq stretch: {freqs_4k[-1] / freqs_128k[-1]:.1f}x")
print(f"High-freq stretch: {freqs_4k[0] / freqs_128k[0]:.1f}x")

6. Evaluation Pitfalls

Evaluating long-context models involves several subtle pitfalls that can produce misleading results:

6.1 Contamination in Long Documents

Long-context benchmarks that use real documents (academic papers, books, Wikipedia articles) risk contamination: the model may have seen these documents during pre-training and can answer questions from parametric memory rather than from the provided context. This inflates apparent long-context performance. Mitigation strategies include using post-training-cutoff documents, synthetic tasks, or "unanswerable question" controls where the context is replaced with irrelevant text to measure the parametric knowledge baseline.

6.2 Position Bias Effects

The "lost-in-the-middle" phenomenon means that where information appears in the context affects whether the model uses it. Evaluations that place key information only at the beginning or end of the context will overestimate performance. Robust evaluation requires systematic position sweeps (as in NIAH) or randomized placement with averaging.

6.3 Length vs. Difficulty Confounds

Longer contexts often correlate with harder tasks (more documents to synthesize, more complex reasoning chains). Separating the effect of context length from task difficulty requires controlled experiments where the same task is embedded in contexts of varying length with irrelevant padding. RULER's parametric design and LongBench v2's length-instruction protocol both attempt to address this confound.

# Code Fragment 29.11.7: Detecting contamination in long-context evaluation
# Key operations: baseline measurement, contamination detection
def measure_contamination_baseline(
 model: str,
 questions: list[dict],
 client: OpenAI,
) -> dict:
 """Measure how many questions the model can answer without context.

 If the model answers correctly without the long context,
 the question may be contaminated (answerable from parametric memory).
 """
 results = {"answerable_without_context": 0, "total": 0}

 for q in questions:
 # Ask without providing the document context
 response = client.chat.completions.create(
 model=model,
 messages=[
 {"role": "system",
 "content": "Answer the question. If you are unsure, "
 "say 'I don't know'."},
 {"role": "user", "content": q["question"]},
 ],
 max_tokens=100,
 temperature=0.0,
 )

 answer = response.choices[0].message.content
 if check_answer(answer, q["expected"]):
 results["answerable_without_context"] += 1
 results["total"] += 1

 contamination_rate = (
 results["answerable_without_context"] / results["total"]
 )
 print(f"Contamination baseline: {contamination_rate:.1%} of questions "
 f"answerable without context")
 return results

Lab: Reproducing RULER Curves

# Code Fragment 29.11.7: RULER curve reproduction for open-weight models
# Key operations: model loading, task generation, batch evaluation
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
from collections import defaultdict

# Load an open-weight long-context model
model_name = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
 model_name,
 torch_dtype=torch.bfloat16,
 device_map="auto",
 attn_implementation="flash_attention_2",
)

# Define evaluation grid
context_lengths = [4_096, 8_192, 16_384, 32_768, 65_536, 131_072]
task_types = ["single_niah", "multi_niah", "tracing", "aggregation"]
num_samples_per_cell = 20

# Generate and evaluate tasks
results = defaultdict(lambda: defaultdict(list))

for task_type in task_types:
 for ctx_len in context_lengths:
 for _ in range(num_samples_per_cell):
 # Generate task based on type and target length
 task = generate_ruler_task(
 task_type=task_type,
 target_length=ctx_len,
 tokenizer=tokenizer,
 )

 # Tokenize and generate
 inputs = tokenizer(
 task["prompt"],
 return_tensors="pt",
 truncation=True,
 max_length=ctx_len + 256,
 ).to(model.device)

 with torch.no_grad():
 outputs = model.generate(
 **inputs,
 max_new_tokens=64,
 temperature=0.0,
 do_sample=False,
 )

 response = tokenizer.decode(
 outputs[0][inputs.input_ids.shape[1]:],
 skip_special_tokens=True,
 )

 correct = evaluate_ruler_answer(
 response, task["expected"], task_type
 )
 results[task_type][ctx_len].append(correct)

# Print RULER curves
print(f"{'Task Type':<16} " + " ".join(
 f"{l//1000:>5}K" for l in context_lengths
))
print("-" * 60)
for task_type in task_types:
 scores = [
 np.mean(results[task_type][l]) for l in context_lengths
 ]
 row = " ".join(f"{s:>5.0%}" for s in scores)
 print(f"{task_type:<16} {row}")

# Code Fragment 29.11.7: Visualizing RULER curves with matplotlib
# Key operations: plotting, multi-line chart, annotation
import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(10, 6))

colors = {
 "single_niah": "#2196F3",
 "multi_niah": "#FF9800",
 "tracing": "#F44336",
 "aggregation": "#4CAF50",
}

labels = {
 "single_niah": "Single-Needle Retrieval",
 "multi_niah": "Multi-Needle Retrieval",
 "tracing": "Multi-Hop Tracing",
 "aggregation": "Aggregation",
}

x_labels = [f"{l//1000}K" for l in context_lengths]

for task_type in task_types:
 scores = [np.mean(results[task_type][l]) for l in context_lengths]
 ax.plot(
 range(len(context_lengths)), scores,
 marker="o", linewidth=2,
 color=colors[task_type],
 label=labels[task_type],
 )

ax.set_xticks(range(len(context_lengths)))
ax.set_xticklabels(x_labels)
ax.set_xlabel("Context Length (tokens)")
ax.set_ylabel("Accuracy")
ax.set_title(f"RULER Curves: {model_name}")
ax.set_ylim(0, 1.05)
ax.legend(loc="lower left")
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig("ruler_curves.png", dpi=150)
plt.show()

What Comes Next

This section completes the evaluation chapter's coverage of long-context assessment. For practical guidance on managing context windows in RAG systems, see Section 20.1's treatment of chunking strategies and context window management. For the attention mechanisms that underlie positional encoding, see Section 4.2: Attention Mechanisms.