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

Aperta Adatta ai principianti
  • Qualcosa non va?
#1,979 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
76/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
machine-learning

Direzione di ricerca

Inizia da transformer_lens/head_detector.py, in particolare da detect_head intorno alle righe 120–140, e leggi come vengono gestiti gli input di tipo lista e i pattern dei tensori. Controlla tests/integration/test_head_detector.py per i casi esistenti di sequenze in batch, poi aggiungi una verifica per una lista con un pattern di tensore personalizzato. Il risultato atteso è che la chiamata restituisca il punteggio medio di rilevamento con forma (cfg.n_layers, cfg.n_heads) senza l’asserzione sulla dimensione del batch.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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

Lingua principale
Python
Stelle
3.9k
Fork
708
Merge medio
1g 10h
PR unite (30g)
82

Preparare l'ambiente

Apri in Codespaces

Avvia il container di sviluppo del progetto nel browser, con il tuo account GitHub.

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di TransformerLensOrg/TransformerLens

Tutte le issue di TransformerLensOrg/TransformerLens

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.