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

オープン 初心者向け
  • 何か問題がありますか？
#1,979 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
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 の本文から書いたものです。

説明

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日 10時間
マージ済み PR（30日）
82

環境構築

Codespaces で開く

このプロジェクトの開発コンテナを、あなたの GitHub アカウントでブラウザ上に起動します。

  • Dockerfile・Docker Compose ファイルなし
  • プルリクエストのテンプレートあり
  • コントリビューションガイドなし

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

TransformerLensOrg/TransformerLens のほかの issue

TransformerLensOrg/TransformerLens の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。