Long-context evaluation - NIAH and RULER

Open in Colab ▶️ Run this notebook in Colab

Long-context evaluation - NIAH and RULER#

Track 06 - Evaluation · Notebook 06 · Runtime: ≈1 min on CPU

Prerequisites: 01_inference/01 (KV cache), 06_eval/01 (perplexity).

Papers: gkamradt’s Needle in a Haystack repo; Hsieh et al. 2024, RULER: What’s the Real Context Size of Your Long-Context Language Models? (2404.06654).

What#

Long-context accuracy doesn’t scale smoothly with context window. Models that claim 128 K often degrade sharply past 16 K. Two standard benchmarks probe this:

  • Needle-in-a-Haystack (NIAH). Insert a short distinctive fact (“The magic number is 7891”) at depth d in a long haystack; ask the model to recall the fact. Sweep depth and context length; plot a heatmap.

  • RULER. Extends NIAH with multi-hop aggregation, variable tracking, word frequency, and question-answering tasks. The goal is to stress every axis of long-context competence, not just needle recall.

We can’t run a real long-context model here, so we model the effect via a simple decay function: recall probability drops with distance from the end (models often forget the middle). We verify the harness reproduces known failure patterns.

from llm_systems_cookbook.nb import bootstrap

import numpy as np

s = bootstrap("06_eval_06_long_context_niah_ruler")

NIAH setup#

Construct haystacks at context_len ∈ {1k, 4k, 16k, 64k}. Insert a needle at depth d ∈ {0.0, 0.25, 0.5, 0.75, 1.0} (fraction of context from the start). Define a recall probability model that captures two empirically observed patterns:

  • Lost-in-the-middle: models recall needles near the start and end of context better than the middle.

  • Context-length decay: recall drops as context grows, faster for middle depths than edges.

def recall_prob(context_len: int, depth: float, max_supported_len: int = 4096) -> float:
    '''Simulated needle-retrieval probability.

    Models:
    - Full recall when context_len <= max_supported_len AND depth in {0, 1} (anchor positions).
    - Gradual decay with context length for middle depths.
    - Stronger at the end of the context than in the middle.
    '''
    # Distance from nearest anchor (start or end).
    middle_distance = min(depth, 1.0 - depth)  # 0 at edge, 0.5 at middle
    length_penalty = np.clip(np.log2(context_len / max_supported_len + 1), 0.0, None)
    # Recall falls off exponentially with middle_distance * length_penalty.
    return float(np.clip(np.exp(-middle_distance * length_penalty * 4.0), 0.05, 1.0))


CONTEXT_LENS = [1024, 4096, 16384, 65536]
DEPTHS = [0.0, 0.25, 0.50, 0.75, 1.0]

heatmap = np.array([[recall_prob(L, d) for d in DEPTHS] for L in CONTEXT_LENS])
for L, row in zip(CONTEXT_LENS, heatmap, strict=True):
    row_str = "  ".join(f"{v:.2f}" for v in row)
    print(f"L={L:>6}  depths={DEPTHS}  recall={row_str}")

NIAH checks#

  • Recall is always ≥ some floor (model doesn’t completely fail).

  • Edge depths (0.0 and 1.0) outperform middle depths (0.5) at every non-trivial context length.

  • Average recall at 64K is strictly worse than at 1K.

avg_recall_per_L = heatmap.mean(axis=1)
for L, avg in zip(CONTEXT_LENS, avg_recall_per_L, strict=True):
    print(f"L={L:>6}  avg recall = {avg:.3f}")

s.check(
    "recall_always_above_floor",
    lambda: float(heatmap.min()) > 0.0,
    msg=f"min recall = {heatmap.min():.3f}",
)
s.check(
    "recall_at_16k_edges_above_middle",
    lambda: heatmap[2, 0] > heatmap[2, 2] and heatmap[2, 4] > heatmap[2, 2],
    msg=f"L=16k row = {heatmap[2]}",
)
s.check(
    "avg_recall_monotone_decreasing_with_context",
    lambda: all(avg_recall_per_L[i] >= avg_recall_per_L[i + 1] for i in range(len(avg_recall_per_L) - 1)),
    msg=f"avg per L = {avg_recall_per_L}",
)

RULER-style composite score#

RULER reports eight task categories averaged at each context length. We simulate three: single-NIAH (as above), multi-key NIAH (recall K of N needles), and variable tracking (follow a variable through reassignments). Harder tasks degrade faster with context.

def multi_key_niah_recall(context_len: int, n_keys: int) -> float:
    '''Average recall over n_keys independently inserted needles.'''
    return float(np.mean([recall_prob(context_len, (i + 1) / (n_keys + 1)) for i in range(n_keys)]))


def variable_tracking_accuracy(context_len: int, chain_length: int) -> float:
    '''Probability of correctly tracking a variable through ``chain_length``
    reassignments: each reassignment step must be retrieved correctly, and
    the chain succeeds only if all of them do.
    '''
    per_step = recall_prob(context_len, 0.5)
    return float(per_step ** chain_length)


ruler = {}
for L in CONTEXT_LENS:
    niah_1 = recall_prob(L, 0.5)
    niah_3 = multi_key_niah_recall(L, n_keys=3)
    var_track = variable_tracking_accuracy(L, chain_length=4)
    ruler[L] = {"niah_1": niah_1, "niah_3": niah_3, "var_track": var_track,
                "avg": (niah_1 + niah_3 + var_track) / 3}
    print(f"L={L:>6}  niah_1={niah_1:.3f}  niah_3={niah_3:.3f}  var_track={var_track:.3f}  avg={ruler[L]['avg']:.3f}")

s.check(
    "ruler_composite_drops_with_context",
    lambda: ruler[CONTEXT_LENS[-1]]["avg"] < ruler[CONTEXT_LENS[0]]["avg"],
    msg=f"L={CONTEXT_LENS[0]} avg = {ruler[CONTEXT_LENS[0]]['avg']:.3f}  "
        f"L={CONTEXT_LENS[-1]} avg = {ruler[CONTEXT_LENS[-1]]['avg']:.3f}",
)
s.check(
    "harder_task_degrades_faster_than_niah",
    lambda: ruler[1024]["var_track"] < ruler[1024]["niah_1"],
    msg=f"at 1k: niah_1={ruler[1024]['niah_1']:.3f}  var_track={ruler[1024]['var_track']:.3f}",
)

The canonical NIAH heatmap#

The Needle-in-a-Haystack chart everyone has seen: needle depth on the x-axis, context length on the y-axis, cell colour = recall probability. The characteristic “lost-in-the-middle” bowtie appears as a darker central column at longer contexts, with the two edge columns staying bright.

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(6.8, 4.0))
im = ax.imshow(heatmap, aspect="auto", cmap="RdYlGn", vmin=0.0, vmax=1.0, origin="lower")
ax.set_xticks(range(len(DEPTHS)))
ax.set_xticklabels([f"{d:.2f}" for d in DEPTHS])
ax.set_yticks(range(len(CONTEXT_LENS)))
ax.set_yticklabels([f"{L//1024}k" if L >= 1024 else str(L) for L in CONTEXT_LENS])
ax.set_xlabel("needle depth (fraction of context)")
ax.set_ylabel("context length (tokens)")
ax.set_title("NIAH recall heatmap  (green = retrieved, red = missed)")
for i in range(heatmap.shape[0]):
    for j in range(heatmap.shape[1]):
        ax.text(j, i, f"{heatmap[i, j]:.2f}", ha="center", va="center",
                color="black" if heatmap[i, j] > 0.5 else "white", fontsize=9)
fig.colorbar(im, ax=ax, label="recall probability")
fig.tight_layout()
plt.show()

Exercises#

  1. Hook this up to a real long-context model (Llama-3.2-1B claims 128K but performance above 8K is a fair test). Replace recall_prob with actual inference.

  2. Vary the needle pattern: exact-string match vs paraphrase-recall vs numerical arithmetic on the needle. Harder patterns degrade at shorter contexts.

  3. Plot the NIAH heatmap with matplotlib’s imshow; the diagonal stripe pattern is the standard “lost in the middle” visual.

References#

  • gkamradt’s NIAH script - the tooling everyone uses for the raw needle task.

  • RULER paper §3 for the eight-task breakdown and the “effective context window” definition.

s.summary()
s.save()