Metadata-Version: 2.4
Name: retrieval-heads
Version: 0.1.2
Summary: Retrieval Head detection in LLMs with vLLM
Author-email: Max Zuo <zuo@brown.edu>
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: matplotlib>=3.11.0
Requires-Dist: names>=0.3.0
Requires-Dist: nnsight>=0.7.0
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: rouge-score>=0.1.2
Requires-Dist: seaborn>=0.13.2
Requires-Dist: tabulate>=0.9.0
Requires-Dist: torch>=2.10.0
Requires-Dist: tqdm>=4.68.2
Requires-Dist: tyro>=1.0.13
Requires-Dist: vllm==0.19.0
Dynamic: license-file

# retrieval-heads

Retrieval head detection in LLMs using vLLM and nnsight activation tracing.

This is my attempt to faithfully reproduce [Retrieval Head Mechanistically Explains Long-Context Factuality](https://arxiv.org/abs/2404.15574), and should work out of the box with any model that uses vLLM's `Attention` or `GatedDeltaNetAttention` implementations.

Main workflows:

-  **Needle-in-a-haystack (NIAH)** – insert a known fact into a long context at varying depths and lengths, then measure retrieval accuracy (ROUGE-L).
-  **Retrieval head detection** – trace query/key activations through every attention head on NIAH results to identify which heads are responsible for retrieval.
-  **Code evaluation** – generate synthetic variable-assignment programs and ask the model for the final value of a target variable.
-  **Phonebook evaluation** – build a synthetic phonebook and ask for entries at varying depths.

## Example Results

### NIAH Heatmap

![NIAH Heatmap](imgs/heatmap.png)

### Retrieval Head Detection

![Retrieval Head Detection Heatmap](imgs/detect_heatmap.png)

## Setup

Installation:
```bash
git clone https://github.com/maxzuo/retrieval-heads.git
pip install -e .
```
Tested using Python 3.12 and vLLM 0.19.0.

## Usage

### NIAH sweep

```bash
retrieval-heads.niah --config configs/qwen3_5_9b.yaml
```

Runs the needle-in-a-haystack evaluation across a grid of context lengths and
document depths. Results are written to `output_dir` as `results.jsonl` (one
JSON record per cell) alongside the resolved `config.yaml`.

Any config field can be overridden via CLI flags:

```bash
retrieval-heads.niah --config configs/qwen3_5_9b.yaml \
    --model.max-model-len 16384 \
    --output-dir ./results/short
```

### Retrieval head detection

```bash
retrieval-heads.detect --config configs/detect.yaml
```

Takes NIAH result files as input, traces each forward pass with nnsight to
capture per-head query/key matrices, and scores each head on whether it attends
to the needle span. Outputs `detected.json` and `detected-agg.json`.

### Code evaluation

```bash
retrieval-heads.code-eval --config configs/code_eval_olmo3.yaml
```

Runs the code-evaluation experiment: generates synthetic variable-assignment
programs across a grid of variable counts and step counts, asks the model for
the final value of a target variable, and scores each answer as correct or not.
Results are written to `output_dir` as `results.jsonl` (one record per task)
alongside the resolved `config.yaml`. Pass `--report` (or `report: true` in the
YAML) to also build an accuracy report, printed to stdout and saved as
`report.txt`.

Any config field can be overridden via CLI flags:

```bash
retrieval-heads.code-eval --config configs/code_eval_olmo3.yaml \
    --code-eval.num-tasks 50 \
    --output-dir ./results/code_eval_short
```

### Report

```bash
retrieval-heads.report results/code_eval_olmo3/results.jsonl
```

Builds a report from any saved `results.jsonl` file. The experiment type is
auto-detected from each record's `experiment_type` field and dispatched to the
matching report (code-eval reports break down accuracy by task type, variable
count, and step count; NIAH and phonebook reports are not yet implemented).
Prints to stdout unless `--output-path` is given, in which case the report is
written there instead.

### Phonebook evaluation

```bash
retrieval-heads.phonebook --config configs/phonebook.yaml
```

Builds a synthetic phonebook for each configured length, asks for entries at
the selected depths, and scores answers by exact digit-normalized phone-number
match. The resolved configuration and one JSONL record per evaluated entry are
written to `output_dir` as `config.yaml` and `results.jsonl`.

Phonebook `depths` are percentages from 0 through 100. Each percentage is
mapped to the nearest zero-based entry position (with half-up rounding), and
positions that collide are evaluated once. If `depths` is omitted, every entry
at every configured length is evaluated. That exhaustive default can produce a
large batch because all prompts are submitted in one generation call.

### Visualization

```bash
retrieval-heads.visualize niah --results results/qwen3_5_9b/results.jsonl
retrieval-heads.visualize detect --results results/detect/detected-agg.json
```

`code-eval` renders one figure with three panels (Recall, State Tracking, State
Based Recall) from a single `results.jsonl`, written to `--output`:

```bash
retrieval-heads.visualize code-eval \
    --results results/code_eval_olmo3/results.jsonl \
    --output results/code_eval_olmo3/code_eval.png --title olmo3
```

`lines` overlays multiple models on one figure as line plots — for either
code-eval or phonebook results (all files must be the same experiment kind).
`--labels` defaults to each file's parent-directory name and `--colors` to the
matplotlib color cycle:

```bash
retrieval-heads.visualize lines \
    --results results/code_eval_olmo3/results.jsonl \
             results/code_eval_olmo_hybrid/results.jsonl \
    --labels olmo3 olmo-hybrid --colors '#1f77b4' '#ff7f0e' \
    --output results/code_eval_lines.png

retrieval-heads.visualize lines \
    --results results/phonebook_olmo3/results.jsonl \
             results/phonebook_olmo_hybrid/results.jsonl \
    --output results/phonebook_lines.png
```

## Configuration

Configs are YAML files with the following sections:

```yaml
model:
  model: Qwen/Qwen3.5-9B
  max_model_len: 32768
  dtype: bfloat16
  chat_template: path/to/template.jinja
  language_model_only: true

haystack:
  haystack_dir: ./PaulGrahamEssays
  needle: "\nThe best thing to do in San Francisco is eat a sandwich...\n"
  retrieval_question: "What is the best thing to do in San Francisco?"

sweep:
  context_lengths: {min: 1000, max: 32000, intervals: 31}
  document_depths: {min: 0, max: 100, intervals: 10}

output_dir: ./results/qwen3_5_9b
```

Sweep dimensions accept either a `{min, max, intervals}` shorthand or an
explicit list of values.

Code-eval configs use `model`, a `code_eval` section, and `output_dir`:

```yaml
model:
  model: allenai/Olmo-3-7B-Instruct
  max_model_len: 32768
  dtype: bfloat16
  language_model_only: true

code_eval:
  num_vars: [4, 8, 16, 32, 64, 128]
  num_steps: [4, 8, 16, 32, 64, 128]
  num_tasks: 100
  task_types: [RECALL, STATE_TRACKING, STATE_BASED_RECALL]
  seed: 42

output_dir: ./results/code_eval_olmo3
```

`num_vars` and `num_steps` accept either the `{min, max, intervals}` shorthand
or an explicit list of values.

Phonebook configs use `model`, a `phonebook` section, and `output_dir`; see
`configs/phonebook.yaml` for a runnable example. `lengths` accepts an explicit
list or range shorthand, while `depths` is an optional explicit percentage
list.
