retrieval-heads 0.1.0__tar.gz

retrieval_heads-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Max Zuo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

retrieval_heads-0.1.0/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: retrieval-heads
+Version: 0.1.0
+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: 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: 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.
+
+Two main workflows:
+
+1. **Needle-in-a-haystack (NIAH)** – insert a known fact into a long context at varying depths and lengths, then measure retrieval accuracy (ROUGE-L).
+2. **Retrieval head detection** – trace query/key activations through every attention head on NIAH results to identify which heads are responsible for retrieval.
+
+## 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`.
+
+### Visualization
+
+```bash
+retrieval-heads.visualize niah --results results/qwen3_5_9b/results.jsonl
+retrieval-heads.visualize detect --results results/detect/detected-agg.json
+```
+
+## 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.

retrieval_heads-0.1.0/README.md

@@ -0,0 +1,93 @@
+# 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.
+
+Two main workflows:
+
+1. **Needle-in-a-haystack (NIAH)** – insert a known fact into a long context at varying depths and lengths, then measure retrieval accuracy (ROUGE-L).
+2. **Retrieval head detection** – trace query/key activations through every attention head on NIAH results to identify which heads are responsible for retrieval.
+
+## 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`.
+
+### Visualization
+
+```bash
+retrieval-heads.visualize niah --results results/qwen3_5_9b/results.jsonl
+retrieval-heads.visualize detect --results results/detect/detected-agg.json
+```
+
+## 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.

retrieval_heads-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=75"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "retrieval-heads"
+version = "0.1.0"
+description = "Retrieval Head detection in LLMs with vLLM"
+authors = [{name = "Max Zuo", email = "zuo@brown.edu"}]
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "matplotlib>=3.11.0",
+    "nnsight>=0.7.0",
+    "pyyaml>=6.0.3",
+    "rouge-score>=0.1.2",
+    "seaborn>=0.13.2",
+    "torch>=2.10.0",
+    "tqdm>=4.68.2",
+    "tyro>=1.0.13",
+    "vllm==0.19.0",
+]
+
+[project.scripts]
+"retrieval-heads.niah" = "retrieval_heads.cli:niah_cli"
+"retrieval-heads.detect" = "retrieval_heads.cli:detect_cli"
+"retrieval-heads.visualize" = "retrieval_heads.visualize:cli"
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+]
+
+[tool.setuptools.packages.find]
+include = ["retrieval_heads*"]
+
+[tool.pytest.ini_options]
+pythonpath = ["."]
+testpaths = ["tests"]

retrieval_heads-0.1.0/retrieval_heads/__init__.py

@@ -0,0 +1,6 @@
+from .configs import (
+    HaystackConfig,
+    ModelConfig,
+    RangeConfig,
+    SweepConfig,
+)

retrieval_heads-0.1.0/retrieval_heads/cli.py

@@ -0,0 +1,254 @@
+import argparse
+from dataclasses import asdict, dataclass, field
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import tqdm
+import tyro
+import yaml
+
+from .configs import (
+    HaystackConfig,
+    ModelConfig,
+    SweepConfig,
+)
+
+
+# Experiment-level configs
+@dataclass
+class NIAHConfig:
+  """Top-level config for a needle-in-a-haystack run."""
+  model: ModelConfig
+  haystack: HaystackConfig
+  sweep: SweepConfig = field(default_factory=SweepConfig)
+  output_dir: str | None = None
+
+
+@dataclass
+class DetectConfig:
+  """Top-level config for retrieval-head detection."""
+  model: ModelConfig
+  results_files: tuple[str, ...] = ()
+  output_dir: str | None = None
+  sweep: SweepConfig | None = None
+  score_threshold: float | None = None
+
+
+# Config loading functions
+
+
+def load_niah_config(path: str | os.PathLike) -> NIAHConfig:
+  """Load a NIAHConfig from a YAML file."""
+  with open(path) as f:
+    raw = yaml.safe_load(f) or {}
+
+  return NIAHConfig(
+      model=ModelConfig(**raw['model']),
+      haystack=HaystackConfig(**raw['haystack']),
+      sweep=SweepConfig(**raw.get('sweep', {})),
+      output_dir=raw.get('output_dir'),
+  )
+
+
+def load_detect_config(path: str | os.PathLike) -> DetectConfig:
+  """Load a DetectConfig from a YAML file."""
+  with open(path) as f:
+    raw = yaml.safe_load(f) or {}
+
+  results_files = raw.get('results_files', ())
+  if isinstance(results_files, str):
+    results_files = (results_files,)
+  else:
+    results_files = tuple(results_files)
+
+  return DetectConfig(
+      model=ModelConfig(**raw['model']),
+      results_files=results_files,
+      output_dir=raw.get('output_dir'),
+      sweep=SweepConfig(**raw['sweep']) if 'sweep' in raw else None,
+      score_threshold=raw.get('score_threshold'),
+  )
+
+
+# Load results.jsonl file
+
+
+def load_results(path: str | os.PathLike) -> list[dict[str, Any]]:
+  """Load and validate NIAH result records from a JSONL file."""
+  required_fields = {
+      'context_length',
+      'document_depth',
+      'needle',
+      'prompt',
+      'token_ids',
+  }
+  results = []
+
+  with open(path) as f:
+    for line_number, line in enumerate(f, start=1):
+      if not line.strip():
+        continue
+      try:
+        result = json.loads(line)
+      except json.JSONDecodeError as error:
+        raise ValueError(
+            f'Invalid JSON in {path} at line {line_number}: {error.msg}'
+        ) from error
+
+      if not isinstance(result, dict):
+        raise ValueError(
+            f'Expected a JSON object in {path} at line {line_number}.')
+
+      missing = sorted(required_fields - result.keys())
+      if missing:
+        raise ValueError(
+            f'Missing required fields in {path} at line {line_number}: '
+            f'{", ".join(missing)}')
+      results.append(result)
+
+  if not results:
+    raise ValueError(f'No result records found in {path}.')
+
+  return results
+
+
+def filter_results(
+    results: list[dict[str, Any]],
+    sweep: SweepConfig | None,
+) -> list[dict[str, Any]]:
+  """Filter results to exact context-length/document-depth sweep cells."""
+  if sweep is None:
+    return results
+
+  context_lengths = set(sweep.context_lengths)
+  document_depths = set(sweep.document_depths)
+  return [
+      result for result in results
+      if result['context_length'] in context_lengths and
+      result['document_depth'] in document_depths
+  ]
+
+
+# Main experiment functions
+
+
+def niah_main(config: NIAHConfig):
+  """Run a NIAH sweep and save results + resolved config."""
+  from .niah import NeedleInAHaystack
+
+  niah = NeedleInAHaystack(
+      haystack_config=config.haystack,
+      sweep_config=config.sweep,
+      model_config=config.model,
+  )
+  niah.run(output_path=os.path.join(config.output_dir, 'results.jsonl'))
+
+  with open(os.path.join(config.output_dir, 'config.yaml'), 'w') as f:
+    yaml.dump(asdict(config), f)
+
+
+def detect_main(config: DetectConfig, detector_cls=None):
+  """Trace selected NIAH results and save aggregate retrieval-head scores."""
+  if not config.results_files:
+    raise ValueError(
+        'results_files is required: set it in the --config YAML or pass '
+        '--results-files.')
+  if config.output_dir is None:
+    raise ValueError(
+        'output_dir is required: set it in the --config YAML or pass '
+        '--output-dir.')
+
+  all_results = []
+  for path in config.results_files:
+    all_results.extend(load_results(path))
+  results = filter_results(all_results, config.sweep)
+  if not results:
+    raise ValueError('No results matched the configured sweep.')
+
+  if config.score_threshold is not None:
+    for result in results:
+      if 'rougeL' not in result:
+        raise ValueError(
+            'score_threshold is set but result record is missing "rougeL".')
+    results = [
+        result for result in results
+        if result['rougeL'] >= config.score_threshold
+    ]
+    if not results:
+      raise ValueError('No results met the score threshold.')
+
+  if detector_cls is None:
+    from .detect import RetrievalHeadDetector
+    detector_cls = RetrievalHeadDetector
+
+  detector = detector_cls(model_config=config.model)
+  for result in tqdm.tqdm(results, desc='Tracing results'):
+    detector.calculate(
+        prompt=result['prompt'],
+        needle=result['needle'],
+        completion_tokens=result['token_ids'],
+    )
+
+  detector.save(config.output_dir)
+
+
+# cli entry points
+
+def niah_cli():
+  """Entry point for ``retrieval-heads.niah``."""
+  pre = argparse.ArgumentParser(add_help=False)
+  pre.add_argument(
+      '--config',
+      type=str,
+      default=None,
+      help='Path to a YAML config with top-level "model", "haystack", "sweep" '
+      'and optional "output_dir" sections.',
+  )
+  known, remaining = pre.parse_known_args()
+
+  default = load_niah_config(
+      known.config) if known.config else tyro.MISSING_NONPROP
+  config = tyro.cli(
+      NIAHConfig,
+      description='Run Needle-in-a-Haystack (NIAH) file.\n\n'
+      'Pass --config with a YAML file to load config from file.',
+      args=remaining,
+      default=default,
+  )
+
+  if config.output_dir is None:
+    raise SystemExit(
+        'output_dir is required: set it in the --config YAML or pass --output-dir.'
+    )
+
+  niah_main(config)
+
+
+def detect_cli():
+  """Entry point for ``retrieval-heads.detect``."""
+  pre = argparse.ArgumentParser(add_help=False)
+  pre.add_argument(
+      '--config',
+      type=str,
+      default=None,
+      help='Path to a YAML config with top-level "model", "haystack", '
+      'optional "sweep", "results_file", and "output_path" sections.',
+  )
+  known, remaining = pre.parse_known_args()
+
+  default = load_detect_config(
+      known.config) if known.config else tyro.MISSING_NONPROP
+  config = tyro.cli(
+      DetectConfig,
+      description='Run retrieval-head detection over saved NIAH JSONL results.'
+      '\n\nPass --config with a YAML file to load config defaults.',
+      args=remaining,
+      default=default,
+  )
+
+  try:
+    detect_main(config)
+  except ValueError as e:
+    raise SystemExit(str(e)) from e

retrieval_heads-0.1.0/retrieval_heads/configs.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ModelConfig:
+  model: str
+  max_model_len: int | None = None
+  trust_remote_code: bool = False
+  dtype: str = 'auto'
+  tensor_parallel_size: int = 1
+  gpu_memory_utilization: float = 0.9
+  enable_prefix_caching: bool = True
+  chat_template: str | None = None
+  enforce_eager: bool = False
+  language_model_only: bool = False
+
+
+@dataclass
+class RangeConfig:
+  min: int
+  max: int
+  intervals: int
+
+
+@dataclass
+class SweepConfig:
+  context_lengths: list[int] = field(
+      default_factory=lambda: list(range(1_000, 50_000, 1_000)))
+  document_depths: list[int] = field(
+      default_factory=lambda: list(range(0, 100, 10)))
+
+  def __post_init__(self):
+
+    def expand_range(r: list[float] | list[int] | RangeConfig) -> list[int]:
+      if isinstance(r, RangeConfig):
+        return list(range(r.min, r.max + 1, (r.max - r.min) // r.intervals))
+      elif isinstance(r, dict):
+        return list(
+            range(
+                r['min'],
+                r['max'] + 1,
+                (r['max'] - r['min']) // r['intervals'],
+            ))
+      else:
+        return sorted(list(r))
+
+    self.context_lengths = expand_range(self.context_lengths)
+    self.document_depths = expand_range(self.document_depths)
+
+
+@dataclass
+class HaystackConfig:
+  haystack_dir: str
+  needle: str = '\nThe best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.\n'
+  retrieval_question: str = 'What is the best thing to do in San Francisco?'