diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md index e184da557..f55697188 100644 --- a/docs/source/content/analysis_tools.md +++ b/docs/source/content/analysis_tools.md @@ -109,6 +109,101 @@ Both node and edge granularity use plain attribution with `ig_steps=1`. API: {func}`~transformer_lens.tools.analysis.attribution_patching.attribution_patch`. +#### Faithfulness: check that a circuit actually explains the behavior + +A ranked edge list is a hypothesis, not a result. `faithfulness()` tests one aligned +clean/corrupt pair at a time by ablating every edge *outside* the candidate circuit +and reporting how much of the clean-to-corrupt metric gap the circuit recovers: + +```python +from transformer_lens.tools.analysis import ( + EdgeAttributionConfig, + attribution_patch, + faithfulness, +) + +ranked = attribution_patch( + model, clean, corrupt, metric_fn, config=EdgeAttributionConfig(granularity="edge") +) +report = faithfulness(model, clean, corrupt, metric_fn, ranked.top_edges(k=50)) +print(report.recovered, report.circuit_size, report.total_edges) +``` + +`recovered` is a fraction, not a percentage: `1.0` means the circuit reproduces +the clean metric, `0.0` means it reproduces the corrupt metric. Values outside +`[0, 1]` are possible and meaningful, since a circuit can overshoot. + +The residual stream is a running sum, so each reader's input is rebuilt exactly +by subtracting the excluded writers' live contributions and adding their +replacements. This is a forward-only measurement: it takes no gradient, so it is +independent of the linearization that ranked the edges in the first place. + +`FaithfulnessConfig(ablation=...)` selects the replacement. The default, +`"corrupt"`, substitutes the corrupt run's own contribution, which is the +evaluation the external EAP-IG reference reports. `"mean"` requires an explicit, +held-out calibration batch through the keyword-only `mean_tokens` argument: + +```python +from transformer_lens.tools.analysis import FaithfulnessConfig + +calibration_tokens = model.to_tokens([ + "The capital of Spain is", + "The capital of Italy is", +]).to(clean.device) +mean_report = faithfulness( + model, clean, corrupt, metric_fn, ranked.top_edges(k=50), + config=FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration_tokens, +) +``` + +Calibration must be a nonempty `[n_calibration, seq]` batch of `torch.int32` or +`torch.int64` token ids, with the evaluated sequence length and the same device +as the clean/corrupt token tensors. Align the sequence positions before averaging; +the helper averages examples only, preserving each position and attention head. +It captures calibration activations with the same model, hook configuration, and +weight-processing mode used for evaluation. + +Mean-mode calls without calibration now raise `ValueError`; there is no fallback +to the evaluated pair. Exact calibration rows equal to either endpoint are also +rejected. This check cannot establish dataset independence: callers must choose +held-out examples and avoid semantically equivalent leakage. Passing `mean_tokens` +in corrupt mode is rejected rather than silently ignored. + +Recovery is still `(circuit_metric - corrupt_metric) / (clean_metric - corrupt_metric)`. +An empty mean-ablated circuit can have nonzero recovery because its calibration +baseline need not reproduce the corrupt run. A full circuit still reproduces the +clean metric. Compare against a random edge set of the same size, using the same +calibration batch, before claiming a circuit is meaningful. + +`report.edge_class_recovered` is circuit-relative leave-one-out recovery. For a +candidate `C` and graph class `E_t`, each entry evaluates `C` with its retained +edges in `E_t` removed, not the full graph with the class removed. Edges originally +outside `C` stay ablated. Every diagnostic uses the aggregate's replacements, +clean/corrupt endpoints, and normalization: + +```text +edge_class_recovered[t] = recovered(C minus E_t) +``` + +An absent class leaves the candidate unchanged, so its entry equals +`report.recovered`. For an empty circuit every class entry equals aggregate +empty-circuit recovery, including with mean replacements. Duplicate circuit +entries count once; `report.circuit_size` is the number of unique retained edges. +`report.edge_class_counts` remains the full-graph count per class, not the number +retained in `C`. + +The difference `report.recovered - report.edge_class_recovered[t]` measures +conditional dependence on that class's retained edges. Interactions and inhibitory +edges can increase recovery after removal; contributions are not additive or +necessarily positive. This is not a measurement of attribution ranking error. +Q/K inputs feed attention scores, but that alone does not establish which class +is ranked most accurately. LayerNorm, MLPs, and downstream attention can introduce +nonlinearity for the other classes too. The breakout costs one extra forward per +class present in the candidate; absent classes reuse aggregate recovery. + +API: {func}`~transformer_lens.tools.analysis.attribution_patching.faithfulness`. + ### Direct Path Patching: state the path and its approximation `get_act_patch_direct_path` fixes a source head and sweeps later destination heads, diff --git a/tests/integration/test_attribution_patching.py b/tests/integration/test_attribution_patching.py index ca3a8a2dc..8e296bb22 100644 --- a/tests/integration/test_attribution_patching.py +++ b/tests/integration/test_attribution_patching.py @@ -1,31 +1,34 @@ -"""Integration guard: ``attribution_patch`` yields finite node scores on a real Bridge. - -The model-free unit suite runs against ``_LinearToyBridge``, which overrides -``hook_dict``, ``hooks()``, and ``check_hooks_to_add`` — the three behaviours the -two real-Bridge failures depend on. Its hook points have no conversion and its -gate check is a no-op, so a green unit suite says nothing about a real Bridge. - -This test boots a real GPT-2 Bridge and exercises the two paths the toy bridge -hides: - -- gradients captured *through* hook conversions — ``blocks.*.attn.hook_z`` hands a - reshaped ``[batch, seq, n_heads, d_head]`` view to the forward hook, so the - ``attn_head_out`` family only scores if the backward hook delivers the gradient - in that converted shape; and -- the default ``names_filter`` (``None``) falling back to the node hook set on a - real ``hook_dict``, which also exposes gated points (``hook_mlp_in``, - ``attn.hook_result``, split-QKV inputs) that ``add_hook`` would reject. +"""Attribution patching and faithfulness on a raw GPT-2 Bridge. + +Covers converted gradient caches, gated reader forks, independent partial +intervention references, and circuit-relative class removal. """ from __future__ import annotations import math +from contextlib import contextmanager +from typing import Any, Callable, Iterator import pytest import torch from transformer_lens.model_bridge import TransformerBridge -from transformer_lens.tools.analysis.attribution_patching import attribution_patch +from transformer_lens.tools.analysis.attribution_patching import ( + EdgeAttributionConfig, + EdgeClass, + FaithfulnessConfig, + GradientCache, + Node, + NodeKind, + _ablate_edges, + _edge_hook_flags, + _edge_hook_names, + _mean_writer_contributions, + attribution_patch, + enumerate_edges, + faithfulness, +) CLEAN_PROMPT = "The capital of France is" CORRUPT_PROMPT = "The capital of Russia is" @@ -43,6 +46,311 @@ def metric(logits: torch.Tensor) -> torch.Tensor: return metric +def _capture_hook(name: str, cache: dict[str, torch.Tensor]) -> Callable[..., None]: + def capture(tensor: torch.Tensor, *, hook: Any) -> None: + del hook + cache[name] = tensor.detach().clone() + + return capture + + +def _capture_gpt2_forward( + model: TransformerBridge, + tokens: torch.Tensor, + overrides: list[tuple[str, Callable]] | None = None, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + cache: dict[str, torch.Tensor] = {} + observers = [ + (name, _capture_hook(name, cache)) for name in _edge_hook_names(model.cfg.n_layers) + ] + with torch.no_grad(), model.hooks(fwd_hooks=(overrides or []) + observers): + logits = model(tokens) + return logits, cache + + +@pytest.fixture(scope="module") +def gpt2_edge_endpoints(gpt2_bridge): + clean = gpt2_bridge.to_tokens(CLEAN_PROMPT) + corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) + assert clean.shape == corrupt.shape + assert gpt2_bridge.original_model.config._attn_implementation == "eager" + with _edge_hook_flags(gpt2_bridge): + clean_logits, clean_values = _capture_gpt2_forward(gpt2_bridge, clean) + _, corrupt_values = _capture_gpt2_forward(gpt2_bridge, corrupt) + cache = GradientCache(clean_values, {}, torch.tensor(0.0)) + edges = enumerate_edges(gpt2_bridge, cache) + return clean, clean_logits, clean_values, corrupt_values, edges + + +def _override_reader(index: tuple[int, ...], value: torch.Tensor) -> Callable[..., torch.Tensor]: + def override(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor: + del hook + replaced = tensor.clone() + replaced[index] = value + return replaced + + return override + + +def _capture_gpt2_ablation( + model: TransformerBridge, + tokens: torch.Tensor, + edges: list[tuple[Node, Node]], + circuit: list[tuple[Node, Node]], + replacements: dict[str, torch.Tensor], + monkeypatch: pytest.MonkeyPatch, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + cache: dict[str, torch.Tensor] = {} + original_hooks = model.hooks + + @contextmanager + def hooks(fwd_hooks: list[tuple[str, Callable]], **kwargs: Any) -> Iterator[Any]: + # Appending observers preserves the production correction order. + observers = [ + (name, _capture_hook(name, cache)) for name in _edge_hook_names(model.cfg.n_layers) + ] + with original_hooks(fwd_hooks=fwd_hooks + observers, **kwargs): + yield model + + before = _bridge_hook_state(model) + with monkeypatch.context() as patch: + patch.setattr(model, "hooks", hooks) + logits = _ablate_edges(model, tokens, edges, circuit, replacements) + assert _bridge_hook_state(model) == before + return logits, cache + + +@pytest.mark.parametrize( + "kind,name,layer,head", + [ + ("q_input", "blocks.1.attn.hook_q_input", 1, 3), + ("k_input", "blocks.1.attn.hook_k_input", 1, 4), + ("v_input", "blocks.1.attn.hook_v_input", 1, 5), + ("mlp_in", "blocks.1.hook_mlp_in", 1, None), + ("logits", "blocks.11.hook_resid_post", 11, None), + ], +) +def test_partial_gpt2_ablation_matches_explicit_reader_override( + gpt2_bridge, + gpt2_edge_endpoints, + monkeypatch: pytest.MonkeyPatch, + kind: NodeKind, + name: str, + layer: int, + head: int | None, +) -> None: + clean, clean_logits, clean_values, corrupt_values, edges = gpt2_edge_endpoints + position = clean.shape[1] - 1 + writer = Node("attn_head_out", position, layer=0, head=1) + reader = Node(kind, position, layer=layer, head=head) + excluded = (writer, reader) + assert excluded in edges + circuit = [edge for edge in edges if edge != excluded] + assert len(circuit) == len(edges) - 1 + index = (0, position) if head is None else (0, position, head) + delta = ( + corrupt_values["blocks.0.attn.hook_result"][0, position, 1] + - clean_values["blocks.0.attn.hook_result"][0, position, 1] + ) + assert delta.abs().max() > 1e-5 + expected_reader = clean_values[name].clone() + expected_reader[index] = ( + clean_values[name][index] + - clean_values["blocks.0.attn.hook_result"][0, position, 1] + + corrupt_values["blocks.0.attn.hook_result"][0, position, 1] + ) + with _edge_hook_flags(gpt2_bridge): + expected, reference = _capture_gpt2_forward( + gpt2_bridge, clean, [(name, _override_reader(index, expected_reader[index]))] + ) + actual, observed = _capture_gpt2_ablation( + gpt2_bridge, clean, edges, circuit, corrupt_values, monkeypatch + ) + # Both routes apply identical fp32 slice arithmetic to the same clean reader. + torch.testing.assert_close(reference[name], expected_reader, atol=0, rtol=0) + torch.testing.assert_close(observed[name], expected_reader, atol=0, rtol=0) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert (actual - clean_logits).abs().max() > 1e-5 + + +def test_coupled_partial_gpt2_ablation_matches_staged_reader_overrides( + gpt2_bridge, + gpt2_edge_endpoints, + monkeypatch: pytest.MonkeyPatch, +) -> None: + clean, clean_logits, clean_values, corrupt_values, edges = gpt2_edge_endpoints + position = clean.shape[1] - 1 + upstream_name = "blocks.1.attn.hook_q_input" + downstream_name = "blocks.2.attn.hook_k_input" + upstream_index = (0, position, 3) + downstream_index = (0, position, 4) + upstream_value = ( + clean_values[upstream_name][upstream_index] + - clean_values["blocks.0.attn.hook_result"][0, position, 1] + + corrupt_values["blocks.0.attn.hook_result"][0, position, 1] + ) + upstream_override = (upstream_name, _override_reader(upstream_index, upstream_value)) + with _edge_hook_flags(gpt2_bridge): + _, staged = _capture_gpt2_forward(gpt2_bridge, clean, [upstream_override]) + live = staged["blocks.1.attn.hook_result"][0, position, 3] + assert (live - clean_values["blocks.1.attn.hook_result"][0, position, 3]).abs().max() > 1e-6 + downstream_value = ( + staged[downstream_name][downstream_index] + - live + + corrupt_values["blocks.1.attn.hook_result"][0, position, 3] + ) + overrides = [ + upstream_override, + (downstream_name, _override_reader(downstream_index, downstream_value)), + ] + expected, reference = _capture_gpt2_forward(gpt2_bridge, clean, overrides) + excluded = { + ( + Node("attn_head_out", position, layer=0, head=1), + Node("q_input", position, layer=1, head=3), + ), + ( + Node("attn_head_out", position, layer=1, head=3), + Node("k_input", position, layer=2, head=4), + ), + } + assert excluded <= set(edges) + circuit = [edge for edge in edges if edge not in excluded] + assert len(circuit) == len(edges) - 2 + actual, observed = _capture_gpt2_ablation( + gpt2_bridge, clean, edges, circuit, corrupt_values, monkeypatch + ) + for name in (upstream_name, downstream_name, "blocks.1.attn.hook_result"): + torch.testing.assert_close(observed[name], reference[name], atol=0, rtol=0) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert (actual - clean_logits).abs().max() > 1e-5 + + +def test_partial_mean_gpt2_ablation_matches_held_out_reader_override( + gpt2_bridge, + gpt2_edge_endpoints, + monkeypatch: pytest.MonkeyPatch, +) -> None: + clean, clean_logits, clean_values, _, edges = gpt2_edge_endpoints + corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) + calibration = gpt2_bridge.to_tokens(["The capital of Spain is", "The capital of Italy is"]) + assert calibration.shape == (2, clean.shape[1]) + position = clean.shape[1] - 1 + name = "blocks.1.attn.hook_q_input" + index = (0, position, 3) + excluded = ( + Node("attn_head_out", position, layer=0, head=1), + Node("q_input", position, layer=1, head=3), + ) + assert excluded in edges + circuit = [edge for edge in edges if edge != excluded] + with _edge_hook_flags(gpt2_bridge): + _, first = _capture_gpt2_forward(gpt2_bridge, calibration[:1]) + _, second = _capture_gpt2_forward(gpt2_bridge, calibration[1:]) + names = ["hook_embed"] + for layer in range(gpt2_bridge.cfg.n_layers): + names.extend([f"blocks.{layer}.attn.hook_result", f"blocks.{layer}.hook_mlp_out"]) + reference_means = {hook: (first[hook] + second[hook]) / 2 for hook in names} + actual_means = _mean_writer_contributions(gpt2_bridge, calibration) + assert actual_means.keys() == reference_means.keys() + for hook in names: + torch.testing.assert_close(actual_means[hook], reference_means[hook], atol=0, rtol=0) + expected_reader = clean_values[name].clone() + expected_reader[index] = ( + clean_values[name][index] + - clean_values["blocks.0.attn.hook_result"][0, position, 1] + + reference_means["blocks.0.attn.hook_result"][0, position, 1] + ) + assert (expected_reader - clean_values[name]).abs().max() > 1e-5 + expected, reference = _capture_gpt2_forward( + gpt2_bridge, clean, [(name, _override_reader(index, expected_reader[index]))] + ) + actual, observed = _capture_gpt2_ablation( + gpt2_bridge, clean, edges, circuit, actual_means, monkeypatch + ) + corrupt_logits, _ = _capture_gpt2_forward(gpt2_bridge, corrupt) + answer_id = int(gpt2_bridge.to_tokens(" Paris")[0, -1]) + wrong_id = int(gpt2_bridge.to_tokens(" Moscow")[0, -1]) + metric = _logit_diff_metric(answer_id, wrong_id) + report = faithfulness( + gpt2_bridge, + clean, + corrupt, + metric, + circuit, + FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration, + ) + torch.testing.assert_close(reference[name], expected_reader, atol=0, rtol=0) + torch.testing.assert_close(observed[name], expected_reader, atol=0, rtol=0) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert (actual - clean_logits).abs().max() > 1e-5 + recovered = float( + (metric(expected) - metric(corrupt_logits)) + / (metric(clean_logits) - metric(corrupt_logits)) + ) + assert report.recovered == pytest.approx(recovered, abs=1e-6) + + +def _bridge_hook_state(model: TransformerBridge) -> tuple: + flags = tuple( + getattr(model.cfg, name) + for name in ("use_attn_result", "use_split_qkv_input", "use_hook_mlp_in", "use_attn_in") + ) + handles = { + name: ( + tuple(id(handle) for handle in point.fwd_hooks), + tuple(id(handle) for handle in point.bwd_hooks), + ) + for name, point in model.hook_dict.items() + } + return flags, handles, model.context_level + + +@pytest.mark.parametrize("fail_at", [None, 1, 3]) +def test_forward_only_faithfulness_restores_gpt2_hooks_after_metric_error( + gpt2_bridge, + gpt2_edge_endpoints, + fail_at: int | None, +) -> None: + clean, _, _, _, edges = gpt2_edge_endpoints + corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) + answer_id = int(gpt2_bridge.to_tokens(" Paris")[0, -1]) + wrong_id = int(gpt2_bridge.to_tokens(" Moscow")[0, -1]) + metric = _logit_diff_metric(answer_id, wrong_id) + calls = 0 + + def checked_metric(logits: torch.Tensor) -> torch.Tensor: + nonlocal calls + calls += 1 + assert not logits.requires_grad + if calls == fail_at: + raise RuntimeError("metric failed") + return metric(logits) + + original_state = _bridge_hook_state(gpt2_bridge) + with _edge_hook_flags(gpt2_bridge): + gpt2_bridge.set_use_split_qkv_input(False) + gpt2_bridge.set_use_hook_mlp_in(False) + gpt2_bridge.set_use_attn_in(True) + try: + with gpt2_bridge.hooks(fwd_hooks=[("hook_embed", lambda tensor, **kwargs: None)]): + before = _bridge_hook_state(gpt2_bridge) + with torch.no_grad(): + if fail_at is None: + report = faithfulness( + gpt2_bridge, clean, corrupt, checked_metric, edges[:-1] + ) + assert math.isfinite(report.recovered) + else: + with pytest.raises(RuntimeError, match="metric failed"): + faithfulness(gpt2_bridge, clean, corrupt, checked_metric, edges[:-1]) + assert _bridge_hook_state(gpt2_bridge) == before + finally: + gpt2_bridge.set_use_attn_in(False) + assert _bridge_hook_state(gpt2_bridge) == original_state + + def test_attribution_patch_scores_every_node_finite_on_real_bridge(gpt2_bridge) -> None: clean = gpt2_bridge.to_tokens(CLEAN_PROMPT) corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) @@ -67,7 +375,146 @@ def test_attribution_patch_scores_every_node_finite_on_real_bridge(gpt2_bridge) for node, score in result.node_scores.items(): assert math.isfinite(score), f"non-finite score for {node}" - # All three node families are present — attn_head_out is the family the hook + # All three node families are present -- attn_head_out is the family the hook # conversion bug broke, mlp_out and embed round out the node graph. families = {node.kind for node in result.node_scores} assert families == {"embed", "attn_head_out", "mlp_out"} + + +def test_faithfulness_recovers_most_of_the_metric_on_gpt2_small(gpt2_bridge) -> None: + """A ranked edge circuit recovers far more of the gap than a random one. + + The ablation rewrites each reader's input at its pre-LN fork hook, so this + only works if the real Bridge's split-QKV and MLP-entry hooks sit where the + rewrite assumes. A circuit built from the top-ranked edges should recover a + large share of the clean-to-corrupt gap; a random edge set of the same size + should recover essentially none, which is what makes the number meaningful + rather than an artifact of ablating almost nothing. + + Thresholds are set from measured values on this exact model and prompt pair + (gpt2-small, fp32, CPU): the top 200 edges recover about 0.62 of the gap, + while a random 200 of the 194,946 edges recover about 0.00. Recovery grows + with the budget -- roughly 0.29 at 50 edges and 0.997 at 1000 -- so the + budget is fixed here rather than left to drift. + """ + clean = gpt2_bridge.to_tokens(CLEAN_PROMPT) + corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) + assert clean.shape == corrupt.shape, "prompts must tokenize to the same length" + + answer_id = int(gpt2_bridge.to_tokens(" Paris")[0, -1].item()) + wrong_id = int(gpt2_bridge.to_tokens(" Moscow")[0, -1].item()) + metric_fn = _logit_diff_metric(answer_id, wrong_id) + + result = attribution_patch( + gpt2_bridge, + clean, + corrupt, + metric_fn, + config=EdgeAttributionConfig(granularity="edge"), + ) + budget = 200 + ranked = result.top_edges(k=budget) + assert len(ranked) == budget + + report = faithfulness(gpt2_bridge, clean, corrupt, metric_fn, ranked) + + assert math.isfinite(report.recovered) + assert report.circuit_size == budget + assert report.total_edges > report.circuit_size + assert report.full_metric != pytest.approx(report.corrupt_metric) + assert report.recovered > 0.5 + + # A random edge set of the same size recovers essentially nothing, so the + # ranking is doing real work rather than the number being an artifact of the + # circuit's size. + all_edges = list(result.edge_scores) + generator = torch.Generator().manual_seed(0) + sampled = torch.randperm(len(all_edges), generator=generator)[:budget].tolist() + random_report = faithfulness( + gpt2_bridge, clean, corrupt, metric_fn, [all_edges[index] for index in sampled] + ) + + assert random_report.circuit_size == report.circuit_size + assert random_report.recovered < 0.1 + assert report.recovered > random_report.recovered + 0.4 + + +def _gpt2_class_reader_names(edge_class: EdgeClass, n_layers: int) -> list[str]: + suffixes = { + "into_qk": ("attn.hook_q_input", "attn.hook_k_input"), + "into_v": ("attn.hook_v_input",), + "into_mlp": ("hook_mlp_in",), + } + if edge_class == "into_logits": + return [f"blocks.{n_layers - 1}.hook_resid_post"] + return [ + f"blocks.{layer}.{suffix}" for layer in range(n_layers) for suffix in suffixes[edge_class] + ] + + +def _gpt2_class_overrides( + clean: torch.Tensor, + clean_values: dict[str, torch.Tensor], + corrupt_values: dict[str, torch.Tensor], + edge_class: EdgeClass, + n_layers: int, +) -> tuple[tuple[Node, Node], list[tuple[str, Callable]]]: + position = clean.shape[1] - 2 + kind: NodeKind = "q_input" if edge_class == "into_v" else "v_input" + name = "blocks.0.attn.hook_q_input" if kind == "q_input" else "blocks.0.attn.hook_v_input" + index = (0, position, 0) + base_value = ( + clean_values[name][index] + - clean_values["hook_embed"][0, position] + + corrupt_values["hook_embed"][0, position] + ) + base_edge = (Node("embed", position), Node(kind, position, layer=0, head=0)) + overrides = [(name, _override_reader(index, base_value))] + # Replacing every writer in a class preserves the prompt-independent residual background. + overrides.extend( + (reader_name, _override_reader((), corrupt_values[reader_name])) + for reader_name in _gpt2_class_reader_names(edge_class, n_layers) + ) + return base_edge, overrides + + +@pytest.mark.parametrize("edge_class", ["into_qk", "into_v", "into_mlp", "into_logits"]) +def test_edge_class_gpt2_removal_keeps_candidate_exclusions( + gpt2_bridge, + gpt2_edge_endpoints, + monkeypatch: pytest.MonkeyPatch, + edge_class: EdgeClass, +) -> None: + clean, _, clean_values, corrupt_values, edges = gpt2_edge_endpoints + corrupt = gpt2_bridge.to_tokens(CORRUPT_PROMPT) + base_edge, overrides = _gpt2_class_overrides( + clean, clean_values, corrupt_values, edge_class, gpt2_bridge.cfg.n_layers + ) + class_names = _gpt2_class_reader_names(edge_class, gpt2_bridge.cfg.n_layers) + assert base_edge in edges + assert base_edge[1].hook_name not in class_names + circuit = [edge for edge in edges if edge != base_edge] + reduced = [edge for edge in circuit if edge[1].hook_name not in class_names] + answer_id = int(gpt2_bridge.to_tokens(" Paris")[0, -1]) + wrong_id = int(gpt2_bridge.to_tokens(" Moscow")[0, -1]) + metric = _logit_diff_metric(answer_id, wrong_id) + with _edge_hook_flags(gpt2_bridge): + expected, reference = _capture_gpt2_forward(gpt2_bridge, clean, overrides) + restored, _ = _capture_gpt2_forward(gpt2_bridge, clean, overrides[1:]) + actual, observed = _capture_gpt2_ablation( + gpt2_bridge, clean, edges, reduced, corrupt_values, monkeypatch + ) + report = faithfulness(gpt2_bridge, clean, corrupt, metric, circuit) + # Dense corrections and cached reader overrides use different reduction orders. + for name in class_names + [base_edge[1].hook_name]: + torch.testing.assert_close(observed[name], reference[name], atol=2e-4, rtol=0) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=0) + expected_recovery = (float(metric(expected)) - report.corrupt_metric) / ( + report.full_metric - report.corrupt_metric + ) + assert report.edge_class_recovered[edge_class] == pytest.approx(expected_recovery, abs=1e-5) + assert report.circuit_size == len(edges) - 1 + assert sum(report.edge_class_counts.values()) == len(edges) + if edge_class != "into_logits": + gap = report.full_metric - report.corrupt_metric + assert abs(float(metric(expected) - metric(restored)) / gap) > 1e-4 diff --git a/tests/unit/tools/test_attribution_patching.py b/tests/unit/tools/test_attribution_patching.py index 6b201c1a7..bbd5bfd7b 100644 --- a/tests/unit/tools/test_attribution_patching.py +++ b/tests/unit/tools/test_attribution_patching.py @@ -1,9 +1,7 @@ """Model-free unit tests for the attribution-patching substrate. -These tests target the ``TransformerBridge`` API exclusively (TransformerLens v4 -deprecates ``HookedTransformer``). They use a tiny, deliberately *linear* -``TransformerBridge`` subclass so gradients have a closed form and, later, the -first-order attribution identity holds exactly. +These tests target the ``TransformerBridge`` API using tiny linear subclasses +with closed-form gradients and exact first-order attribution identities. """ from __future__ import annotations @@ -22,21 +20,34 @@ from transformer_lens.tools.analysis.attribution_patching import ( AttributionResult, EdgeAttributionConfig, + EdgeClass, + FaithfulnessConfig, + FaithfulnessResult, GradientCache, Node, + NodeKind, + _ablate_edges, _assert_edges_unique, _check_required_hooks, + _edge_class, _edge_effects, _edge_hook_flags, _edge_hook_names, _ensure_edge_hook_flags, + _excluded_writers_by_reader, + _mean_writer_contributions, _node_effects, + _normalize_circuit, + _reader_hook_names, _required_hook_names, _writer_hook_name, + _writer_hook_names, attribution_patch, cache_activation_and_gradient, + capture_writer_outputs, enumerate_edges, enumerate_nodes, + faithfulness, ) D_MODEL = 4 @@ -227,7 +238,7 @@ def test_gradient_cache_activation_only_skips_gradients() -> None: # --------------------------------------------------------------------------- -# Commit 2 — typed computational-graph node model +# Typed computational-graph node model # --------------------------------------------------------------------------- N_HEADS = 2 @@ -311,7 +322,7 @@ def test_enumerate_nodes_raises_on_missing_hook() -> None: # --------------------------------------------------------------------------- -# Commit 3 — config + result API +# Configuration and result API # --------------------------------------------------------------------------- @@ -383,7 +394,7 @@ def test_top_edges_empty_when_no_edge_sweep_has_run() -> None: # --------------------------------------------------------------------------- -# Commit 4 — node attribution_patch entry point +# Node attribution entry point # --------------------------------------------------------------------------- @@ -391,7 +402,7 @@ class _AttnMlpBlock(nn.Module): """A block exposing the standard node-granularity hook points. ``hook_z`` carries the per-head output ``[batch, seq, n_heads, d_head]`` and - ``hook_mlp_out`` the MLP write ``[batch, seq, d_model]`` — the two per-layer + ``hook_mlp_out`` the MLP write ``[batch, seq, d_model]`` -- the two per-layer hook families ``enumerate_nodes`` reads. The block is linear; the test only needs a real hook graph with gradients flowing to those points, not attention. """ @@ -609,10 +620,10 @@ def test_attribution_patch_identical_inputs_are_exactly_zero_in_eval_mode() -> N # --------------------------------------------------------------------------- -# Commit 5 — linear-model reconstruction identity +# Linear-model reconstruction identity # --------------------------------------------------------------------------- # -# ``_NodeGraphToyBridge`` is fully linear by construction — the remaining +# ``_NodeGraphToyBridge`` is fully linear by construction -- the remaining # nonlinearities on the residual->metric path are neutralized without extra # freezing: ``ln_final`` is ``nn.Identity``, the # attention projection is a plain linear map with no softmax pattern, the MLP has @@ -662,8 +673,8 @@ def test_linear_reconstruction_identity_holds_on_a_complete_cut() -> None: post-attention residual, which already contains the embed and attention writes), so a node's full-model gradient re-counts the paths of every node upstream of it: summing embed + attention + MLP scores overcounts. The embed layer is the - input-side complete cut — every path to the metric passes through exactly one - embed position — so its scores alone reconstruct the exact metric delta for a + input-side complete cut -- every path to the metric passes through exactly one + embed position -- so its scores alone reconstruct the exact metric delta for a linear model. """ model = _NodeGraphToyBridge() @@ -774,7 +785,7 @@ def test_nonlinear_node_scores_read_the_corrupt_run_gradient() -> None: With an input-dependent Jacobian the clean and corrupt caches hold *different* gradients, so pairing the clean-activation delta with the clean gradient (the reversed convention) yields different scores from the corrupt-gradient one the - docstring promises. This pins ``attribution_patch`` to the corrupt gradient — a + docstring promises. This pins ``attribution_patch`` to the corrupt gradient -- a guard the linear tests structurally cannot provide. """ model = _NonlinearNodeGraphToyBridge() @@ -1295,7 +1306,7 @@ def test_enumerate_edges_drops_same_layer_heads_from_mlp_reader_on_parallel_attn # --------------------------------------------------------------------------- -# Commit 3 - edge scoring in attribution_patch +# Edge attribution scoring # --------------------------------------------------------------------------- # # Edge scoring needs every reader hook to actually participate in the forward @@ -1376,7 +1387,7 @@ def forward(self, residual: torch.Tensor) -> torch.Tensor: ) w_o_per_head = self.w_o.weight.reshape(d_model, self.n_heads, self.d_head).permute(1, 2, 0) - per_head_out_raw = torch.einsum("bshd,hdm->bshm", z, w_o_per_head) + per_head_out_raw = self._attention_output(torch.einsum("bshd,hdm->bshm", z, w_o_per_head)) if self.cfg.use_attn_result: per_head_out = self.hook_result(per_head_out_raw) else: @@ -1387,6 +1398,28 @@ def forward(self, residual: torch.Tensor) -> torch.Tensor: mlp_out = self.hook_mlp_out(self.w_mlp(mlp_in)) return self.hook_resid_post(residual + mlp_out) + def _attention_output(self, per_head_out: torch.Tensor) -> torch.Tensor: + """Transform the per-head attention output before it joins the residual. + + Identity here. A subclass can make the block nonlinear, which is what + stops a circuit from reproducing the metric merely by containing every + edge into the logits reader. + """ + return per_head_out + + +class _NonlinearEdgeScoringBlock(_EdgeScoringBlock): + """``_EdgeScoringBlock`` with a GELU on the attention output path. + + The nonlinearity breaks the linear relation between a writer's contribution + and the metric, so a circuit cannot recover the metric just by containing + every edge into the logits reader. That is what makes a random edge set of + the same size a meaningful baseline rather than a coin flip. + """ + + def _attention_output(self, per_head_out: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(per_head_out) + class _EdgeScoringToyBridge(_LinearToyBridge): """A tiny ``TransformerBridge`` whose edge-granularity hooks all sit on the data path. @@ -1466,6 +1499,30 @@ def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None: self.cfg.use_hook_mlp_in = use_hook_mlp_in +class _NonlinearEdgeScoringToyBridge(_EdgeScoringToyBridge): + """``_EdgeScoringToyBridge`` whose attention output passes through a GELU. + + Reuses the parent's hook graph and ``hooks()`` plumbing but swaps in + :class:`_NonlinearEdgeScoringBlock`, so a writer's contribution reaches the + metric nonlinearly. On a linear toy any circuit holding every edge into the + logits reader recovers the metric outright, which makes a random baseline + meaningless; the nonlinearity is what gives the comparison teeth. + """ + + def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: + super().__init__(dtype=dtype) + torch.manual_seed(1) + self.blocks = nn.ModuleList( + [ + _NonlinearEdgeScoringBlock(D_MODEL, N_HEADS, D_HEAD, layer, dtype, self.cfg) + for layer in range(N_LAYERS) + ] + ) + # The parent's eval() ran before these blocks existed, so they would + # otherwise start in training mode and trip require_eval_mode. + self.eval() + + def test_attribution_patch_edge_granularity_scores_every_edge_with_finite_values() -> None: model = _EdgeScoringToyBridge() clean = torch.tensor([[1, 2, 3]]) @@ -1647,10 +1704,10 @@ def test_attribution_patch_edge_sweep_restores_caller_hook_flags() -> None: # --------------------------------------------------------------------------- -# Commit 4 - exact-patch parity and mutation-checked reconstruction +# Exact-patch parity and mutation-checked reconstruction # --------------------------------------------------------------------------- # -# Two guards on the edge-scoring identity Risk 1 warns about: a genuine +# Two guards on the edge-scoring identity: a genuine # activation patch of a single edge must match _edge_effects' estimate (sign # always; magnitude too on this linear toy model), and mutating one writer's # captured contribution must change only the edges that writer feeds, never a @@ -1719,6 +1776,45 @@ def test_edge_scores_sum_to_the_readers_direct_input_change() -> None: assert checked > 0 +def _assert_mutation_only_changes_the_perturbed_writers_edges( + edges: list[tuple[Node, Node]], + writer: Node, + perturbation: torch.Tensor, + baseline_scores: dict[tuple[Node, Node], float], + mutated_scores: dict[tuple[Node, Node], float], + corrupt_cache: GradientCache, +) -> None: + """Assert a single-writer perturbation moved exactly that writer's edges. + + The scorer indexes writer deltas and reader gradients independently. + Perturbing one head at one position must not change sibling writers' scores. + """ + changed = 0 + moved = 0 + for edge in edges: + edge_writer, reader = edge + if edge_writer == writer: + grad = corrupt_cache.gradients[reader.hook_name] + assert grad is not None + if reader.kind in ("q_input", "k_input", "v_input"): + grad_vec = grad[0, reader.position, reader.head] + else: + grad_vec = grad[0, reader.position] + shift = float((perturbation * grad_vec).sum()) + expected = baseline_scores[edge] + shift + # Cancellation in float32 needs an absolute tolerance. + assert mutated_scores[edge] == pytest.approx(expected, abs=1e-6) + changed += 1 + if abs(shift) > 1e-6: + assert mutated_scores[edge] != pytest.approx(baseline_scores[edge]) + moved += 1 + else: + assert mutated_scores[edge] == baseline_scores[edge] + + assert changed > 0 + assert moved > 0 # the perturbation genuinely shifts scores at the readout position + + def test_edge_effects_mutation_only_changes_the_perturbed_writers_edges() -> None: """Perturbing one writer's captured contribution changes only that writer's edges. @@ -1768,30 +1864,9 @@ def test_edge_effects_mutation_only_changes_the_perturbed_writers_edges() -> Non mutated_scores = _edge_effects(perturbed_clean_cache, corrupt_cache, edges) - changed = 0 - moved = 0 - for edge in edges: - edge_writer, reader = edge - if edge_writer == writer: - grad = corrupt_cache.gradients[reader.hook_name] - assert grad is not None - if reader.kind in ("q_input", "k_input", "v_input"): - grad_vec = grad[0, reader.position, reader.head] - else: - grad_vec = grad[0, reader.position] - shift = float((perturbation * grad_vec).sum()) - expected = baseline_scores[edge] + shift - # float32: baseline and shift partly cancel, so the default rel=1e-6 is too tight - assert mutated_scores[edge] == pytest.approx(expected, abs=1e-6) - changed += 1 - if abs(shift) > 1e-6: - assert mutated_scores[edge] != pytest.approx(baseline_scores[edge]) - moved += 1 - else: - assert mutated_scores[edge] == baseline_scores[edge] - - assert changed > 0 - assert moved > 0 # the perturbation genuinely shifts scores at the readout position + _assert_mutation_only_changes_the_perturbed_writers_edges( + edges, writer, perturbation, baseline_scores, mutated_scores, corrupt_cache + ) # --------------------------------------------------------------------------- @@ -1956,3 +2031,1259 @@ def test_top_edges_can_surface_a_head_to_logits_edge() -> None: # At least one carries a real signed effect, so the ranking is not over zeros. assert any(abs(score) > 1e-6 for _writer, _reader, score in head_to_logits) + + +# --------------------------------------------------------------------------- +# Writer-output capture (forward-only) +# --------------------------------------------------------------------------- +# +# Rewriting a reader's input during a live forward pass needs the writers' +# current contributions rather than a gradient estimate. The capture is +# therefore forward-only and names-filtered to the writer families, the +# opposite contract from cache_activation_and_gradient. + + +def test_writer_hook_names_cover_the_three_writer_families() -> None: + names = _writer_hook_names(N_LAYERS) + + assert names[0] == "hook_embed" + for layer in range(N_LAYERS): + assert f"blocks.{layer}.attn.hook_result" in names + assert f"blocks.{layer}.hook_mlp_out" in names + # Writer families only: no reader inputs, and no pre-hook_result attn-z. + assert not any("hook_q_input" in name for name in names) + assert not any("hook_mlp_in" in name for name in names) + assert not any(name.endswith("attn.hook_z") for name in names) + + +def test_reader_hook_names_cover_the_reader_input_families() -> None: + names = _reader_hook_names(N_LAYERS) + + for layer in range(N_LAYERS): + for input_hook in ("hook_q_input", "hook_k_input", "hook_v_input"): + assert f"blocks.{layer}.attn.{input_hook}" in names + assert f"blocks.{layer}.hook_mlp_in" in names + # Reader inputs only: no writer outputs. + assert not any("hook_result" in name for name in names) + assert not any("hook_mlp_out" in name for name in names) + + +def test_capture_writer_outputs_populates_only_the_writer_family() -> None: + model = _EdgeScoringToyBridge() + tokens = model.to_tokens("prompt") + + with _edge_hook_flags(model): + captured = capture_writer_outputs(model, tokens) + + assert set(captured) == set(_writer_hook_names(N_LAYERS)) + # The reader-input family and the pre-hook_result attn-z are not captured. + assert "blocks.0.attn.hook_q_input" not in captured + assert "blocks.0.hook_mlp_in" not in captured + assert "blocks.0.attn.hook_z" not in captured + + assert captured["hook_embed"].shape == (1, SEQ_LEN, D_MODEL) + for layer in range(N_LAYERS): + assert captured[f"blocks.{layer}.attn.hook_result"].shape == ( + 1, + SEQ_LEN, + N_HEADS, + D_MODEL, + ) + assert captured[f"blocks.{layer}.hook_mlp_out"].shape == ( + 1, + SEQ_LEN, + D_MODEL, + ) + + +def test_capture_writer_outputs_returns_detached_tensors() -> None: + model = _EdgeScoringToyBridge() + tokens = model.to_tokens("prompt") + + with _edge_hook_flags(model): + captured = capture_writer_outputs(model, tokens) + + for tensor in captured.values(): + assert not tensor.requires_grad + assert tensor.grad_fn is None + + +def test_capture_writer_outputs_raises_when_the_filter_matches_nothing() -> None: + model = _EdgeScoringToyBridge() + tokens = model.to_tokens("prompt") + + with pytest.raises(ValueError, match="matched no hook points"): + capture_writer_outputs(model, tokens, names_filter=["blocks.0.not_a_hook"]) + + +def test_capture_writer_outputs_needs_no_grad() -> None: + """The capture runs with autograd off, where the gradient helper refuses. + + The two helpers have opposite contracts: this one is forward-only and must + work under ``torch.no_grad()``, while ``cache_activation_and_gradient`` + exists to retain gradients and raises without autograd. + """ + model = _EdgeScoringToyBridge() + tokens = model.to_tokens("prompt") + metric = _metric_fn(answer=1, wrong=2) + + with _edge_hook_flags(model), torch.no_grad(): + captured = capture_writer_outputs(model, tokens) + assert set(captured) == set(_writer_hook_names(N_LAYERS)) + + with pytest.raises(ValueError, match="autograd"): + cache_activation_and_gradient(model, tokens, metric) + + +# --------------------------------------------------------------------------- +# Reader-input rewrite (edge ablation core) +# --------------------------------------------------------------------------- +# +# The residual stream is a running sum, so a reader's input is exactly the sum +# of its writers' contributions. Ablating an edge means subtracting that +# writer's live contribution and adding its replacement at the reader's fork +# hook, which fires before the layer norm -- an exact correction, not an +# approximation. Partial circuits need independent reader and logit references; +# boundary circuits alone cannot expose incorrect upstream corrections. + + +def _ablation_replacement( + model: _EdgeScoringToyBridge, tokens: torch.Tensor +) -> dict[str, torch.Tensor]: + """The value standing in for an ablated writer: the corrupt run's own.""" + with _edge_hook_flags(model): + return capture_writer_outputs(model, tokens) + + +def test_excluded_writers_by_reader_groups_only_out_of_circuit_edges() -> None: + model = _EdgeScoringToyBridge() + _ensure_edge_hook_flags(model) + tokens = model.to_tokens("prompt") + metric = _metric_fn(answer=1, wrong=2) + edge_hook_names = _edge_hook_names(N_LAYERS) + + corrupt_cache = cache_activation_and_gradient( + model, tokens, metric, names_filter=edge_hook_names + ) + edges = enumerate_edges(model, corrupt_cache) + + # Keeping the whole graph excludes nothing. + assert _excluded_writers_by_reader(edges, edges) == {} + + # Keeping nothing excludes every writer, grouped once per reader. + excluded = _excluded_writers_by_reader(edges, []) + readers = {reader for _writer, reader in edges} + assert set(excluded) == readers + for reader, writers in excluded.items(): + expected = [writer for writer, edge_reader in edges if edge_reader == reader] + assert writers == expected + + # A single kept edge excludes exactly that edge's writer from its reader. + # Pick a reader with more than one incoming edge, so the exclusion is + # observable rather than emptying the reader's writer list. + kept_edge = next( + edge for edge in edges if len([writer for writer, reader in edges if reader == edge[1]]) > 1 + ) + excluded = _excluded_writers_by_reader(edges, [kept_edge]) + kept_writer, kept_reader = kept_edge + incoming = [writer for writer, reader in edges if reader == kept_reader] + assert kept_writer not in excluded[kept_reader] + assert len(excluded[kept_reader]) == len(incoming) - 1 + + +def test_ablating_the_full_circuit_reproduces_the_clean_metric() -> None: + """Keeping every edge leaves each reader's input untouched.""" + model = _EdgeScoringToyBridge() + _ensure_edge_hook_flags(model) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + metric = _metric_fn(answer=1, wrong=2) + edge_hook_names = _edge_hook_names(N_LAYERS) + + corrupt_cache = cache_activation_and_gradient( + model, corrupt, metric, names_filter=edge_hook_names + ) + clean_cache = cache_activation_and_gradient( + model, clean, metric, names_filter=edge_hook_names, compute_gradient=False + ) + edges = enumerate_edges(model, corrupt_cache) + replacements = _ablation_replacement(model, corrupt) + + with _edge_hook_flags(model): + ablated_logits = _ablate_edges(model, clean, edges, edges, replacements) + + ablated = float(metric(ablated_logits)) + assert ablated == pytest.approx(float(clean_cache.metric), abs=1e-6) + # Discriminating: the clean and corrupt runs genuinely differ, so this + # cannot pass on a degenerate pair. + assert ablated != pytest.approx(float(corrupt_cache.metric), abs=1e-6) + + +def test_ablation_installs_one_hook_per_reader_node() -> None: + """Each reader node is corrected once, not once per incoming edge. + + Several reader nodes share one hook point (one per head, one per position), + and each rewrites a disjoint slice of it. Installing one hook per excluded + edge would apply the correction repeatedly and overshoot, so the count of + hooks on each fork point must equal the number of reader nodes reading it -- + not the number of edges into those nodes. + """ + model = _EdgeScoringToyBridge() + _ensure_edge_hook_flags(model) + tokens = model.to_tokens("prompt") + metric = _metric_fn(answer=1, wrong=2) + edge_hook_names = _edge_hook_names(N_LAYERS) + + corrupt_cache = cache_activation_and_gradient( + model, tokens, metric, names_filter=edge_hook_names + ) + edges = enumerate_edges(model, corrupt_cache) + replacements = _ablation_replacement(model, tokens) + + readers = {reader for _writer, reader in edges} + assert len(readers) > 0 + # Some readers have several incoming edges, so a per-edge hook would + # register several hooks on the same point. + crowded = { + reader + for reader in readers + if len([writer for writer, edge_reader in edges if edge_reader == reader]) > 1 + } + assert crowded + + # The ablation's hooks live only for the duration of its forward, so count + # them from a pre-hook on the model root, which fires while they are + # registered. + counts: dict[str, int] = {} + + def snapshot(_module: nn.Module, _args: tuple) -> None: + for reader in readers: + counts[reader.hook_name] = len(model.hook_dict[reader.hook_name].fwd_hooks) + + handle = model.register_forward_pre_hook(snapshot) + try: + with _edge_hook_flags(model): + _ablate_edges(model, tokens, edges, [], replacements) + finally: + handle.remove() + + expected = { + name: len([reader for reader in readers if reader.hook_name == name]) + for name in {reader.hook_name for reader in readers} + } + assert counts == expected + # Fewer hooks than edges: the grouping collapsed the per-edge corrections. + assert sum(counts.values()) < len(edges) + + +def _toy_ablation_edges() -> list[tuple[Node, Node]]: + """Explicit topology for the two-layer, two-head sequential toy.""" + edges: list[tuple[Node, Node]] = [] + for position in range(SEQ_LEN): + writers = [ + Node("embed", position), + Node("attn_head_out", position, layer=0, head=0), + Node("attn_head_out", position, layer=0, head=1), + Node("mlp_out", position, layer=0), + Node("attn_head_out", position, layer=1, head=0), + Node("attn_head_out", position, layer=1, head=1), + Node("mlp_out", position, layer=1), + ] + kinds: tuple[NodeKind, ...] = ("q_input", "k_input", "v_input") + for layer, upstream, post_attention in ((0, 1, 3), (1, 4, 6)): + for kind in kinds: + for head in range(N_HEADS): + reader = Node(kind, position, layer=layer, head=head) + edges.extend((writer, reader) for writer in writers[:upstream]) + reader = Node("mlp_in", position, layer=layer) + edges.extend((writer, reader) for writer in writers[:post_attention]) + reader = Node("logits", position, layer=1) + edges.extend((writer, reader) for writer in writers) + return edges + + +@torch.no_grad() +def _functional_toy_ablation( + model: _EdgeScoringToyBridge, + tokens: torch.Tensor, + excluded: set[tuple[Node, Node]], + replacements: dict[Node, torch.Tensor], +) -> tuple[torch.Tensor, dict[Node, torch.Tensor], dict[Node, torch.Tensor]]: + """Evaluate the linear toy by summing edge values, without invoking hooks.""" + writers: dict[Node, torch.Tensor] = {} + readers: dict[Node, torch.Tensor] = {} + + def read(reader: Node, available: list[Node]) -> torch.Tensor: + value = torch.stack( + [ + replacements[writer] if (writer, reader) in excluded else writers[writer] + for writer in available + ] + ).sum(dim=0) + readers[reader] = value + return value + + outputs = [] + for position in range(tokens.shape[1]): + embed = Node("embed", position) + writers[embed] = model.embed.weight[tokens[0, position]] + available = [embed] + for layer, block in enumerate(model.blocks): + heads = [] + for head in range(N_HEADS): + z = torch.zeros(D_HEAD, dtype=model.cfg.dtype) + projections: tuple[tuple[NodeKind, torch.Tensor], ...] = ( + ("q_input", block.w_q.weight), + ("k_input", block.w_k.weight), + ("v_input", block.w_v.weight), + ) + for kind, weight in projections: + reader = Node(kind, position, layer=layer, head=head) + z += weight[head * D_HEAD : (head + 1) * D_HEAD] @ read(reader, available) + writer = Node("attn_head_out", position, layer=layer, head=head) + writers[writer] = block.w_o.weight[:, head * D_HEAD : (head + 1) * D_HEAD] @ z + heads.append(writer) + available.extend(heads) + reader = Node("mlp_in", position, layer=layer) + writer = Node("mlp_out", position, layer=layer) + writers[writer] = block.w_mlp.weight @ read(reader, available) + available.append(writer) + final = read(Node("logits", position, layer=1), available) + outputs.append(model.unembed.weight @ final) + return torch.stack(outputs).unsqueeze(0), writers, readers + + +def _toy_replacement_tensors(writers: dict[Node, torch.Tensor]) -> dict[str, torch.Tensor]: + tensors = {"hook_embed": torch.stack([writers[Node("embed", p)] for p in range(SEQ_LEN)])[None]} + for layer in range(N_LAYERS): + tensors[f"blocks.{layer}.attn.hook_result"] = torch.stack( + [ + torch.stack( + [writers[Node("attn_head_out", p, layer=layer, head=h)] for h in range(N_HEADS)] + ) + for p in range(SEQ_LEN) + ] + )[None] + tensors[f"blocks.{layer}.hook_mlp_out"] = torch.stack( + [writers[Node("mlp_out", p, layer=layer)] for p in range(SEQ_LEN)] + )[None] + return tensors + + +def _observed_toy_ablation( + model: _EdgeScoringToyBridge, + clean: torch.Tensor, + circuit: list[tuple[Node, Node]], + replacements: dict[Node, torch.Tensor], + monkeypatch: pytest.MonkeyPatch, +) -> tuple[torch.Tensor, dict[Node, torch.Tensor]]: + observed: dict[Node, torch.Tensor] = {} + original_hooks = model.hooks + edges = _toy_ablation_edges() + + def observer(reader: Node) -> Callable[..., None]: + def capture(tensor: torch.Tensor, *, hook: Any) -> None: + del hook + if reader.head is None: + observed[reader] = tensor[0, reader.position].clone() + else: + observed[reader] = tensor[0, reader.position, reader.head].clone() + + return capture + + @contextmanager + def hooks(fwd_hooks: list[tuple[str, Any]], **kwargs: Any) -> Iterator[Any]: + # Observers run after all corrections sharing the reader's hook point. + observers = [ + (reader.hook_name, observer(reader)) + for reader in dict.fromkeys(reader for _, reader in edges) + ] + with original_hooks(fwd_hooks=fwd_hooks + observers, **kwargs): + yield model + + with monkeypatch.context() as patch, _edge_hook_flags(model): + patch.setattr(model, "hooks", hooks) + logits = _ablate_edges(model, clean, edges, circuit, _toy_replacement_tensors(replacements)) + return logits, observed + + +@pytest.mark.parametrize("kind", ["q_input", "k_input", "v_input", "mlp_in", "logits"]) +@pytest.mark.parametrize("head,position", [(0, 0), (1, 2)]) +def test_partial_ablation_matches_functional_toy_reference( + kind: NodeKind, head: int, position: int, monkeypatch: pytest.MonkeyPatch +) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + edges = _toy_ablation_edges() + assert edges == enumerate_edges(model, _synthetic_edge_cache()) + clean_logits, clean_writers, clean_readers = _functional_toy_ablation(model, clean, set(), {}) + corrupt_logits, replacements, _ = _functional_toy_ablation(model, corrupt, set(), {}) + with torch.no_grad(): + torch.testing.assert_close(clean_logits, model(clean)) + torch.testing.assert_close(corrupt_logits, model(corrupt)) + writer = Node("attn_head_out", position, layer=0, head=head) + reader = Node(kind, position, layer=1, head=head if kind.endswith("input") else None) + excluded = {(writer, reader)} + assert excluded <= set(edges) + assert not torch.allclose(clean_writers[writer], replacements[writer]) + circuit = [edge for edge in edges if edge not in excluded] + expected, _, expected_readers = _functional_toy_ablation(model, clean, excluded, replacements) + actual, observed = _observed_toy_ablation(model, clean, circuit, replacements, monkeypatch) + torch.testing.assert_close(actual, expected, atol=1e-12, rtol=1e-12) + assert not torch.allclose(actual, clean_logits) + assert not torch.allclose(observed[reader], clean_readers[reader]) + assert observed.keys() == expected_readers.keys() + for node, value in expected_readers.items(): + torch.testing.assert_close(observed[node], value, atol=1e-12, rtol=1e-12) + + +def test_sparse_partial_ablation_matches_functional_toy_reference( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + edges = _toy_ablation_edges() + clean_logits, _, _ = _functional_toy_ablation(model, clean, set(), {}) + corrupt_logits, replacements, _ = _functional_toy_ablation(model, corrupt, set(), {}) + circuit = edges[::2] + excluded = set(edges) - set(circuit) + expected, _, readers = _functional_toy_ablation(model, clean, excluded, replacements) + actual, observed = _observed_toy_ablation(model, clean, circuit, replacements, monkeypatch) + torch.testing.assert_close(actual, expected, atol=1e-12, rtol=1e-12) + assert not torch.allclose(actual, clean_logits) + assert not torch.allclose(actual, corrupt_logits) + for reader, value in readers.items(): + torch.testing.assert_close(observed[reader], value, atol=1e-12, rtol=1e-12) + + +def test_coupled_partial_ablation_uses_intervened_live_writers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + edges = _toy_ablation_edges() + clean_logits, clean_writers, clean_readers = _functional_toy_ablation(model, clean, set(), {}) + _, replacements, _ = _functional_toy_ablation(model, corrupt, set(), {}) + embed = Node("embed", 2) + writer = Node("attn_head_out", 2, layer=0, head=0) + downstream = Node("v_input", 2, layer=1, head=1) + excluded = { + (embed, Node("q_input", 2, layer=0, head=0)), + (writer, Node("k_input", 2, layer=1, head=0)), + } + circuit = [edge for edge in edges if edge not in excluded] + assert (writer, downstream) in circuit + expected, live_writers, readers = _functional_toy_ablation(model, clean, excluded, replacements) + actual, observed = _observed_toy_ablation(model, clean, circuit, replacements, monkeypatch) + shift = live_writers[writer] - clean_writers[writer] + assert not torch.allclose(shift, torch.zeros_like(shift)) + mlp = Node("mlp_out", 2, layer=0) + mlp_shift = live_writers[mlp] - clean_writers[mlp] + torch.testing.assert_close(readers[downstream], clean_readers[downstream] + shift + mlp_shift) + torch.testing.assert_close(actual, expected, atol=1e-12, rtol=1e-12) + assert not torch.allclose(actual, clean_logits) + for reader, value in readers.items(): + torch.testing.assert_close(observed[reader], value, atol=1e-12, rtol=1e-12) + + +def test_ablation_replacement_mutation_isolates_direct_reader_slices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + edges = _toy_ablation_edges() + _, replacements, _ = _functional_toy_ablation(model, corrupt, set(), {}) + writer = Node("attn_head_out", 2, layer=0, head=1) + excluded_readers = { + Node("q_input", 2, layer=1, head=0), + Node("k_input", 2, layer=1, head=1), + Node("v_input", 2, layer=1, head=0), + } + excluded = {(writer, reader) for reader in excluded_readers} + circuit = [edge for edge in edges if edge not in excluded] + perturbation = torch.tensor([0.37, -0.19, 0.23, -0.41], dtype=torch.float64) + mutated = dict(replacements) + mutated[writer] = replacements[writer] + perturbation + _, baseline = _observed_toy_ablation(model, clean, circuit, replacements, monkeypatch) + logits, observed = _observed_toy_ablation(model, clean, circuit, mutated, monkeypatch) + expected, _, readers = _functional_toy_ablation(model, clean, excluded, mutated) + torch.testing.assert_close(logits, expected, atol=1e-12, rtol=1e-12) + for reader, value in readers.items(): + torch.testing.assert_close(observed[reader], value, atol=1e-12, rtol=1e-12) + # Later MLP and logits readers can change through causal propagation. + if reader.kind in ("q_input", "k_input", "v_input"): + shift = perturbation if reader in excluded_readers else torch.zeros_like(perturbation) + torch.testing.assert_close( + observed[reader] - baseline[reader], shift, atol=1e-12, rtol=1e-12 + ) + + +def test_ablation_raises_when_a_writer_contribution_was_not_captured() -> None: + """An uncaptured writer raises rather than silently skipping the correction.""" + model = _EdgeScoringToyBridge() + _ensure_edge_hook_flags(model) + tokens = model.to_tokens("prompt") + metric = _metric_fn(answer=1, wrong=2) + edge_hook_names = _edge_hook_names(N_LAYERS) + + corrupt_cache = cache_activation_and_gradient( + model, tokens, metric, names_filter=edge_hook_names + ) + edges = enumerate_edges(model, corrupt_cache) + + with _edge_hook_flags(model), pytest.raises(ValueError, match="no contribution captured"): + _ablate_edges(model, tokens, edges, [], {}) + + +# --------------------------------------------------------------------------- +# Faithfulness (public API) +# --------------------------------------------------------------------------- +# +# faithfulness() ablates every edge outside a candidate circuit and reports how +# much of the clean-to-corrupt metric gap the circuit recovers. The two boundary +# circuits pin it: keeping nothing must reproduce the corrupt run, keeping +# everything must reproduce the clean run. + + +def _faithfulness_toy() -> tuple[_EdgeScoringToyBridge, torch.Tensor, torch.Tensor, Callable]: + model = _EdgeScoringToyBridge() + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + return model, clean, corrupt, _metric_fn(answer=1, wrong=2) + + +def _graph_edges(model: _EdgeScoringToyBridge, corrupt: torch.Tensor, metric: Callable) -> list: + _ensure_edge_hook_flags(model) + with _edge_hook_flags(model): + corrupt_cache = cache_activation_and_gradient( + model, corrupt, metric, names_filter=_edge_hook_names(N_LAYERS) + ) + return enumerate_edges(model, corrupt_cache) + + +def test_faithfulness_recovers_the_full_metric_for_the_full_circuit() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + edges = _graph_edges(model, corrupt, metric) + + report = faithfulness(model, clean, corrupt, metric, edges) + + assert isinstance(report, FaithfulnessResult) + assert report.recovered == pytest.approx(1.0, abs=1e-6) + assert report.circuit_size == len(edges) + assert report.total_edges == len(edges) + + +def test_faithfulness_recovers_nothing_for_the_empty_circuit() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + + report = faithfulness(model, clean, corrupt, metric, []) + + assert report.recovered == pytest.approx(0.0, abs=1e-6) + assert report.circuit_size == 0 + assert report.total_edges > 0 + + +def test_faithfulness_full_and_corrupt_metrics_match_direct_evaluation() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + + report = faithfulness(model, clean, corrupt, metric, []) + + with torch.no_grad(): + assert report.full_metric == pytest.approx(float(metric(model(clean))), abs=1e-6) + assert report.corrupt_metric == pytest.approx(float(metric(model(corrupt))), abs=1e-6) + # The pair genuinely separates, so the recovered fraction is well defined. + assert report.full_metric != pytest.approx(report.corrupt_metric) + + +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_faithfulness_is_forward_only_and_preserves_parameter_gradients( + config: FaithfulnessConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + circuit = _toy_ablation_edges()[::2] + mean_tokens = _toy_mean_calibration() if config.ablation == "mean" else None + baseline = faithfulness(model, clean, corrupt, metric, circuit, config, mean_tokens=mean_tokens) + parameters = list(model.parameters()) + for index, parameter in enumerate(parameters): + parameter.grad = torch.full_like(parameter, 0.17) if index % 2 else None + saved = [None if parameter.grad is None else parameter.grad.clone() for parameter in parameters] + + def forbidden_backward(*args: Any, **kwargs: Any) -> None: + raise AssertionError("faithfulness must not request a backward pass") + + def forward_metric(logits: torch.Tensor) -> torch.Tensor: + assert not logits.requires_grad + return metric(logits) + + monkeypatch.setattr(torch.autograd, "grad", forbidden_backward) + monkeypatch.setattr(torch.Tensor, "backward", forbidden_backward) + for enabled in (False, True): + with torch.set_grad_enabled(enabled): + report = faithfulness( + model, clean, corrupt, forward_metric, circuit, config, mean_tokens=mean_tokens + ) + assert report == baseline + for parameter, previous in zip(parameters, saved): + if previous is None: + assert parameter.grad is None + else: + torch.testing.assert_close(parameter.grad, previous, atol=0, rtol=0) + + +def test_faithfulness_defaults_to_corrupt_ablation() -> None: + assert FaithfulnessConfig().ablation == "corrupt" + + +def test_faithfulness_mean_ablation_requires_calibration() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + with pytest.raises(ValueError, match="mean_tokens"): + faithfulness(model, clean, corrupt, metric, [], FaithfulnessConfig(ablation="mean")) + + +def _toy_mean_calibration() -> torch.Tensor: + return torch.tensor([[0, 4, 5], [2, 5, 0], [4, 0, 2]]) + + +def _functional_mean_writers( + model: _EdgeScoringToyBridge, calibration: torch.Tensor +) -> dict[Node, torch.Tensor]: + examples = [_functional_toy_ablation(model, row[None], set(), {})[1] for row in calibration] + return { + writer: torch.stack([example[writer] for example in examples]).mean(dim=0) + for writer in examples[0] + } + + +def test_mean_writer_contributions_preserve_positions_and_heads() -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + calibration = _toy_mean_calibration() + expected = _toy_replacement_tensors(_functional_mean_writers(model, calibration)) + first = _toy_replacement_tensors(_functional_toy_ablation(model, calibration[:1], set(), {})[1]) + with _edge_hook_flags(model): + actual = _mean_writer_contributions(model, calibration) + assert actual.keys() == expected.keys() + for name, value in actual.items(): + assert value.shape == expected[name].shape + assert value.shape[0] == 1 + assert value.shape[1] == calibration.shape[1] + assert value.dtype == torch.float64 + assert value.device == calibration.device + assert not value.requires_grad + torch.testing.assert_close(value, expected[name], atol=1e-12, rtol=1e-12) + assert not torch.allclose(value, first[name]) + + +@pytest.mark.parametrize("budget", [0, 3, None]) +def test_faithfulness_mean_ablation_matches_functional_reference(budget: int | None) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + calibration = _toy_mean_calibration() + edges = _toy_ablation_edges() + readout = [edge for edge in edges if edge[1] == Node("logits", SEQ_LEN - 1, layer=1)] + circuit = edges if budget is None else readout[:budget] + replacements = _functional_mean_writers(model, calibration) + clean_logits, _, _ = _functional_toy_ablation(model, clean, set(), {}) + corrupt_logits, _, _ = _functional_toy_ablation(model, corrupt, set(), {}) + expected_logits, _, _ = _functional_toy_ablation( + model, clean, set(edges) - set(circuit), replacements + ) + metric = _metric_fn(1, 2) + expected = float( + (metric(expected_logits) - metric(corrupt_logits)) + / (metric(clean_logits) - metric(corrupt_logits)) + ) + report = faithfulness( + model, + clean, + corrupt, + metric, + circuit, + FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration, + ) + assert report.recovered == pytest.approx(expected, abs=1e-12) + assert report.full_metric == pytest.approx(float(metric(clean_logits)), abs=1e-12) + assert report.corrupt_metric == pytest.approx(float(metric(corrupt_logits)), abs=1e-12) + if budget is None: + assert report.recovered == pytest.approx(1.0, abs=1e-12) + elif budget == 0: + mean_final = sum( + value for writer, value in replacements.items() if writer.position == clean.shape[1] - 1 + ) + mean_logits = model.unembed.weight.detach() @ mean_final + torch.testing.assert_close(expected_logits[0, -1], mean_logits, atol=1e-12, rtol=1e-12) + else: + empty_logits, _, _ = _functional_toy_ablation(model, clean, set(edges), replacements) + assert not torch.allclose(expected_logits[0, -1], empty_logits[0, -1]) + assert not torch.allclose(expected_logits[0, -1], clean_logits[0, -1]) + + +def test_faithfulness_mean_ablation_differs_from_corrupt_ablation() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + edges = _toy_ablation_edges() + circuit = edges[::2] + calibration = _toy_mean_calibration() + corrupt_report = faithfulness(model, clean, corrupt, metric, circuit) + mean_report = faithfulness( + model, + clean, + corrupt, + metric, + circuit, + config=FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration, + ) + expected_logits, _, _ = _functional_toy_ablation( + model, clean, set(edges) - set(circuit), _functional_mean_writers(model, calibration) + ) + expected = (float(metric(expected_logits)) - corrupt_report.corrupt_metric) / ( + corrupt_report.full_metric - corrupt_report.corrupt_metric + ) + assert mean_report.recovered == pytest.approx(expected, abs=1e-6) + assert corrupt_report.recovered != pytest.approx(mean_report.recovered) + assert corrupt_report.full_metric == pytest.approx(mean_report.full_metric) + assert corrupt_report.corrupt_metric == pytest.approx(mean_report.corrupt_metric) + + +def test_mean_calibration_does_not_depend_on_evaluated_pair( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + calibration = _toy_mean_calibration() + captured: list[tuple[torch.Tensor, dict[str, torch.Tensor]]] = [] + + def capture(model: Any, tokens: torch.Tensor) -> dict[str, torch.Tensor]: + values = capture_writer_outputs(model, tokens) + captured.append((tokens.clone(), {name: value.clone() for name, value in values.items()})) + return values + + monkeypatch.setattr( + "transformer_lens.tools.analysis.attribution_patching.capture_writer_outputs", capture + ) + for clean, corrupt in ( + (clean, corrupt), + (torch.tensor([[0, 2, 5]]), torch.tensor([[5, 2, 0]])), + ): + faithfulness( + model, + clean, + corrupt, + metric, + _toy_ablation_edges()[::2], + FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration, + ) + assert len(captured) == 2 * len(calibration) + for index, row in enumerate(calibration): + first_tokens, first_values = captured[index] + second_tokens, second_values = captured[index + len(calibration)] + torch.testing.assert_close(first_tokens, row[None], atol=0, rtol=0) + torch.testing.assert_close(second_tokens, row[None], atol=0, rtol=0) + for name in first_values: + torch.testing.assert_close(first_values[name], second_values[name], atol=0, rtol=0) + + +@pytest.mark.parametrize( + "tokens,error", + [ + (torch.empty(0, SEQ_LEN, dtype=torch.long), "at least one"), + (torch.tensor([0, 4, 5]), "2D"), + (torch.tensor([[[0, 4, 5]]]), "2D"), + (torch.tensor([[0, 4]]), "same sequence length"), + (torch.tensor([[0.0, 4.0, 5.0]]), "dtype"), + (torch.tensor([[0, 4, 5]], dtype=torch.int16), "dtype"), + (torch.tensor([[True, False, True]]), "dtype"), + (torch.empty(1, SEQ_LEN, device="meta", dtype=torch.long), "same device"), + (torch.tensor([[0, 4, 5], [1, 2, 3]]), "evaluated.*endpoint"), + (torch.tensor([[0, 4, 5], [3, 2, 1]]), "evaluated.*endpoint"), + ], +) +def test_mean_calibration_rejects_invalid_tokens_without_forward( + tokens: torch.Tensor, + error: str, +) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + model.set_use_attn_result(True) + before = _mean_toy_hook_state(model) + + def forbidden_forward(*args: Any) -> None: + raise AssertionError("invalid calibration must be rejected before a forward") + + handle = model.register_forward_pre_hook(forbidden_forward) + try: + with pytest.raises(ValueError, match=error): + faithfulness( + model, + clean, + corrupt, + metric, + [], + FaithfulnessConfig(ablation="mean"), + mean_tokens=tokens, + ) + assert _mean_toy_hook_state(model) == before + finally: + handle.remove() + + +def test_corrupt_ablation_rejects_mean_calibration() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + with pytest.raises(ValueError, match="only accepted with mean ablation"): + faithfulness(model, clean, corrupt, metric, [], mean_tokens=_toy_mean_calibration()) + + +def test_mean_calibration_accepts_int32_token_ids() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + config = FaithfulnessConfig(ablation="mean") + first = faithfulness( + model, clean, corrupt, metric, [], config, mean_tokens=_toy_mean_calibration() + ) + second = faithfulness( + model, + clean, + corrupt, + metric, + [], + config, + mean_tokens=_toy_mean_calibration().to(torch.int32), + ) + assert first == second + + +def _mean_toy_hook_state(model: _EdgeScoringToyBridge) -> tuple: + return ( + tuple( + getattr(model.cfg, name) + for name in ("use_attn_result", "use_split_qkv_input", "use_hook_mlp_in") + ), + { + name: ( + tuple(id(handle) for handle in point.fwd_hooks), + tuple(id(handle) for handle in point.bwd_hooks), + ) + for name, point in model.hook_dict.items() + }, + ) + + +@pytest.mark.parametrize("failure", ["capture", "metric"]) +def test_mean_calibration_restores_hooks_after_error(failure: str) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + calibration = _toy_mean_calibration() + model.set_use_attn_result(True) + calls = 0 + + def capture_guard(tensor: torch.Tensor, *, hook: Any) -> None: + nonlocal calls + calls += 1 + if failure == "capture" and calls == 3: + raise RuntimeError("calibration capture failed") + + def checked_metric(logits: torch.Tensor) -> torch.Tensor: + if failure == "metric" and calls == 6: + raise RuntimeError("mean metric failed") + return metric(logits) + + with model.hooks(fwd_hooks=[("hook_embed", capture_guard)]): + before = _mean_toy_hook_state(model) + with pytest.raises(RuntimeError, match="failed"): + faithfulness( + model, + clean, + corrupt, + checked_metric, + _toy_ablation_edges()[::2], + FaithfulnessConfig(ablation="mean"), + mean_tokens=calibration, + ) + assert _mean_toy_hook_state(model) == before + + +def test_faithfulness_accepts_top_edges_output() -> None: + """A ranked edge list feeds straight in, scores and all.""" + model, clean, corrupt, metric = _faithfulness_toy() + ranked = attribution_patch( + model, clean, corrupt, metric, config=EdgeAttributionConfig(granularity="edge") + ).top_edges(k=5) + + report = faithfulness(model, clean, corrupt, metric, ranked) + + assert report.circuit_size == len(ranked) + assert math.isfinite(report.recovered) + + +def test_faithfulness_raises_on_an_edge_outside_the_graph() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + # The graph never mixes positions, so a cross-position pair is not an edge. + bogus = (Node(kind="embed", position=1), Node(kind="mlp_in", layer=0, position=0)) + + with pytest.raises(ValueError, match="not in the graph"): + faithfulness(model, clean, corrupt, metric, [bogus]) + + +def test_faithfulness_raises_when_the_metric_gap_is_zero() -> None: + model, clean, _, metric = _faithfulness_toy() + + with pytest.raises(ValueError, match="same metric"): + faithfulness(model, clean, clean, metric, []) + + +def test_faithfulness_raises_on_token_length_mismatch() -> None: + model, _, corrupt, metric = _faithfulness_toy() + shorter = torch.tensor([[1, 2]]) + + with pytest.raises(ValueError, match="same length"): + faithfulness(model, shorter, corrupt, metric, []) + + +def test_faithfulness_raises_on_more_than_one_pair() -> None: + """Batched input is rejected rather than silently half-ablated. + + The metric reads a single example, so a batch would need per-pair + aggregation the caller should own. Silently ablating only the first row + would return a plausible but wrong number. + """ + model, clean, corrupt, metric = _faithfulness_toy() + batched_clean = torch.cat([clean, clean], dim=0) + batched_corrupt = torch.cat([corrupt, corrupt], dim=0) + + with pytest.raises(ValueError, match="one clean/corrupt pair at a time"): + faithfulness(model, batched_clean, batched_corrupt, metric, []) + + +def test_faithfulness_rejects_model_in_training_mode() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + model.train() + + with pytest.raises(ValueError, match="evaluation mode"): + faithfulness(model, clean, corrupt, metric, []) + + +def test_faithfulness_restores_caller_hook_flags() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + assert model.cfg.use_attn_result is False + assert model.cfg.use_split_qkv_input is False + assert model.cfg.use_hook_mlp_in is False + + faithfulness(model, clean, corrupt, metric, []) + + assert model.cfg.use_attn_result is False + assert model.cfg.use_split_qkv_input is False + assert model.cfg.use_hook_mlp_in is False + + +def test_faithfulness_restores_hook_flags_when_it_raises() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + bogus = (Node(kind="embed", position=1), Node(kind="mlp_in", layer=0, position=0)) + + with pytest.raises(ValueError, match="not in the graph"): + faithfulness(model, clean, corrupt, metric, [bogus]) + + assert model.cfg.use_attn_result is False + assert model.cfg.use_split_qkv_input is False + assert model.cfg.use_hook_mlp_in is False + + +def test_faithfulness_recovery_is_graded_not_binary() -> None: + """A partial circuit lands strictly between the corrupt and clean metrics. + + Only edges into the readout position can move this toy's metric, since it is + position-wise and the metric reads the last position. Keeping half of those + must therefore recover part of the gap rather than all or none of it. + """ + model, clean, corrupt, metric = _faithfulness_toy() + edges = _graph_edges(model, corrupt, metric) + readout = [edge for edge in edges if edge[1].position == SEQ_LEN - 1] + assert readout + partial = readout[: len(readout) // 2] + + report = faithfulness(model, clean, corrupt, metric, partial) + + assert 0.0 < report.recovered < 1.0 + assert report.circuit_size == len(partial) + + +# --------------------------------------------------------------------------- +# Edge-class breakout +# --------------------------------------------------------------------------- +# +# Class removal stays inside the evaluated circuit, with a shared baseline. + + +def test_edge_class_partitions_edges_exhaustively_and_without_overlap() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + edges = _graph_edges(model, corrupt, metric) + + classes = [_edge_class(edge) for edge in edges] + + # Every edge maps to exactly one class, and the classes partition the graph. + assert set(classes) == {"into_qk", "into_v", "into_mlp", "into_logits"} + assert len(classes) == len(edges) + + report = faithfulness(model, clean, corrupt, metric, []) + assert sum(report.edge_class_counts.values()) == report.total_edges + assert set(report.edge_class_counts) == set(classes) + + +def test_edge_class_maps_each_reader_kind_to_its_class() -> None: + writer = Node(kind="embed", position=0) + + assert _edge_class((writer, Node(kind="q_input", layer=0, head=0, position=0))) == "into_qk" + assert _edge_class((writer, Node(kind="k_input", layer=0, head=0, position=0))) == "into_qk" + assert _edge_class((writer, Node(kind="v_input", layer=0, head=0, position=0))) == "into_v" + assert _edge_class((writer, Node(kind="mlp_in", layer=0, position=0))) == "into_mlp" + assert _edge_class((writer, Node(kind="logits", layer=0, position=0))) == "into_logits" + + # A writer kind is not a reader, so it has no class. + with pytest.raises(ValueError, match="no edge class"): + _edge_class((writer, Node(kind="mlp_out", layer=0, position=0))) + + +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_edge_class_empty_circuit_is_aggregate(config: FaithfulnessConfig) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + calibration = _toy_mean_calibration() if config.ablation == "mean" else None + report = faithfulness(model, clean, corrupt, metric, [], config, mean_tokens=calibration) + assert report.edge_class_recovered + assert all(value == report.recovered for value in report.edge_class_recovered.values()) + + +def test_edge_class_recovered_reconciles_with_the_aggregate() -> None: + model, clean, corrupt, metric = _faithfulness_toy() + edges = _graph_edges(model, corrupt, metric) + + report = faithfulness(model, clean, corrupt, metric, edges) + + assert set(report.edge_class_recovered) == { + "into_qk", + "into_v", + "into_mlp", + "into_logits", + } + for edge_class, recovered in report.edge_class_recovered.items(): + assert math.isfinite(recovered), f"non-finite recovery for {edge_class}" + assert report.edge_class_counts[edge_class] > 0 + # Keeping every edge is the aggregate, so the full-circuit report recovers + # the whole gap. + assert report.recovered == pytest.approx(1.0, abs=1e-6) + + +_CLASS_READER_KINDS: dict[EdgeClass, tuple[NodeKind, ...]] = { + "into_qk": ("q_input", "k_input"), + "into_v": ("v_input",), + "into_mlp": ("mlp_in",), + "into_logits": ("logits",), +} + + +def _functional_circuit_recovery( + model: _EdgeScoringToyBridge, + clean: torch.Tensor, + corrupt: torch.Tensor, + circuit: list[tuple[Node, Node]], + replacements: dict[Node, torch.Tensor], +) -> float: + clean_logits, _, _ = _functional_toy_ablation(model, clean, set(), {}) + corrupt_logits, _, _ = _functional_toy_ablation(model, corrupt, set(), {}) + logits, _, _ = _functional_toy_ablation( + model, clean, set(_toy_ablation_edges()) - set(circuit), replacements + ) + metric = _metric_fn(1, 2) + return float( + (metric(logits) - metric(corrupt_logits)) / (metric(clean_logits) - metric(corrupt_logits)) + ) + + +@pytest.mark.parametrize("edge_class", list(_CLASS_READER_KINDS)) +@pytest.mark.parametrize("partial", [False, True]) +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_edge_class_removal_matches_functional_candidate_reference( + edge_class: EdgeClass, + partial: bool, + config: FaithfulnessConfig, +) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + edges = _toy_ablation_edges() + circuit = edges[::2] if partial else edges + calibration = _toy_mean_calibration() if config.ablation == "mean" else None + replacements = ( + _functional_mean_writers(model, calibration) + if calibration is not None + else _functional_toy_ablation(model, corrupt, set(), {})[1] + ) + report = faithfulness( + model, clean, corrupt, _metric_fn(1, 2), circuit, config, mean_tokens=calibration + ) + kinds = _CLASS_READER_KINDS[edge_class] + reduced = [edge for edge in circuit if edge[1].kind not in kinds] + expected = _functional_circuit_recovery(model, clean, corrupt, reduced, replacements) + assert report.edge_class_recovered[edge_class] == pytest.approx(expected, abs=1e-12) + assert report.recovered == pytest.approx( + _functional_circuit_recovery(model, clean, corrupt, circuit, replacements), abs=1e-12 + ) + assert report.edge_class_counts[edge_class] == sum(edge[1].kind in kinds for edge in edges) + + +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_edge_class_absent_from_candidate_reuses_aggregate( + config: FaithfulnessConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + circuit = [edge for edge in _toy_ablation_edges() if edge[1].kind == "logits"][:3] + calibration = _toy_mean_calibration() if config.ablation == "mean" else None + calls: list[list[tuple[Node, Node]]] = [] + + def ablate(model: Any, tokens: torch.Tensor, edges: list, kept: list, replacements: dict): + calls.append(list(kept)) + return _ablate_edges(model, tokens, edges, kept, replacements) + + monkeypatch.setattr( + "transformer_lens.tools.analysis.attribution_patching._ablate_edges", ablate + ) + report = faithfulness(model, clean, corrupt, metric, circuit, config, mean_tokens=calibration) + assert calls == [circuit, []] + for edge_class in ("into_qk", "into_v", "into_mlp"): + assert report.edge_class_recovered[edge_class] == report.recovered + assert sum(report.edge_class_counts.values()) == report.total_edges + + +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_edge_class_breakout_changes_with_candidate(config: FaithfulnessConfig) -> None: + model = _EdgeScoringToyBridge(dtype=torch.float64) + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + calibration = _toy_mean_calibration() if config.ablation == "mean" else None + replacements = ( + _functional_mean_writers(model, calibration) + if calibration is not None + else _functional_toy_ablation(model, corrupt, set(), {})[1] + ) + q_edge = (Node("embed", 2), Node("q_input", 2, layer=0, head=0)) + logits = Node("logits", 2, layer=1) + circuits = [ + [q_edge, (Node("embed", 2), logits)], + [q_edge, (Node("mlp_out", 2, layer=0), logits)], + ] + expected = [ + _functional_circuit_recovery(model, clean, corrupt, circuit[1:], replacements) + for circuit in circuits + ] + assert expected[0] != pytest.approx(expected[1], abs=1e-12) + for circuit, value in zip(circuits, expected): + report = faithfulness( + model, clean, corrupt, _metric_fn(1, 2), circuit, config, mean_tokens=calibration + ) + assert report.edge_class_recovered["into_qk"] == pytest.approx(value, abs=1e-12) + + +@pytest.mark.parametrize("ranked", [False, True]) +@pytest.mark.parametrize("config", [FaithfulnessConfig(), FaithfulnessConfig(ablation="mean")]) +def test_faithfulness_duplicate_circuit_entries_are_equivalent( + ranked: bool, + config: FaithfulnessConfig, +) -> None: + model, clean, corrupt, metric = _faithfulness_toy() + circuit = _toy_ablation_edges()[::2] + calibration = _toy_mean_calibration() if config.ablation == "mean" else None + duplicates = circuit[:2] + circuit + circuit[:1] + entries = [(writer, reader, float(index)) for index, (writer, reader) in enumerate(duplicates)] + assert _normalize_circuit(entries if ranked else duplicates) == circuit + reference = faithfulness( + model, clean, corrupt, metric, circuit, config, mean_tokens=calibration + ) + actual = faithfulness( + model, + clean, + corrupt, + metric, + entries if ranked else duplicates, + config, + mean_tokens=calibration, + ) + assert actual == reference + assert actual.circuit_size == len(circuit) + + +# --------------------------------------------------------------------------- +# Random-edge-set baseline +# --------------------------------------------------------------------------- +# +# A ranked circuit is only meaningful if it beats an arbitrary edge set of the +# same size. The toy is nonlinear here on purpose: on a linear model any circuit +# holding every edge into the logits reader recovers the metric outright, so the +# baseline could tie by accident. + + +def _random_edge_sets( + edges: list[tuple[Node, Node]], size: int, draws: int, seed: int +) -> list[list[tuple[Node, Node]]]: + """``draws`` distinct random edge sets of ``size`` edges, from a fixed seed.""" + generator = torch.Generator().manual_seed(seed) + return [ + [edges[index] for index in torch.randperm(len(edges), generator=generator)[:size].tolist()] + for _ in range(draws) + ] + + +def test_random_edge_set_recovers_markedly_less_than_the_attribution_circuit() -> None: + """The ranked circuit beats random edge sets of the same size by a wide margin. + + This is the honesty guard: a ranking that does no better than chance would + still produce a plausible-looking ``recovered`` number, so the comparison + against same-size random sets is what makes the number mean anything. + """ + model = _NonlinearEdgeScoringToyBridge() + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + metric = _metric_fn(answer=1, wrong=2) + + result = attribution_patch( + model, clean, corrupt, metric, config=EdgeAttributionConfig(granularity="edge") + ) + edges = list(result.edge_scores) + budget = 8 + ranked = result.top_edges(k=budget) + ranked_report = faithfulness(model, clean, corrupt, metric, ranked) + + random_reports = [ + faithfulness(model, clean, corrupt, metric, circuit) + for circuit in _random_edge_sets(edges, budget, draws=3, seed=0) + ] + random_recovered = [report.recovered for report in random_reports] + best_random = max(random_recovered) + + # The baseline is not degenerate: if every random set recovered the whole + # gap, the toy would be too easy to discriminate and the comparison would be + # vacuous. + assert all(value < 1.0 for value in random_recovered), random_recovered + assert ranked_report.recovered > best_random + 0.2, ( + f"ranked recovered {ranked_report.recovered:.4f} vs random " + f"{[round(value, 4) for value in random_recovered]}" + ) + + +def test_random_edge_set_baseline_is_reproducible_under_a_fixed_seed() -> None: + model = _NonlinearEdgeScoringToyBridge() + clean = torch.tensor([[1, 2, 3]]) + corrupt = torch.tensor([[3, 2, 1]]) + metric = _metric_fn(answer=1, wrong=2) + + result = attribution_patch( + model, clean, corrupt, metric, config=EdgeAttributionConfig(granularity="edge") + ) + edges = list(result.edge_scores) + + first = _random_edge_sets(edges, 8, draws=3, seed=0) + second = _random_edge_sets(edges, 8, draws=3, seed=0) + + assert first == second + # Different seeds give different draws, so the seed is actually used. + assert first != _random_edge_sets(edges, 8, draws=3, seed=1) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 13ac09858..7ba0fb2b3 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -8,15 +8,16 @@ - attribution_patching: Attribution patching (gradient-linearized activation patching) over residual-stream nodes and edges: typed computational graph, a names-filtered manual-backward gradient cache, and signed node and edge - scores (EAP). Integrated gradients (EAP-IG) and ablate-outside faithfulness - are not implemented yet. + scores (EAP), plus forward-only ablate-outside faithfulness with corrupt + or held-out mean replacements and circuit-relative class diagnostics. + Integrated gradients (EAP-IG) are not implemented. - backward_lens: GPT-2 MLP weight-gradient factors projected into vocabulary space with explicit raw-gradient sign semantics. - direct_logit_attribution: Direct Logit Attribution (DLA) over components, layers, or attention heads. - direct_path_patching: Direct path patching for head-to-head circuit analysis. - - jacobian_lens: The Jacobian lens (J-lens) — per-layer causal transport to + - jacobian_lens: The Jacobian lens (J-lens) -- per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, J-space sparse decomposition, and anchored coordinate patching (offline and dynamic/hooked). @@ -43,8 +44,11 @@ from transformer_lens.tools.analysis.attribution_patching import ( AttributionResult, EdgeAttributionConfig, + FaithfulnessConfig, + FaithfulnessResult, Node, attribution_patch, + faithfulness, ) from transformer_lens.tools.analysis.backward_lens import ( BackwardLens, @@ -145,6 +149,8 @@ "DegenerateDirectionError", "DirectLogitAttribution", "EdgeAttributionConfig", + "FaithfulnessConfig", + "FaithfulnessResult", "FunctionSpec", "HeadAffinityPair", "HeadAffinityResult", @@ -179,6 +185,7 @@ "decompose_head", "direct_logit_attribution", "estimate_occupancy", + "faithfulness", "fit_sparse_probe", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", diff --git a/transformer_lens/tools/analysis/attribution_patching.py b/transformer_lens/tools/analysis/attribution_patching.py index 6cf392d34..9ea6143a5 100644 --- a/transformer_lens/tools/analysis/attribution_patching.py +++ b/transformer_lens/tools/analysis/attribution_patching.py @@ -1,4 +1,4 @@ -"""Attribution patching — linearized activation patching on ``TransformerBridge``. +"""Attribution patching -- linearized activation patching on ``TransformerBridge``. Attribution patching estimates the causal effect of every model component on a task metric with a *gradient-based linearization* of activation patching: rather @@ -9,8 +9,7 @@ first-order Taylor estimate ``effect(node) = (a_clean - a_corrupt) . g``. Scores over a batch of clean/corrupt pairs are averaged before ranking. -Only the ``TransformerBridge`` API is targeted; TransformerLens v4 deprecates -``HookedTransformer``. +Only the ``TransformerBridge`` API is targeted. Sign/direction convention (denoising form): the gradient is taken on the *corrupt* run and the estimate points *toward* the clean activation, so a @@ -23,17 +22,37 @@ should filter to the hook families their analysis actually reads. Scope: this build ships node and edge granularity with plain attribution -(``ig_steps=1``). The integrated-gradient path (EAP-IG, ``ig_steps>1``) and -ablate-outside faithfulness are not implemented yet; their API is declared here, -and ``ig_steps>1`` raises :class:`NotImplementedError`, so downstream code can pin +(``ig_steps=1``), plus ablate-outside faithfulness. The integrated-gradient path +(EAP-IG, ``ig_steps>1``) is not implemented yet; its API is declared here, and +``ig_steps>1`` raises :class:`NotImplementedError`, so downstream code can pin against a stable surface now. + +Faithfulness intervenes rather than estimating: it ablates every edge outside a +candidate circuit and reports how much of the clean-to-corrupt metric gap the +circuit recovers. The residual stream is a running sum, so a reader's input is +exactly the sum of its writers' contributions, and rebuilding that input from +the included writers plus replacements for the excluded ones is an exact +correction rather than an approximation. The reader's fork hooks +(``attn.hook_q_input`` / ``hook_k_input`` / ``hook_v_input``, ``hook_mlp_in``) +fire before the layer norm, so the correction is a plain add and subtract in +``d_model`` space with no norm scale to divide out. """ from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Any, Callable, Iterator, Literal, Optional, Sequence, Union +from typing import ( + Any, + Callable, + Collection, + Iterator, + Literal, + Optional, + Sequence, + Union, + cast, +) import torch @@ -46,6 +65,8 @@ "embed", "attn_head_out", "mlp_out", "q_input", "k_input", "v_input", "mlp_in", "logits" ] Granularity = Literal["node", "edge"] +AblationMode = Literal["corrupt", "mean"] +EdgeClass = Literal["into_qk", "into_v", "into_mlp", "into_logits"] @dataclass @@ -172,9 +193,8 @@ class EdgeAttributionConfig: This build implements node and edge granularity with plain attribution. ``ig_steps>1`` is accepted by the type but raises :class:`NotImplementedError` at construction, so downstream code can import and reference this API now - while the integrated-gradient path is not implemented yet. Once EAP-IG lands, - the default flips to ``ig_steps=5`` (EAP-IG is the faithful default); until - then the default is the only executable value, ``ig_steps=1``. + while the integrated-gradient path is not implemented. The default is the + only executable value, ``ig_steps=1``. Attributes: granularity: ``"node"`` or ``"edge"``. Defaults to ``"node"``. @@ -219,8 +239,8 @@ def top_nodes(self, k: int = 10) -> list[tuple[Node, float]]: """The ``k`` nodes with the largest effect magnitude, strongest first. Ranking is by absolute score: a node with a large negative effect is as - causally important as one with a large positive effect, so magnitude — not - signed value — orders the circuit. Ties keep enumeration order (stable + causally important as one with a large positive effect, so magnitude -- not + signed value -- orders the circuit. Ties keep enumeration order (stable sort). Requesting more than the available nodes returns all of them. """ ranked = sorted(self.node_scores.items(), key=lambda item: abs(item[1]), reverse=True) @@ -238,6 +258,57 @@ def top_edges(self, k: int = 10) -> list[tuple[Node, Node, float]]: return [(writer, reader, score) for (writer, reader), score in ranked[:k]] +@dataclass(frozen=True) +class FaithfulnessConfig: + """Configuration for an ablate-outside faithfulness measurement. + + Attributes: + ablation: How an out-of-circuit edge's writer is replaced. + ``"corrupt"`` substitutes the corrupt run's own contribution, which + is the evaluation the pinned external reference reports and the + default here so the two compare like quantities. ``"mean"`` + substitutes that writer's mean contribution over the independent + calibration batch supplied as ``mean_tokens`` to :func:`faithfulness`. + The evaluated clean/corrupt pair is not used to estimate this mean. + """ + + ablation: AblationMode = "corrupt" + + +@dataclass +class FaithfulnessResult: + """How much of the clean-to-corrupt metric gap a candidate circuit recovers. + + Attributes: + recovered: Fraction of the clean-to-corrupt metric gap the circuit + recovers, ``(m_ablated - m_corrupt) / (m_clean - m_corrupt)``. A + fraction, not a percentage: ``1.0`` means ablating every edge + outside the circuit reproduces the clean metric, ``0.0`` means it + reproduces the corrupt metric. Values outside ``[0, 1]`` are + possible and meaningful -- a circuit can overshoot the clean run. + full_metric: The clean run's metric, the top of the gap. + corrupt_metric: The corrupt run's metric, the bottom of the gap. + circuit_size: Number of edges kept. + total_edges: Number of edges in the graph, so the budget is visible. + edge_class_recovered: Recovery after removing this class's retained + edges from the candidate circuit. All edges originally outside the + circuit stay ablated, with the same replacements and metric gap. + A class absent from the circuit has the aggregate recovery. This + measures conditional circuit dependence, not attribution ranking + accuracy. Interactions can increase recovery after removal; class + contributions are not additive and need not lie in ``[0, 1]``. + edge_class_counts: Full-graph edge counts per class, not retained counts. + """ + + recovered: float + full_metric: float + corrupt_metric: float + circuit_size: int + total_edges: int + edge_class_recovered: dict[EdgeClass, float] = field(default_factory=dict) + edge_class_counts: dict[EdgeClass, int] = field(default_factory=dict) + + def _required_hook_names(n_layers: int, granularity: Granularity = "node") -> list[str]: """Hook points a sweep at ``granularity`` reads. @@ -384,7 +455,7 @@ def enumerate_nodes(model: Any, cache: GradientCache) -> list[Node]: The node list, ordered embed-then-layerwise for deterministic ranking. Raises: - ValueError: if any required hook point is absent from ``cache`` — the + ValueError: if any required hook point is absent from ``cache`` -- the graph is never silently truncated. """ n_layers = int(model.cfg.n_layers) @@ -453,6 +524,40 @@ def _edge_hook_names(n_layers: int) -> list[str]: ) +def _writer_hook_names(n_layers: int) -> list[str]: + """The hook points holding a writer's own residual-stream contribution. + + The three writer families an edge sweep scores: the token embedding write, + each attention head's per-head output (``attn.hook_result``, the + decomposition of the head's contribution after it is projected into the + residual stream), and each layer's MLP output. These are the hook points + :func:`_writer_hook_name` resolves a single writer node to, listed as a + family so a capture can filter to exactly them. + """ + names = ["hook_embed"] + for layer in range(n_layers): + names.append(f"blocks.{layer}.attn.hook_result") + names.append(f"blocks.{layer}.hook_mlp_out") + return names + + +def _reader_hook_names(n_layers: int) -> list[str]: + """The hook points holding a reader's input. + + The residual each reader consumes: the split ``attn.hook_q_input`` / + ``hook_k_input`` / ``hook_v_input`` per head, and each layer's MLP entry + ``hook_mlp_in``. The terminal logits reader reads the final + ``hook_resid_post``, which :func:`_required_edge_reader_hook_names` lists. + """ + names: list[str] = [] + for layer in range(n_layers): + names.append(f"blocks.{layer}.attn.hook_q_input") + names.append(f"blocks.{layer}.attn.hook_k_input") + names.append(f"blocks.{layer}.attn.hook_v_input") + names.append(f"blocks.{layer}.hook_mlp_in") + return names + + def enumerate_edges(model: Any, cache: GradientCache) -> list[tuple[Node, Node]]: """Enumerate every writer -> reader edge in the residual-stream graph. @@ -586,7 +691,7 @@ def cache_activation_and_gradient( inert and its ``.grad`` stays ``None``. A backward hook goes through the same conversion and delivers the real gradient in canonical shape. Driving the backward with :func:`torch.autograd.grad` instead of ``metric.backward()`` - keeps it off every parameter's ``.grad`` buffer — no caller grads clobbered, no + keeps it off every parameter's ``.grad`` buffer -- no caller grads clobbered, no model-sized buffer allocated. With ``compute_gradient=False`` the backward is skipped entirely: only forward @@ -601,7 +706,7 @@ def cache_activation_and_gradient( tokens: Input token ids for a single forward pass. metric_fn: Maps the model logits to a scalar to differentiate. names_filter: Restricts which hook points are cached (and have gradients - retained). ``None`` (the default) caches the node-granularity hook set — + retained). ``None`` (the default) caches the node-granularity hook set -- ``hook_embed`` plus each layer's ``attn.hook_z`` and ``hook_mlp_out``; pass an explicit filter to cache any hook point outside that set. On a real Bridge, ``None`` cannot mean "every hook point": the gated points @@ -696,8 +801,8 @@ def hook(grad: torch.Tensor, *, hook: Any) -> None: else () ) - # The most-upstream cached point's own backward hook never fires — nothing - # requested is upstream of it, so the backward stops at it — but its cached + # The most-upstream cached point's own backward hook never fires -- nothing + # requested is upstream of it, so the backward stops at it -- but its cached # tensor is on-path and unconverted, so torch.autograd.grad returns that # gradient directly. Backfill any point the hooks missed from that return. for name, grad in zip(captured_names, returned): @@ -707,6 +812,76 @@ def hook(grad: torch.Tensor, *, hook: Any) -> None: return GradientCache(activations=activations, gradients=gradients, metric=metric.detach()) +def capture_writer_outputs( + model: Any, + tokens: torch.Tensor, + names_filter: NamesFilter = None, +) -> dict[str, torch.Tensor]: + """Capture writer contributions in one forward-only pass. + + Rewriting a reader's input during a live forward pass needs the writers' + *current* contributions rather than a gradient estimate. This runs a single + forward under :func:`torch.no_grad` with forward hooks only: no backward is + driven and no gradient is retained. That is the opposite contract from + :func:`cache_activation_and_gradient`, which exists to retain gradients and + refuses to run without autograd, so the two are kept as separate helpers + rather than one helper with a mode switch. + + The default filter is the writer-output family only -- ``hook_embed``, + ``blocks.*.attn.hook_result``, and ``blocks.*.hook_mlp_out``. Caching the + reader-input family as well would roughly double the captured memory for no + benefit here, since a caller rewriting a reader's input reads that input + from the live forward rather than from a cache. + + Memory caveat: ``attn.hook_result`` is a per-head + ``[batch, seq, n_heads, d_model]`` tensor, so the capture is fine on a model + the size of gpt2-small and does not scale to models with many heads or + layers. The hook point only fires while ``cfg.use_attn_result`` is on; see + :func:`_edge_hook_flags`. + + Args: + model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``, + ``hook_dict``, and the ``hooks()`` context manager. + tokens: Input token ids for a single forward pass. + names_filter: Restricts which hook points are captured. ``None`` (the + default) captures the writer-output family. On a real Bridge the + gated points (``attn.hook_result``, the split-QKV inputs, + ``hook_mlp_in``) raise in ``add_hook`` unless their ``set_use_*`` + flag is on, so a filter reaching them must be paired with + :func:`_edge_hook_flags`. + + Returns: + Detached activation tensors keyed by hook name, for every hook point + matching ``names_filter``. + + Raises: + ValueError: if ``names_filter`` matches no hook point. + """ + if names_filter is None: + names_filter = _writer_hook_names(int(model.cfg.n_layers)) + predicate = _as_predicate(names_filter) + names = [name for name in model.hook_dict if predicate(name)] + if not names: + raise ValueError("names_filter matched no hook points") + + captured: dict[str, torch.Tensor] = {} + + def make_fwd_hook(name: str) -> Callable[..., None]: + def hook(tensor: torch.Tensor, *, hook: Any) -> None: + del hook + if isinstance(tensor, torch.Tensor): + captured[name] = tensor.detach().clone() + return None + + return hook + + fwd_hooks = [(name, make_fwd_hook(name)) for name in names] + with torch.no_grad(), model.hooks(fwd_hooks=fwd_hooks): + model(tokens) + + return captured + + def _node_effects( clean_cache: GradientCache, corrupt_cache: GradientCache, @@ -846,6 +1021,170 @@ def _aggregate_edge_scores_to_writer_nodes( return totals +def _writer_contribution_slice(writer: Node, tensor: torch.Tensor) -> torch.Tensor: + """The writer's own contribution at its position, and head if it has one. + + A captured writer tensor holds every position (and, for an attention head, + every head); an edge carries only the writer's contribution at the edge's + position, so this selects that slice. The result is a ``d_model`` vector, + which broadcasts into a reader's input -- a per-head reader receives the + writer once per head. + """ + if writer.kind == "attn_head_out": + return tensor[0, writer.position, writer.head] + return tensor[0, writer.position] + + +def _excluded_writers_by_reader( + edges: Sequence[tuple[Node, Node]], + circuit: Collection[tuple[Node, Node]], +) -> dict[Node, list[Node]]: + """Group the out-of-circuit writers feeding each reader node. + + A reader node's input is rebuilt once, from the writers whose edges into it + lie outside ``circuit``. Grouping by reader node keeps that rebuild to a + single hook per node: one hook per excluded edge would apply the correction + once per edge instead of once per node. + """ + circuit_set = set(circuit) + excluded: dict[Node, list[Node]] = {} + for writer, reader in edges: + if (writer, reader) not in circuit_set: + excluded.setdefault(reader, []).append(writer) + return excluded + + +def _make_writer_capture_hook(name: str, sink: dict[str, torch.Tensor]) -> Callable[..., None]: + """Record a writer's live contribution for the reader hooks that follow it. + + A writer always precedes the readers it feeds, so by the time a reader's + fork hook fires the contributions it needs are already in ``sink``. The + tensor is detached rather than cloned: it is only read, and the reader hook + clones its own input before rewriting it. + """ + + def hook(tensor: torch.Tensor, *, hook: Any) -> None: + del hook + if isinstance(tensor, torch.Tensor): + sink[name] = tensor.detach() + return None + + return hook + + +def _make_edge_ablation_hook( + reader: Node, + excluded_writers: Sequence[Node], + live_writers: dict[str, torch.Tensor], + replacement_writers: dict[str, torch.Tensor], +) -> Callable[..., torch.Tensor]: + """Rebuild one reader's input with its out-of-circuit writers replaced. + + The residual stream is a running sum, so a reader's input is exactly the sum + of its writers' contributions. Removing an edge therefore means subtracting + that writer's live contribution and adding its replacement -- an exact + correction rather than an approximation. The reader's fork hook fires before + the layer norm, so the correction is a plain add and subtract in ``d_model`` + space with no norm scale to divide out. + + A per-head reader's input is the residual replicated across heads, so the + correction is applied to that reader's own head slot. Several reader nodes + share one hook point (one per head, one per position), and each rewrites a + disjoint slice, so no writer's contribution is subtracted twice. + + Args: + reader: The reader whose input is rebuilt. + excluded_writers: Writers whose edges into ``reader`` lie outside the + circuit. + live_writers: Contributions captured during the current forward, keyed + by writer hook name. + replacement_writers: The value standing in for each excluded writer, + keyed by writer hook name. + + Returns: + A forward hook returning the rebuilt input. + + Raises: + ValueError: if an excluded writer's contribution was not captured, which + means its hook family was not enabled for this forward. + """ + if reader.kind in ("q_input", "k_input", "v_input"): + # Node.__post_init__ guarantees a head for the per-head reader kinds. + index: tuple[int, ...] = (0, reader.position, cast(int, reader.head)) + else: + index = (0, reader.position) + + def hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor: + del hook + rebuilt = tensor.clone() + for writer in excluded_writers: + name = _writer_hook_name(writer) + live = live_writers.get(name) + replacement = replacement_writers.get(name) + if live is None or replacement is None: + raise ValueError( + f"cannot ablate the edge from {writer} into {reader}: no " + f"contribution captured at {name!r}. Enable the writer's hook " + "family before running the ablation." + ) + rebuilt[index] = ( + rebuilt[index] + - _writer_contribution_slice(writer, live) + + _writer_contribution_slice(writer, replacement) + ) + return rebuilt + + return hook + + +def _ablate_edges( + model: Any, + tokens: torch.Tensor, + edges: Sequence[tuple[Node, Node]], + circuit: Collection[tuple[Node, Node]], + replacement_writers: dict[str, torch.Tensor], +) -> torch.Tensor: + """Run one forward with every edge outside ``circuit`` ablated. + + Captures each writer's live contribution and rewrites each reader's input at + its fork hook, so the forward sees a model in which only the circuit's edges + carry the run's own values. Forward-only: no gradient is involved, and the + caller must have the writer and reader hook families enabled (see + :func:`_edge_hook_flags`). + + Args: + model: A ``TransformerBridge`` (or compatible). + tokens: Input token ids for the ablated forward. + edges: The full edge list, as :func:`enumerate_edges` returns it. + circuit: The edges to keep. Every other edge is ablated. + replacement_writers: The value standing in for each ablated writer, + keyed by writer hook name. + + Returns: + The model's logits under the ablation, detached. The rewrite is a + forward-only intervention, so the pass runs under :func:`torch.no_grad` + and no gradient is retained. + """ + excluded_by_reader = _excluded_writers_by_reader(edges, circuit) + + live_writers: dict[str, torch.Tensor] = {} + fwd_hooks: list[tuple[str, Callable[..., Any]]] = [ + (name, _make_writer_capture_hook(name, live_writers)) + for name in _writer_hook_names(int(model.cfg.n_layers)) + ] + fwd_hooks.extend( + ( + reader.hook_name, + _make_edge_ablation_hook(reader, writers, live_writers, replacement_writers), + ) + for reader, writers in excluded_by_reader.items() + ) + + with torch.no_grad(), model.hooks(fwd_hooks=fwd_hooks): + logits = model(tokens) + return logits + + def attribution_patch( model: Any, clean: torch.Tensor, @@ -975,3 +1314,246 @@ def attribution_patch( node_scores = {node: total / batch for node, total in totals.items()} return AttributionResult(node_scores=node_scores) + + +def _normalize_circuit( + circuit: Sequence[tuple[Node, Node]] | Sequence[tuple[Node, Node, float]], +) -> list[tuple[Node, Node]]: + """Deduplicate circuit pairs in first-occurrence order, dropping any scores. + + Accepts either the bare edge list :func:`enumerate_edges` returns or the + ranked ``(writer, reader, score)`` triples ``AttributionResult.top_edges`` + returns, so a discovered edge set feeds straight into + :func:`faithfulness` without reshaping. + """ + normalized: list[tuple[Node, Node]] = [] + seen: set[tuple[Node, Node]] = set() + for entry in circuit: + if len(entry) == 2: + writer, reader = entry + elif len(entry) == 3: + writer, reader, _score = entry + else: + raise ValueError( + f"circuit entries must be (writer, reader) or (writer, reader, score), " + f"got a {len(entry)}-tuple" + ) + edge = (writer, reader) + if edge not in seen: + normalized.append(edge) + seen.add(edge) + return normalized + + +def _edge_class(edge: tuple[Node, Node]) -> EdgeClass: + """The class an edge belongs to, keyed by what its reader consumes. + + Q and K share a class because both feed attention scores. Class removal + measures the candidate circuit's dependence on those retained edges; it + does not establish relative ranking fidelity or downstream linearity. + """ + reader = edge[1] + if reader.kind in ("q_input", "k_input"): + return "into_qk" + if reader.kind == "v_input": + return "into_v" + if reader.kind == "mlp_in": + return "into_mlp" + if reader.kind == "logits": + return "into_logits" + raise ValueError(f"{reader.kind} is a writer kind and has no edge class") + + +def _mean_writer_contributions( + model: Any, + mean_tokens: torch.Tensor, +) -> dict[str, torch.Tensor]: + """Average calibration examples, preserving position, head, and singleton batch axes.""" + count = int(mean_tokens.shape[0]) + if count == 0: + raise ValueError("mean_tokens must contain at least one calibration example") + totals: dict[str, torch.Tensor] = {} + for index in range(count): + captured = capture_writer_outputs(model, mean_tokens[index : index + 1]) + for name, tensor in captured.items(): + running = totals.get(name) + totals[name] = tensor if running is None else running + tensor + return {name: total / count for name, total in totals.items()} + + +def faithfulness( + model: Any, + clean: torch.Tensor, + corrupt: torch.Tensor, + metric_fn: MetricFn, + circuit: Sequence[tuple[Node, Node]] | Sequence[tuple[Node, Node, float]], + config: FaithfulnessConfig = FaithfulnessConfig(), + *, + mean_tokens: Optional[torch.Tensor] = None, +) -> FaithfulnessResult: + """Measure how much of the clean-to-corrupt metric gap a circuit recovers. + + Ablates every edge *outside* ``circuit`` and reports the metric the model + then produces. The residual stream is a running sum, so each reader's input + is rebuilt exactly -- subtracting the excluded writers' live contributions + and adding their replacements -- rather than approximated. A circuit that + recovers most of the gap is a faithful account of the behavior; one that + recovers little is not, however well it ranks. + + This is a forward-only measurement: no gradient is taken, so it is + independent of the linearization :func:`attribution_patch` uses to rank + edges. Feed a ranked edge set straight in from + ``attribution_patch(..., config=EdgeAttributionConfig(granularity="edge")).top_edges(k=...)``. + + Mean ablation requires held-out ``mean_tokens`` rather than the evaluated + endpoints. Recovery is still normalized to the clean/corrupt metric gap, + so an empty mean-ablated circuit need not have zero recovery. + + Args: + model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``, + ``hook_dict``, and ``hooks()``. + clean: Clean token ids for one evaluated example, shape ``[1, seq]``. + corrupt: Paired corrupt token ids, shape ``[1, seq]``. + metric_fn: Maps single-example logits to a scalar. + circuit: The edges to keep, as ``(writer, reader)`` pairs or the + ``(writer, reader, score)`` triples ``top_edges`` returns. Every + other edge in the graph is ablated. Duplicate entries count once. + config: Ablation configuration. Defaults to replacing an ablated writer + with the corrupt run's own contribution. + mean_tokens: Required only for ``ablation="mean"``. Held-out calibration + token ids, shape ``[n_calibration, seq]`` with a nonempty batch and + the evaluated sequence length. Must have dtype ``torch.int32`` or + ``torch.int64`` and share the clean/corrupt token device. Writer + activations are averaged over examples, retaining sequence positions + and heads. Exact rows matching either evaluated endpoint are rejected; + callers must also avoid semantically equivalent dataset leakage. + Calibration uses the same model and processing mode as evaluation. + + Returns: + A :class:`FaithfulnessResult` with the recovered fraction, the clean and + corrupt metrics bounding it, the circuit's size against the graph's, and + a per-edge-class breakout. Each class entry removes only that class's + retained edges from the candidate. It shares the aggregate's replacements + and normalization, and costs one extra forward per retained class. + Classes absent from the candidate reuse aggregate recovery. Class counts + describe the full graph. This is conditional circuit dependence, not + an estimate of ranking error; removals need not reduce recovery. + + Raises: + ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number + of pairs, hold more than one pair, a pair tokenizes to different + lengths, the model or a submodule is in training mode, a circuit edge + is not in the graph, the clean and corrupt metrics are equal (so + the recovered fraction is undefined), mean calibration is missing, + empty, misaligned, on a different input device, has an unsupported + dtype, contains an evaluated endpoint, or is supplied in corrupt mode. + """ + if clean.ndim != 2 or corrupt.ndim != 2: + raise ValueError( + "faithfulness expects 2D [batch, seq] token tensors, got clean " + f"{tuple(clean.shape)} and corrupt {tuple(corrupt.shape)}" + ) + if clean.shape[0] != corrupt.shape[0]: + raise ValueError( + "clean and corrupt must hold the same number of prompt pairs, got " + f"{clean.shape[0]} and {corrupt.shape[0]}" + ) + if clean.shape[0] != 1: + raise ValueError( + "faithfulness measures one clean/corrupt pair at a time, since the " + f"metric reads a single example; got {clean.shape[0]} pairs. Loop over " + "pairs and aggregate the recovered fractions yourself." + ) + if clean.shape[1] != corrupt.shape[1]: + raise ValueError( + "each clean/corrupt pair must tokenize to the same length; got clean " + f"length {clean.shape[1]} and corrupt length {corrupt.shape[1]}. " + "Faithfulness aligns activations position-by-position." + ) + + if config.ablation == "mean": + if mean_tokens is None: + raise ValueError("mean ablation requires independent calibration via mean_tokens") + if mean_tokens.ndim != 2: + raise ValueError("mean_tokens must be a 2D [n_calibration, seq] token tensor") + if mean_tokens.shape[0] == 0: + raise ValueError("mean_tokens must contain at least one calibration example") + if mean_tokens.shape[1] != clean.shape[1]: + raise ValueError("mean_tokens must have the same sequence length as the evaluated pair") + if mean_tokens.dtype not in (torch.int32, torch.int64): + raise ValueError( + "mean_tokens must contain token ids with dtype torch.int32 or torch.int64" + ) + if mean_tokens.device != clean.device or mean_tokens.device != corrupt.device: + raise ValueError("mean_tokens must be on the same device as the clean/corrupt tokens") + if (mean_tokens == clean).all(dim=1).any() or (mean_tokens == corrupt).all(dim=1).any(): + raise ValueError("mean_tokens must not contain either evaluated clean/corrupt endpoint") + elif mean_tokens is not None: + raise ValueError("mean_tokens is only accepted with mean ablation") + + require_eval_mode(model, operation="faithfulness()") + + kept = _normalize_circuit(circuit) + n_layers = int(model.cfg.n_layers) + hook_names = _edge_hook_names(n_layers) + + with torch.no_grad(), _edge_hook_flags(model): + corrupt_cache = cache_activation_and_gradient( + model, corrupt, metric_fn, names_filter=hook_names, compute_gradient=False + ) + edges = enumerate_edges(model, corrupt_cache) + graph = set(edges) + unknown = [edge for edge in kept if edge not in graph] + if unknown: + raise ValueError( + f"circuit holds {len(unknown)} edge(s) that are not in the graph, " + f"starting with {unknown[0]}; a stale or mistyped circuit would " + "otherwise ablate nothing." + ) + + if mean_tokens is not None: + replacements = _mean_writer_contributions(model, mean_tokens) + else: + replacements = capture_writer_outputs(model, corrupt) + + # Forward-only throughout: no gradient is taken, so the clean endpoint + # is evaluated without building a graph. + with torch.no_grad(): + clean_metric = float(metric_fn(model(clean))) + corrupt_metric = float(corrupt_cache.metric) + gap = clean_metric - corrupt_metric + if gap == 0.0: + raise ValueError( + "the clean and corrupt runs produce the same metric, so the recovered " + "fraction is undefined; choose a pair the metric separates." + ) + ablated_metric = float(metric_fn(_ablate_edges(model, clean, edges, kept, replacements))) + + recovered = (ablated_metric - corrupt_metric) / gap + class_counts: dict[EdgeClass, int] = {} + for edge in edges: + edge_class = _edge_class(edge) + class_counts[edge_class] = class_counts.get(edge_class, 0) + 1 + + kept_classes = {_edge_class(edge) for edge in kept} + class_recovered: dict[EdgeClass, float] = {} + for edge_class in class_counts: + if edge_class not in kept_classes: + class_recovered[edge_class] = recovered + continue + # Removing a class must not restore edges excluded by the candidate. + reduced = [edge for edge in kept if _edge_class(edge) != edge_class] + class_metric = float( + metric_fn(_ablate_edges(model, clean, edges, reduced, replacements)) + ) + class_recovered[edge_class] = (class_metric - corrupt_metric) / gap + + return FaithfulnessResult( + recovered=recovered, + full_metric=clean_metric, + corrupt_metric=corrupt_metric, + circuit_size=len(kept), + total_edges=len(edges), + edge_class_recovered=class_recovered, + edge_class_counts=class_counts, + )