Coverage for transformer_lens/head_detector.py: 96%
89 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Head Detector.
3Utilities for detecting specific types of heads (e.g. previous token heads).
4"""
6import logging
7from collections import defaultdict
8from typing import Dict, List, Optional, Tuple, Union, cast
10import numpy as np
11import torch
12from typing_extensions import Literal, get_args
14from transformer_lens.ActivationCache import ActivationCache
15from transformer_lens.model_protocol import TransformerLensModel
16from transformer_lens.utilities import is_lower_triangular, is_square
18HeadName = Literal["previous_token_head", "duplicate_token_head", "induction_head"]
19HEAD_NAMES = cast(List[HeadName], get_args(HeadName))
20ErrorMeasure = Literal["abs", "mul"]
22LayerHeadTuple = Tuple[int, int]
23LayerToHead = Dict[int, List[int]]
25INVALID_HEAD_NAME_ERR = (
26 f"detection_pattern must be a Tensor or one of head names: {HEAD_NAMES}; got %s"
27)
29CACHE_WITH_SEQ_LIST_ERR = (
30 "A single cache cannot be reused across multiple prompts, so `cache` is not\n"
31 "supported when `seq` is a list. Pass one prompt at a time, or omit `cache`\n"
32 "and let each prompt be run separately."
33)
35SEQ_LEN_ERR = "The sequence must be non-empty and must fit within the model's context window."
37DET_PAT_NOT_SQUARE_ERR = "The detection pattern must be a lower triangular matrix of shape (sequence_length, sequence_length); sequence_length=%d; got detection pattern of shape %s"
40def detect_head(
41 model: TransformerLensModel,
42 seq: Union[str, List[str]],
43 detection_pattern: Union[torch.Tensor, HeadName],
44 heads: Optional[Union[List[LayerHeadTuple], LayerToHead]] = None,
45 cache: Optional[ActivationCache] = None,
46 *,
47 exclude_bos: bool = False,
48 exclude_current_token: bool = False,
49 error_measure: ErrorMeasure = "mul",
50) -> torch.Tensor:
51 """Search for a Particular Type of Attention Head.
53 Searches the model (or a set of specific heads, for circuit analysis) for a particular type of
54 attention head. This head is specified by a detection pattern, a (sequence_length,
55 sequence_length) tensor representing the attention pattern we expect that type of attention head
56 to show. The detection pattern can be also passed not as a tensor, but as a name of one of
57 pre-specified types of attention head (see `HeadName` for available patterns), in which case the
58 tensor is computed within the function itself.
60 There are two error measures available for quantifying the match between the detection pattern
61 and the actual attention pattern.
63 1. `"mul"` (default) multiplies both tensors element-wise and divides the sum of the result by
64 the sum of the attention pattern. Typically, the detection pattern should in this case
65 contain only ones and zeros, which allows a straightforward interpretation of the score: how
66 big fraction of this head's attention is allocated to these specific query-key pairs? Using
67 values other than 0 or 1 is not prohibited but will raise a warning (which can be disabled,
68 of course).
70 2. `"abs"` calculates the mean element-wise absolute difference between the detection pattern
71 and the actual attention pattern. The "raw result" ranges from 0 to 2 where lower score
72 corresponds to greater accuracy. Subtracting it from 1 maps that range to (-1, 1) interval,
73 with 1 being perfect match and -1 perfect mismatch.
75 Which one should you use?
77 `"mul"` is likely better for quick or exploratory investigations. For precise examinations where
78 you're trying to reproduce as much functionality as possible or really test your understanding
79 of the attention head, you probably want to switch to `"abs"`.
81 The advantage of `"abs"` is that you can make more precise predictions, and have that measured
82 in the score. You can predict, for instance, 0.2 attention to X, and 0.8 attention to Y, and
83 your score will be better if your prediction is closer. The "mul" metric does not allow this,
84 you'll get the same score if attention is 0.2, 0.8 or 0.5, 0.5 or 0.8, 0.2.
86 Args:
87 model: Model being used.
88 seq: String or list of strings being fed to the model.
89 head_name: Name of an existing head in HEAD_NAMES we want to check. Must pass either a
90 head_name or a detection_pattern, but not both!
91 detection_pattern: (sequence_length, sequence_length)nTensor representing what attention
92 pattern corresponds to the head we're looking for or the name of a pre-specified head.
93 Currently available heads are: `["previous_token_head", "duplicate_token_head",
94 "induction_head"]`.
95 heads: If specific attention heads is given here, all other heads' score is set to -1.
96 Useful for IOI-style circuit analysis. Heads can be specified as a list of tuples (layer,
97 head) or a dictionary mapping a layer to heads within that layer that we want to
98 analyze. cache: Include the cache to save time if you want.
99 exclude_bos: Exclude attention paid to the beginning of sequence token.
100 exclude_current_token: Exclude attention paid to the current token.
101 error_measure: `"mul"` for using element-wise multiplication. `"abs"` for using absolute
102 values of element-wise differences as the error measure.
104 Returns:
105 Tensor representing the score for each attention head.
106 """
108 cfg = model.cfg
109 tokens = model.to_tokens(seq).to(cfg.device)
110 seq_len = tokens.shape[-1]
112 # Validate error_measure
114 assert error_measure in get_args(
115 ErrorMeasure
116 ), f"Invalid error_measure={error_measure}; valid values are {get_args(ErrorMeasure)}"
118 # Validate detection pattern if it's a string
119 if isinstance(detection_pattern, str):
120 assert detection_pattern in HEAD_NAMES, INVALID_HEAD_NAME_ERR % detection_pattern
121 if isinstance(seq, list):
122 # Every other argument is forwarded below. `cache` deliberately is
123 # not, because one cache holds the activations of one prompt and
124 # cannot serve the rest. Say so rather than dropping it quietly.
125 if cache is not None:
126 raise ValueError(CACHE_WITH_SEQ_LIST_ERR)
127 batch_scores = [
128 detect_head(
129 model,
130 batch_seq,
131 detection_pattern,
132 heads=heads,
133 exclude_bos=exclude_bos,
134 exclude_current_token=exclude_current_token,
135 error_measure=error_measure,
136 )
137 for batch_seq in seq
138 ]
139 return torch.stack(batch_scores).mean(0)
140 detection_pattern = cast(
141 torch.Tensor,
142 eval(f"get_{detection_pattern}_detection_pattern(tokens.cpu())"),
143 ).to(cfg.device)
145 # if we're using "mul", detection_pattern should consist of zeros and ones
146 if error_measure == "mul" and not set(detection_pattern.unique().tolist()).issubset({0, 1}): 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true
147 logging.warning(
148 "Using detection pattern with values other than 0 or 1 with error_measure 'mul'"
149 )
151 # Validate inputs and detection pattern shape
152 assert 1 < tokens.shape[-1] < cfg.n_ctx, SEQ_LEN_ERR
153 assert (
154 is_lower_triangular(detection_pattern) and seq_len == detection_pattern.shape[0]
155 ), DET_PAT_NOT_SQUARE_ERR % (seq_len, detection_pattern.shape)
157 if cache is None:
158 _, cache = model.run_with_cache(tokens, remove_batch_dim=True)
160 if heads is None:
161 layer2heads = {layer_i: list(range(cfg.n_heads)) for layer_i in range(cfg.n_layers)}
162 elif isinstance(heads, list): 162 ↛ 167line 162 didn't jump to line 167 because the condition on line 162 was always true
163 layer2heads = defaultdict(list)
164 for layer, head in heads:
165 layer2heads[layer].append(head)
166 else:
167 layer2heads = heads
169 matches = -torch.ones(cfg.n_layers, cfg.n_heads, dtype=cfg.dtype)
171 for layer, layer_heads in layer2heads.items():
172 # [n_heads q_pos k_pos]
173 layer_attention_patterns = cache["pattern", layer, "attn"]
174 for head in layer_heads:
175 head_attention_pattern = layer_attention_patterns[head, :, :]
176 head_score = compute_head_attention_similarity_score(
177 head_attention_pattern,
178 detection_pattern=detection_pattern,
179 exclude_bos=exclude_bos,
180 exclude_current_token=exclude_current_token,
181 error_measure=error_measure,
182 )
183 matches[layer, head] = head_score
184 return matches
187# Previous token head
188def get_previous_token_head_detection_pattern(
189 tokens: torch.Tensor, # [batch (1) x pos]
190) -> torch.Tensor:
191 """Outputs a detection score for [previous token heads](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=0O5VOHe9xeZn8Ertywkh7ioc).
193 Args:
194 tokens: Tokens being fed to the model.
195 """
196 detection_pattern = torch.zeros(tokens.shape[-1], tokens.shape[-1])
197 # Adds a diagonal of 1's below the main diagonal.
198 detection_pattern[1:, :-1] = torch.eye(tokens.shape[-1] - 1)
199 return torch.tril(detection_pattern)
202# Duplicate token head
203def get_duplicate_token_head_detection_pattern(
204 tokens: torch.Tensor, # [batch (1) x pos]
205) -> torch.Tensor:
206 """Outputs a detection score for [duplicate token heads](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=2UkvedzOnghL5UHUgVhROxeo).
208 Args:
209 sequence: String being fed to the model.
210 """
211 # [pos x pos]
212 token_pattern = tokens.repeat(tokens.shape[-1], 1).numpy()
214 # If token_pattern[i][j] matches its transpose, then token j and token i are duplicates.
215 eq_mask = np.equal(token_pattern, token_pattern.T).astype(int)
217 np.fill_diagonal(eq_mask, 0) # Current token is always a duplicate of itself. Ignore that.
218 detection_pattern = eq_mask.astype(int)
219 return torch.tril(torch.as_tensor(detection_pattern).float())
222# Induction head
223def get_induction_head_detection_pattern(
224 tokens: torch.Tensor, # [batch (1) x pos]
225) -> torch.Tensor:
226 """Outputs a detection score for [induction heads](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=_tFVuP5csv5ORIthmqwj0gSY).
228 Args:
229 sequence: String being fed to the model.
230 """
231 duplicate_pattern = get_duplicate_token_head_detection_pattern(tokens)
233 # Shift all items one to the right
234 shifted_tensor = torch.roll(duplicate_pattern, shifts=1, dims=1)
236 # Replace first column with 0's
237 # we don't care about bos but shifting to the right moves the last column to the first,
238 # and the last column might contain non-zero values.
239 zeros_column = torch.zeros(duplicate_pattern.shape[0], 1)
240 result_tensor = torch.cat((zeros_column, shifted_tensor[:, 1:]), dim=1)
241 return torch.tril(result_tensor)
244def get_supported_heads() -> None:
245 """Print the supported head names."""
246 print(f"Supported heads: {HEAD_NAMES}")
249def compute_head_attention_similarity_score(
250 attention_pattern: torch.Tensor, # [q_pos k_pos]
251 detection_pattern: torch.Tensor, # [seq_len seq_len] (seq_len == q_pos == k_pos)
252 *,
253 exclude_bos: bool,
254 exclude_current_token: bool,
255 error_measure: ErrorMeasure,
256) -> float:
257 """Compute the similarity between `attention_pattern` and `detection_pattern`.
259 Args:
260 attention_pattern: Lower triangular matrix (Tensor) representing the attention pattern of a particular attention head.
261 detection_pattern: Lower triangular matrix (Tensor) representing the attention pattern we are looking for.
262 exclude_bos: `True` if the beginning-of-sentence (BOS) token should be omitted from comparison. `False` otherwise.
263 exclude_bcurrent_token: `True` if the current token at each position should be omitted from comparison. `False` otherwise.
264 error_measure: "abs" for using absolute values of element-wise differences as the error measure. "mul" for using element-wise multiplication (legacy code).
265 """
266 assert is_square(
267 attention_pattern
268 ), f"Attention pattern is not square; got shape {attention_pattern.shape}"
270 # mul
272 if error_measure == "mul":
273 # Clone before masking. attention_pattern is a view into the caller's
274 # ActivationCache, so masking in place permanently corrupts the cached
275 # pattern: its rows stop summing to 1 and every later use of that cache
276 # silently reads modified activations.
277 if exclude_bos or exclude_current_token:
278 attention_pattern = attention_pattern.clone()
279 if exclude_bos:
280 attention_pattern[:, 0] = 0
281 if exclude_current_token:
282 attention_pattern.fill_diagonal_(0)
283 score = attention_pattern * detection_pattern
284 return (score.sum() / attention_pattern.sum()).item()
286 # abs
288 abs_diff = (attention_pattern - detection_pattern).abs()
289 assert (abs_diff - torch.tril(abs_diff).to(abs_diff.device)).sum() == 0
291 size = len(abs_diff)
292 if exclude_bos:
293 abs_diff[:, 0] = 0
294 if exclude_current_token:
295 abs_diff.fill_diagonal_(0)
297 return 1 - round((abs_diff.mean() * size).item(), 3)
298 return 1 - round((abs_diff.mean() * size).item(), 3)