[Bug Report] detect_head fails with a list of sequences and a custom tensor detection pattern

Open Beginner friendly
  • Is something wrong?
#1,979 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
76/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start with transformer_lens/head_detector.py, especially detect_head around lines 120–140, and read how list inputs and tensor patterns are handled. Check tests/integration/test_head_detector.py for the existing batched sequence cases, then add coverage for a list with a custom tensor pattern. Done means the call returns the average detection score with shape (cfg.n_layers, cfg.n_heads) without the batch dimension assertion.

Written by the indexing model from the issue text.

Description

Description of the bug

In transformer_lens.head_detector.detect_head, passing a list of sequences (seq: List[str]) together with a custom detection pattern tensor (detection_pattern: torch.Tensor) causes an AssertionError:

AssertionError: Cannot remove batch dimension from cache with batch size 2

The function supports both List[str] inputs and custom torch.Tensor detection patterns. However, the list-handling logic (if isinstance(seq, list):) is incorrectly nested inside the if isinstance(detection_pattern, str): condition.

As a result, when a tensor detection pattern is provided, the list-handling logic is skipped. The function processes all prompts as a batch and attempts to remove the batch dimension, which fails when the batch contains more than one sequence.

Code example

The following example demonstrates the issue:

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.head_detector import (
    detect_head,
    get_previous_token_head_detection_pattern,
)

# Initialize a small model on CPU
model = TransformerBridge.boot_transformers(
    "hf-internal-testing/tiny-random-gpt2",
    device="cpu",
)

# Create two prompts of equal length
prompt1 = "hello world"
prompt2 = "hello world"

tokens = model.to_tokens(prompt1)

# Generate a valid detection pattern tensor
pattern = get_previous_token_head_detection_pattern(tokens)

# Single prompt with tensor pattern works
result_single = detect_head(model, prompt1, pattern)
print("Single prompt:", result_single.shape)

# Multiple prompts with string pattern work
result_string = detect_head(
    model,
    [prompt1, prompt2],
    "previous_token_head",
)
print("Batched prompts with string pattern:", result_string.shape)

# Multiple prompts with tensor pattern trigger the bug
result_tensor = detect_head(model, [prompt1, prompt2], pattern)

Actual output

The first two calls execute successfully. The third call raises:

Traceback (most recent call last):
  File ".../transformer_lens/head_detector.py", line 158, in detect_head
    _, cache = model.run_with_cache(tokens, remove_batch_dim=True)

  File ".../transformer_lens/model_bridge/bridge_core.py", line 1779, in run_with_cache
    activation_cache.remove_batch_dim(batch_size=batch_size)

  File ".../transformer_lens/ActivationCache.py", line 185, in remove_batch_dim
    batch_size == 1

AssertionError: Cannot remove batch dimension from cache with batch size 2

Expected behaviour

detect_head(model, [prompt1, prompt2], pattern) should process each prompt individually using the supplied detection pattern and return the average detection score across all prompts.

The returned tensor should have shape (cfg.n_layers, cfg.n_heads), consistent with the behaviour when a string detection pattern is provided.

System Info

  • TransformerLens version / Git commit: 0.0.0 / 6bea1fa7cfcbf9064aa8ac99bbd5fb1e4a6c2240
  • Installation method: Source installation (uv sync)
  • Operating system: Linux (Ubuntu 24.04)
  • Python version: 3.11.18
  • PyTorch version: 2.11.0+cu130
  • Device: CPU (reproduction)

Additional context

The issue is in transformer_lens/head_detector.py, inside the detect_head function, around lines 120–140.

The list-handling block is currently nested inside the string-pattern validation block:

if isinstance(detection_pattern, str):
    assert detection_pattern in HEAD_NAMES

    if isinstance(seq, list):
        # Process each sequence individually
        ...

This means that when detection_pattern is a tensor, the entire list-handling block is skipped.

Moving the if isinstance(seq, list): block outside and before the string-pattern validation would allow both string and tensor detection patterns to work with lists of sequences.

Existing integration tests in tests/integration/test_head_detector.py cover batched sequences with string pattern names, but do not appear to cover batched sequences with custom tensor detection patterns.

Checklist

Dominant language
Python
Stars
3.9k
Forks
708
Avg merge
1d 10h
Merged PRs (30d)
82

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

  • No Dockerfile or Docker Compose file
  • Has a pull request template
  • No contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from TransformerLensOrg/TransformerLens

All issues in TransformerLensOrg/TransformerLens

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.