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

未关闭 适合新手
  • 哪里不对？
#1,979 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

已经有一个关联 PR 被合并了。

  • #1982 来自 @koriyoshi2041 —— 已合并

评估

难度
2/5
预计耗时
1-3 小时
新手友好度
76/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python, pytorch
领域
machine-learning

调研方向

从 transformer_lens/head_detector.py 开始，重点查看第 120–140 行附近的 detect_head，并了解列表输入和张量模式的处理方式。检查 tests/integration/test_head_detector.py 中现有的批量序列用例，然后添加一个带有自定义张量模式的列表用例。完成标准是：调用返回形状为 (cfg.n_layers, cfg.n_heads) 的平均检测分数，且不触发批次维度断言。

由索引模型根据 Issue 内容生成。

描述

bug complexity-moderate TransformerBridge

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

主要语言
Python
星标
3.9k
派生
708
平均合并
1 天 4 小时
30 天内合并 PR
87

环境准备

在 Codespaces 中打开

在浏览器里用你自己的 GitHub 账号启动这个项目的开发容器。

  • 没有 Dockerfile 或 Docker Compose 文件
  • 有 Pull Request 模板
  • 没有贡献指南

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

TransformerLensOrg/TransformerLens 的其他 Issue

查看 TransformerLensOrg/TransformerLens 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。