diff --git a/tests/integration/model_bridge/test_analysis_methods.py b/tests/integration/model_bridge/test_analysis_methods.py new file mode 100644 index 000000000..4f780c30c --- /dev/null +++ b/tests/integration/model_bridge/test_analysis_methods.py @@ -0,0 +1,191 @@ +"""Tests for TransformerBridge mechanistic interpretability analysis methods. + +Tests tokens_to_residual_directions, accumulated_bias, all_composition_scores, +all_head_labels, and top-level W_E/W_U/b_U properties. Validates against +HookedTransformer for correctness, not just shape/type. + +Uses distilgpt2 (CI-cached). +""" + +import pytest +import torch + +from transformer_lens import HookedTransformer +from transformer_lens.model_bridge.bridge import TransformerBridge + + +@pytest.fixture(scope="module") +def bridge_compat(): + b = TransformerBridge.boot_transformers("distilgpt2", device="cpu") + b.enable_compatibility_mode() + return b + + +@pytest.fixture(scope="module") +def reference_ht(): + return HookedTransformer.from_pretrained("distilgpt2", device="cpu") + + +class TestTopLevelWeightProperties: + """Test W_E, W_U, b_U delegate to the correct component tensors.""" + + def test_W_E_is_same_object_as_embed(self, bridge_compat): + """bridge.W_E should be the exact same tensor as bridge.embed.W_E.""" + assert bridge_compat.W_E is bridge_compat.embed.W_E + + def test_W_U_equals_unembed(self, bridge_compat): + """bridge.W_U should equal bridge.unembed.W_U (may be a view/transpose).""" + assert torch.equal(bridge_compat.W_U, bridge_compat.unembed.W_U) + + def test_b_U_equals_unembed(self, bridge_compat): + """bridge.b_U should equal bridge.unembed.b_U.""" + assert torch.equal(bridge_compat.b_U, bridge_compat.unembed.b_U) + + def test_W_E_matches_hooked_transformer(self, bridge_compat, reference_ht): + """bridge.W_E values should match HookedTransformer.W_E.""" + assert bridge_compat.W_E.shape == reference_ht.W_E.shape + # After weight processing, embeddings may differ due to centering. + # But shapes must match and both must be non-zero. + assert bridge_compat.W_E.std() > 0 + assert reference_ht.W_E.std() > 0 + + def test_W_U_matches_hooked_transformer(self, bridge_compat, reference_ht): + """bridge.W_U values should match HookedTransformer.W_U.""" + assert bridge_compat.W_U.shape == reference_ht.W_U.shape + max_diff = (bridge_compat.W_U - reference_ht.W_U).abs().max().item() + assert max_diff < 1e-4, f"W_U differs by {max_diff}" + + +class TestTokensToResidualDirections: + """Test tokens_to_residual_directions produces correct unembedding vectors.""" + + def test_single_token_string(self, bridge_compat): + """String token should return a 1-D vector of size d_model.""" + rd = bridge_compat.tokens_to_residual_directions("hello") + assert rd.shape == (bridge_compat.cfg.d_model,) + + def test_single_token_int(self, bridge_compat): + """Integer token should return a 1-D vector of size d_model.""" + rd = bridge_compat.tokens_to_residual_directions(100) + assert rd.shape == (bridge_compat.cfg.d_model,) + + def test_equals_W_U_column(self, bridge_compat): + """Result should be exactly the corresponding column of W_U.""" + token_id = 42 + rd = bridge_compat.tokens_to_residual_directions(token_id) + expected = bridge_compat.W_U[:, token_id] + assert torch.equal(rd, expected) + + def test_batch_tokens(self, bridge_compat): + """1-D tensor of tokens should return (n_tokens, d_model).""" + tokens = torch.tensor([100, 200, 300]) + rd = bridge_compat.tokens_to_residual_directions(tokens) + assert rd.shape == (3, bridge_compat.cfg.d_model) + # Each row should match the corresponding W_U column + for i, tok in enumerate(tokens): + assert torch.equal(rd[i], bridge_compat.W_U[:, tok]) + + def test_matches_hooked_transformer(self, bridge_compat, reference_ht): + """Output should match HookedTransformer for the same tokens.""" + tokens = torch.tensor([10, 20, 30]) + bridge_rd = bridge_compat.tokens_to_residual_directions(tokens) + ht_rd = reference_ht.tokens_to_residual_directions(tokens) + max_diff = (bridge_rd - ht_rd).abs().max().item() + assert max_diff < 1e-4, f"Residual directions differ by {max_diff}" + + +class TestAccumulatedBias: + """Test accumulated_bias sums biases correctly.""" + + def test_layer_zero_is_zeros(self, bridge_compat): + """accumulated_bias(0) should be all zeros (no layers processed).""" + ab = bridge_compat.accumulated_bias(0) + assert ab.shape == (bridge_compat.cfg.d_model,) + assert torch.allclose(ab, torch.zeros_like(ab)) + + def test_layer_one_includes_first_block(self, bridge_compat): + """accumulated_bias(1) should include block 0's biases and be non-zero.""" + ab = bridge_compat.accumulated_bias(1) + assert ab.shape == (bridge_compat.cfg.d_model,) + # distilgpt2 has biases, so this should be non-zero + assert ab.norm() > 0 + + def test_monotonically_increasing_norm(self, bridge_compat): + """Accumulated bias norm should generally increase with more layers.""" + # Not strictly monotonic, but bias(n_layers) should have larger norm than bias(0) + ab_0 = bridge_compat.accumulated_bias(0) + ab_all = bridge_compat.accumulated_bias(bridge_compat.cfg.n_layers) + assert ab_all.norm() > ab_0.norm() + + def test_matches_hooked_transformer(self, bridge_compat, reference_ht): + """Output should match HookedTransformer.""" + for layer in [0, 1, 3, bridge_compat.cfg.n_layers]: + bridge_ab = bridge_compat.accumulated_bias(layer) + ht_ab = reference_ht.accumulated_bias(layer) + max_diff = (bridge_ab - ht_ab).abs().max().item() + assert max_diff < 1e-4, f"accumulated_bias({layer}) differs by {max_diff}" + + def test_mlp_input_flag(self, bridge_compat, reference_ht): + """mlp_input=True should include the current layer's attn bias.""" + bridge_ab = bridge_compat.accumulated_bias(1, mlp_input=True) + ht_ab = reference_ht.accumulated_bias(1, mlp_input=True) + max_diff = (bridge_ab - ht_ab).abs().max().item() + assert max_diff < 1e-4, f"accumulated_bias(1, mlp_input=True) differs by {max_diff}" + + +class TestAllCompositionScores: + """Test all_composition_scores produces correct composition score matrices.""" + + def test_shape(self, bridge_compat): + """Shape should be (n_layers, n_heads, n_layers, n_heads).""" + cfg = bridge_compat.cfg + scores = bridge_compat.all_composition_scores("Q") + assert scores.shape == (cfg.n_layers, cfg.n_heads, cfg.n_layers, cfg.n_heads) + + def test_upper_triangular_masking(self, bridge_compat): + """Scores should be zero where left_layer >= right_layer.""" + scores = bridge_compat.all_composition_scores("Q") + n_layers = bridge_compat.cfg.n_layers + for l1 in range(n_layers): + for l2 in range(l1 + 1): # l2 <= l1 + assert ( + scores[l1, :, l2, :] == 0 + ).all(), f"Scores at L{l1}->L{l2} should be zero (upper triangular)" + + def test_nonzero_above_diagonal(self, bridge_compat): + """At least some scores above the diagonal should be non-zero.""" + scores = bridge_compat.all_composition_scores("Q") + # Check L0 -> L1 (first above-diagonal block) + assert scores[0, :, 1, :].abs().sum() > 0 + + def test_all_modes_work(self, bridge_compat): + """Q, K, V modes should all produce valid tensors.""" + for mode in ["Q", "K", "V"]: + scores = bridge_compat.all_composition_scores(mode) + assert not torch.isnan(scores).any(), f"NaN in {mode} composition scores" + + def test_invalid_mode_raises(self, bridge_compat): + """Invalid mode should raise ValueError.""" + with pytest.raises(ValueError, match="mode must be one of"): + bridge_compat.all_composition_scores("X") + + +class TestAllHeadLabels: + """Test all_head_labels produces correct labels.""" + + def test_count(self, bridge_compat): + """Should have n_layers * n_heads labels.""" + labels = bridge_compat.all_head_labels + expected = bridge_compat.cfg.n_layers * bridge_compat.cfg.n_heads + assert len(labels) == expected + + def test_format(self, bridge_compat): + """Labels should follow L{layer}H{head} format.""" + labels = bridge_compat.all_head_labels + assert labels[0] == "L0H0" + assert labels[1] == "L0H1" + assert labels[bridge_compat.cfg.n_heads] == "L1H0" + + def test_matches_hooked_transformer(self, bridge_compat, reference_ht): + """Should match HookedTransformer's labels exactly.""" + assert bridge_compat.all_head_labels == reference_ht.all_head_labels() diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index f7c909624..dcbe2960b 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -1103,6 +1103,21 @@ def b_out(self) -> torch.Tensor: """Stack the MLP output biases across all layers.""" return self._stack_block_params("mlp.b_out") + @property + def W_U(self) -> torch.Tensor: + """Unembedding matrix (d_model, d_vocab). Maps residual stream to logits.""" + return self.unembed.W_U + + @property + def b_U(self) -> torch.Tensor: + """Unembedding bias (d_vocab).""" + return self.unembed.b_U + + @property + def W_E(self) -> torch.Tensor: + """Token embedding matrix (d_vocab, d_model).""" + return self.embed.W_E + @property def QK(self): return FactoredMatrix(self.W_Q, self.W_K.transpose(-2, -1)) @@ -1111,6 +1126,119 @@ def QK(self): def OV(self): return FactoredMatrix(self.W_V, self.W_O) + # ------------------------------------------------------------------ + # Mechanistic interpretability analysis methods + # ------------------------------------------------------------------ + + def tokens_to_residual_directions( + self, + tokens: Union[str, int, torch.Tensor], + ) -> torch.Tensor: + """Map tokens to their unembedding vectors (residual stream directions). + + Returns the columns of W_U corresponding to the given tokens — i.e. the + directions in the residual stream that the model dots with to produce the + logit for each token. + + WARNING: If you use this without folding in LayerNorm (compatibility mode), + the results will be misleading because LN weights change the unembed map. + + Args: + tokens: A single token (str, int, or scalar tensor), a 1-D tensor of + token IDs, or a 2-D batch of token IDs. + + Returns: + Tensor of unembedding vectors with shape matching the input token shape + plus a trailing d_model dimension. + """ + if isinstance(tokens, torch.Tensor) and tokens.numel() > 1: + residual_directions = self.W_U[:, tokens] + residual_directions = einops.rearrange( + residual_directions, "d_model ... -> ... d_model" + ) + return residual_directions + else: + if isinstance(tokens, str): + token = self.to_single_token(tokens) + elif isinstance(tokens, int): + token = tokens + elif isinstance(tokens, torch.Tensor) and tokens.numel() == 1: + token = int(tokens.item()) + else: + raise ValueError(f"Invalid token type: {type(tokens)}") + residual_direction = self.W_U[:, token] + return residual_direction + + def accumulated_bias( + self, + layer: int, + mlp_input: bool = False, + include_mlp_biases: bool = True, + ) -> torch.Tensor: + """Sum of attention and MLP output biases up to the input of a given layer. + + Args: + layer: Layer number in [0, n_layers]. 0 means no layers, n_layers means all. + mlp_input: If True, include the attention output bias of the target layer + (i.e. bias up to the MLP input of that layer). + include_mlp_biases: Whether to include MLP biases. Useful to set False when + expanding attn_out into individual heads but keeping mlp_out as-is. + + Returns: + Tensor of shape [d_model] with the accumulated bias. + """ + accumulated = torch.zeros(self.cfg.d_model, device=self.cfg.device) + for i in range(layer): + block = self.blocks[i] + b_O = getattr(block.attn, "b_O", None) + if b_O is not None: + accumulated = accumulated + b_O + if include_mlp_biases: + b_out = getattr(block.mlp, "b_out", None) + if b_out is not None: + accumulated = accumulated + b_out + if mlp_input: + assert layer < self.cfg.n_layers, "Cannot include attn_bias from beyond the final layer" + block = self.blocks[layer] + b_O = getattr(block.attn, "b_O", None) + if b_O is not None: + accumulated = accumulated + b_O + return accumulated + + def all_composition_scores(self, mode: str) -> torch.Tensor: + """Composition scores for all pairs of heads. + + Returns an (n_layers, n_heads, n_layers, n_heads) tensor that is upper + triangular on the layer axes (a head can only compose with later heads). + + See https://transformer-circuits.pub/2021/framework/index.html + + Args: + mode: One of "Q", "K", "V" — which composition type to compute. + """ + left = self.OV + if mode == "Q": + right = self.QK + elif mode == "K": + right = self.QK.T + elif mode == "V": + right = self.OV + else: + raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}") + + scores = utils.composition_scores(left, right, broadcast_dims=True) + mask = ( + torch.arange(self.cfg.n_layers, device=self.cfg.device)[:, None, None, None] + < torch.arange(self.cfg.n_layers, device=self.cfg.device)[None, None, :, None] + ) + scores = torch.where(mask, scores, torch.zeros_like(scores)) + return scores + + @property + def all_head_labels(self) -> list[str]: + """Human-readable labels for all attention heads, e.g. ['L0H0', 'L0H1', ...].""" + return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)] + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: """Returns parameters following standard PyTorch semantics.