From 0c62ac66c452f8fd3a596475789fdb87c8de685f Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 29 Aug 2026 15:18:00 +0530 Subject: [PATCH 1/6] test: establish exact and attribution-patching baselines --- tests/test_atp_star.py | 100 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 tests/test_atp_star.py diff --git a/tests/test_atp_star.py b/tests/test_atp_star.py new file mode 100644 index 00000000..2c325bdd --- /dev/null +++ b/tests/test_atp_star.py @@ -0,0 +1,100 @@ +"""Correctness tests for the public-API AtP* research pattern.""" + +from collections import OrderedDict +import torch +from nnsight import NNsight + + +def _linear_model() -> NNsight: + """A deterministic model where first-order attribution is exact.""" + network = torch.nn.Sequential( + OrderedDict( + [ + ("components", torch.nn.Linear(3, 3, bias=False)), + ("readout", torch.nn.Linear(3, 1, bias=False)), + ] + ) + ) + with torch.no_grad(): + network.components.weight.copy_( + torch.tensor( + [ + [1.0, 0.5, -0.5], + [0.0, 2.0, 1.0], + [-1.0, 0.0, 0.25], + ] + ) + ) + network.readout.weight.copy_(torch.tensor([[2.0, -3.0, 0.5]])) + return NNsight(network) + + +def test_basic_atp_matches_exact_component_patching(): + """AtP and exact patching agree when the downstream model is linear.""" + model = _linear_model() + clean = torch.tensor([[2.0, -1.0, 0.5]]) + noise = torch.tensor([[-1.0, 0.5, 2.0]]) + + with model.trace(clean): + clean_components = model.components.output.save() + clean_metric = model.output.sum().save() + + with model.trace(noise): + noise_ref = model.components.output + noise_ref.requires_grad_(True) + noise_components = noise_ref.save() + noise_metric_ref = model.output.sum() + noise_metric = noise_metric_ref.save() + with noise_metric_ref.backward(): + noise_gradient = noise_ref.grad.clone().save() + + atp_effects = ((clean_components - noise_components) * noise_gradient)[0] + exact_effects = [] + for component in range(clean_components.shape[-1]): + with model.trace(noise): + model.components.output[:, component] = clean_components[:, component] + patched_metric = model.output.sum().save() + exact_effects.append(patched_metric - noise_metric) + exact_effects = torch.stack(exact_effects) + + torch.testing.assert_close(atp_effects, exact_effects) + torch.testing.assert_close(atp_effects.sum(), clean_metric - noise_metric) + assert torch.equal( + atp_effects.abs().argsort(descending=True), + exact_effects.abs().argsort(descending=True), + ) + + +def test_basic_atp_is_only_an_approximation_after_a_nonlinearity(): + """The test suite must not imply that AtP scores are causal effects.""" + network = torch.nn.Sequential( + OrderedDict( + [ + ("components", torch.nn.Linear(1, 1, bias=False)), + ("saturation", torch.nn.Tanh()), + ] + ) + ) + with torch.no_grad(): + network.components.weight.fill_(1.0) + model = NNsight(network) + + clean = torch.tensor([[0.0]]) + noise = torch.tensor([[4.0]]) + with model.trace(clean): + clean_component = model.components.output.save() + with model.trace(noise): + noise_ref = model.components.output + noise_ref.requires_grad_(True) + noise_component = noise_ref.save() + noise_metric_ref = model.output.sum() + noise_metric = noise_metric_ref.save() + with noise_metric_ref.backward(): + noise_gradient = noise_ref.grad.clone().save() + with model.trace(noise): + model.components.output = clean_component + patched_metric = model.output.sum().save() + + atp_effect = ((clean_component - noise_component) * noise_gradient).sum() + exact_effect = patched_metric - noise_metric + assert atp_effect.abs() < exact_effect.abs() * 0.01 From 16c73308613e362352b9a9ed32dcdb7fa492bf0b Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 29 Aug 2026 15:18:10 +0530 Subject: [PATCH 2/6] test: validate AtP* attention QK correction --- tests/test_atp_star.py | 245 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 245 insertions(+) diff --git a/tests/test_atp_star.py b/tests/test_atp_star.py index 2c325bdd..9ca3350d 100644 --- a/tests/test_atp_star.py +++ b/tests/test_atp_star.py @@ -2,9 +2,108 @@ from collections import OrderedDict import torch +import pytest +from transformers import ( + GPT2Config, + GPT2LMHeadModel, + GPTNeoXConfig, + GPTNeoXForCausalLM, +) from nnsight import NNsight +def _attention_probabilities(query, key, mask=None, scale=None): + scores = torch.einsum("...qd,...kd->...qk", query, key) + if scale is None: + scale = query.shape[-1] ** -0.5 + scores = scores * scale + if mask is not None: + scores = scores + mask + return scores.softmax(dim=-1) + + +def _query_output_delta( + clean_query, noise_key, noise_value, noise_output, mask=None, scale=None +): + patched_probabilities = _attention_probabilities( + clean_query, noise_key, mask, scale + ) + patched_output = torch.einsum( + "...qk,...kd->...qd", patched_probabilities, noise_value + ) + return patched_output - noise_output + + +def _key_output_delta( + noise_query, + clean_key, + noise_key, + noise_value, + noise_probabilities, + noise_output, + scale=None, +): + """Exact local output deltas for every single-key patch in O(T²D).""" + score_delta = torch.einsum("...qd,...kd->...qk", noise_query, clean_key - noise_key) + if scale is None: + scale = noise_query.shape[-1] ** -0.5 + score_delta = score_delta * scale + + log_odds = torch.log(noise_probabilities) - torch.log1p(-noise_probabilities) + patched_probability = torch.sigmoid(log_odds + score_delta) + denominator = torch.where( + noise_probabilities == 1, + torch.ones_like(noise_probabilities), + 1 - noise_probabilities, + ) + probability_scale = (patched_probability - noise_probabilities) / denominator + probability_scale = torch.where( + noise_probabilities == 1, + torch.zeros_like(probability_scale), + probability_scale, + ) + value_difference = noise_value.unsqueeze(-3) - noise_output.unsqueeze(-2) + return probability_scale.unsqueeze(-1) * value_difference + + +def _attention_source_call(attention, family): + operation_names = { + "gpt2": ("attention_interface_0",), + "pythia": ("attention_interface_0", "unknown_0"), + } + if family not in operation_names: + raise ValueError(f"Unsupported AtP* model family: {family}") + for name in operation_names[family]: + try: + return getattr(attention.source, name) + except AttributeError: + pass + raise ValueError( + f"Unsupported {family} attention source layout; expected one of " + f"{operation_names[family]}" + ) + + +def _unpack_attention_call(call): + args, kwargs = call.inputs + offset = 0 if torch.is_tensor(args[0]) else 1 + if len(args) < offset + 3: + raise ValueError("Attention call does not expose Q, K, and V inputs") + query, key, value = args[offset : offset + 3] + output, probabilities = call.output + if output.shape == query.shape: + normalized_output = output + elif output.shape[-3:] == ( + query.shape[-2], + query.shape[-3], + query.shape[-1], + ): + normalized_output = output.transpose(-3, -2) + else: + raise ValueError("Unsupported per-head attention output layout") + return query, key, value, normalized_output, probabilities + + def _linear_model() -> NNsight: """A deterministic model where first-order attribution is exact.""" network = torch.nn.Sequential( @@ -29,6 +128,38 @@ def _linear_model() -> NNsight: return NNsight(network) +def _tiny_transformer(family): + with torch.random.fork_rng(): + torch.manual_seed(704) + if family == "gpt2": + config = GPT2Config( + vocab_size=32, + n_positions=16, + n_embd=8, + n_layer=2, + n_head=2, + ) + config._attn_implementation = "eager" + model = NNsight(GPT2LMHeadModel(config).eval()) + return model, model.transformer.h[0].attn + + if family == "pythia": + config = GPTNeoXConfig( + vocab_size=32, + max_position_embeddings=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + rotary_pct=0.5, + ) + config._attn_implementation = "eager" + model = NNsight(GPTNeoXForCausalLM(config).eval()) + return model, model.gpt_neox.layers[0].attention + + raise ValueError(f"Unsupported AtP* model family: {family}") + + def test_basic_atp_matches_exact_component_patching(): """AtP and exact patching agree when the downstream model is linear.""" model = _linear_model() @@ -98,3 +229,117 @@ def test_basic_atp_is_only_an_approximation_after_a_nonlinearity(): atp_effect = ((clean_component - noise_component) * noise_gradient).sum() exact_effect = patched_metric - noise_metric assert atp_effect.abs() < exact_effect.abs() * 0.01 + + +def test_qk_correction_matches_brute_force_attention_patches(): + """Vectorized Q/K corrections equal explicit one-node recomputation.""" + generator = torch.Generator().manual_seed(704) + shape = (2, 3, 4, 5) + noise_query = torch.randn(shape, generator=generator, dtype=torch.float64) + clean_query = torch.randn(shape, generator=generator, dtype=torch.float64) + noise_key = torch.randn(shape, generator=generator, dtype=torch.float64) + clean_key = torch.randn(shape, generator=generator, dtype=torch.float64) + noise_value = torch.randn(shape, generator=generator, dtype=torch.float64) + mask = torch.full((4, 4), float("-inf"), dtype=torch.float64).triu(1) + + noise_probabilities = _attention_probabilities(noise_query, noise_key, mask) + noise_output = torch.einsum("...qk,...kd->...qd", noise_probabilities, noise_value) + + query_delta = _query_output_delta( + clean_query, noise_key, noise_value, noise_output, mask + ) + for position in range(noise_query.shape[-2]): + patched_query = noise_query.clone() + patched_query[..., position, :] = clean_query[..., position, :] + patched_probabilities = _attention_probabilities(patched_query, noise_key, mask) + patched_output = torch.einsum( + "...qk,...kd->...qd", patched_probabilities, noise_value + ) + torch.testing.assert_close( + query_delta[..., position, :], + (patched_output - noise_output)[..., position, :], + ) + + key_delta = _key_output_delta( + noise_query, + clean_key, + noise_key, + noise_value, + noise_probabilities, + noise_output, + ) + for position in range(noise_key.shape[-2]): + patched_key = noise_key.clone() + patched_key[..., position, :] = clean_key[..., position, :] + patched_probabilities = _attention_probabilities(noise_query, patched_key, mask) + patched_output = torch.einsum( + "...qk,...kd->...qd", patched_probabilities, noise_value + ) + torch.testing.assert_close( + key_delta[..., position, :], patched_output - noise_output + ) + + +def test_key_correction_preserves_sub_epsilon_probabilities(): + """Small, nonzero probabilities must not be clamped before correction.""" + noise_query = torch.ones((1, 1, 1, 1), dtype=torch.float64) + noise_key = torch.tensor([[[[0.0], [-40.0]]]], dtype=torch.float64) + clean_key = noise_key.clone() + clean_key[..., 1, :] = 0 + noise_value = torch.tensor([[[[2.0], [-3.0]]]], dtype=torch.float64) + noise_probabilities = _attention_probabilities(noise_query, noise_key) + noise_output = torch.einsum("...qk,...kd->...qd", noise_probabilities, noise_value) + + key_delta = _key_output_delta( + noise_query, + clean_key, + noise_key, + noise_value, + noise_probabilities, + noise_output, + ) + patched_probabilities = _attention_probabilities(noise_query, clean_key) + patched_output = torch.einsum( + "...qk,...kd->...qd", patched_probabilities, noise_value + ) + + assert noise_probabilities[..., 1].item() < torch.finfo(torch.float64).eps + torch.testing.assert_close(key_delta[..., 1, :], patched_output - noise_output) + + +@pytest.mark.parametrize("family", ("gpt2", "pythia")) +def test_attention_source_adapter_matches_model_attention(family): + """Each MVP adapter reproduces the model's eager attention result.""" + input_ids = torch.tensor([[1, 2, 3, 4]]) + model, attention = _tiny_transformer(family) + with model.trace(input_ids): + call = _attention_source_call(attention, family) + query, key, value, attention_output, probabilities = _unpack_attention_call( + call + ) + query.requires_grad_(True) + saved_query = query.save() + saved_key = key.save() + saved_value = value.save() + saved_output = attention_output.save() + saved_probabilities = probabilities.save() + metric = model.output.logits.float().square().mean() + with metric.backward(): + query_gradient = query.grad.clone().save() + + causal_mask = torch.full((4, 4), float("-inf")).triu(1) + recomputed_probabilities = _attention_probabilities( + saved_query, saved_key, causal_mask + ) + recomputed_output = torch.einsum( + "...qk,...kd->...qd", recomputed_probabilities, saved_value + ) + expected = (1, 2, 4, 4) + assert saved_query.shape == expected + assert saved_key.shape == expected + assert saved_value.shape == expected + assert saved_probabilities.shape == (1, 2, 4, 4) + torch.testing.assert_close(saved_probabilities, recomputed_probabilities) + torch.testing.assert_close(saved_output, recomputed_output) + assert torch.isfinite(query_gradient).all() + assert torch.count_nonzero(query_gradient) From 13937e9858116223f6eefd24f93fca15eda7c250 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 29 Aug 2026 15:18:23 +0530 Subject: [PATCH 3/6] test: cover AtP* GradDrop and subset diagnostics --- tests/test_atp_star.py | 110 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/tests/test_atp_star.py b/tests/test_atp_star.py index 9ca3350d..d2e74ec9 100644 --- a/tests/test_atp_star.py +++ b/tests/test_atp_star.py @@ -1,6 +1,7 @@ """Correctness tests for the public-API AtP* research pattern.""" from collections import OrderedDict + import torch import pytest from transformers import ( @@ -9,9 +10,25 @@ GPTNeoXConfig, GPTNeoXForCausalLM, ) + from nnsight import NNsight +class _CancellationModel(torch.nn.Module): + """A direct path and an equal, opposite indirect path.""" + + def __init__(self): + super().__init__() + self.component = torch.nn.Identity() + self.indirect = torch.nn.Linear(1, 1, bias=False) + with torch.no_grad(): + self.indirect.weight.fill_(-1.0) + + def forward(self, inputs): + component = self.component(inputs) + return component + self.indirect(component) + + def _attention_probabilities(query, key, mask=None, scale=None): scores = torch.einsum("...qd,...kd->...qk", query, key) if scale is None: @@ -104,6 +121,30 @@ def _unpack_attention_call(call): return query, key, value, normalized_output, probabilities +def _aggregate_graddrop(drop_estimates): + """Equation 11: sum absolute layer-drop estimates and divide by L - 1.""" + layers = drop_estimates.shape[0] + if layers < 2: + raise ValueError("GradDrop requires at least two residual layers") + return drop_estimates.abs().sum(dim=0) / (layers - 1) + + +def _subset_statistics(masks, effects): + """Algorithm 1 statistics for included and excluded node subsets.""" + if masks.ndim != 2 or effects.shape != masks.shape[:1]: + raise ValueError("Expected masks [sample, node] and effects [sample]") + membership = torch.stack((masks, ~masks)) + counts = membership.sum(dim=1) + if torch.any(counts < 2): + raise ValueError("Each node needs two included and two excluded samples") + + expanded_effects = effects[None, :, None] + means = (expanded_effects * membership).sum(dim=1) / counts + centered = expanded_effects - means[:, None, :] + variances = (centered.square() * membership).sum(dim=1) / (counts - 1) + return counts, means, variances + + def _linear_model() -> NNsight: """A deterministic model where first-order attribution is exact.""" network = torch.nn.Sequential( @@ -307,6 +348,29 @@ def test_key_correction_preserves_sub_epsilon_probabilities(): torch.testing.assert_close(key_delta[..., 1, :], patched_output - noise_output) +def test_graddrop_exposes_a_cancellation_hidden_component(): + """Dropping an indirect residual gradient reveals the direct path.""" + model = NNsight(_CancellationModel()) + inputs = torch.tensor([[3.0]]) + + with model.trace(inputs): + component = model.component.output + component.requires_grad_(True) + indirect = model.indirect.output + indirect.requires_grad_(True) + metric = model.output.sum() + + with metric.backward(retain_graph=True): + standard_gradient = component.grad.clone().save() + + with metric.backward(): + indirect.grad = torch.zeros_like(indirect.grad) + dropped_gradient = component.grad.clone().save() + + torch.testing.assert_close(standard_gradient, torch.zeros_like(standard_gradient)) + torch.testing.assert_close(dropped_gradient, torch.ones_like(dropped_gradient)) + + @pytest.mark.parametrize("family", ("gpt2", "pythia")) def test_attention_source_adapter_matches_model_attention(family): """Each MVP adapter reproduces the model's eager attention result.""" @@ -343,3 +407,49 @@ def test_attention_source_adapter_matches_model_attention(family): torch.testing.assert_close(saved_output, recomputed_output) assert torch.isfinite(query_gradient).all() assert torch.count_nonzero(query_gradient) + + +def test_graddrop_aggregation_matches_equation_11(): + drop_estimates = torch.tensor( + [ + [1.0, -2.0, 0.0], + [-3.0, 4.0, 2.0], + [2.0, -1.0, -4.0], + ] + ) + expected = torch.tensor([3.0, 3.5, 3.0]) + torch.testing.assert_close(_aggregate_graddrop(drop_estimates), expected) + + +def test_subset_statistics_recover_additive_node_effects(): + """Algorithm 1 is exact on a balanced, additive subset experiment.""" + nodes = 4 + integers = torch.arange(2**nodes) + bit_positions = torch.arange(nodes) + masks = integers[:, None].bitwise_and(1 << bit_positions).bool() + contributions = torch.tensor([3.0, -2.0, 0.5, 4.0]) + effects = masks.float() @ contributions + + counts, means, variances = _subset_statistics(masks, effects) + estimates = means[0] - means[1] + + assert torch.equal(counts, torch.full_like(counts, 2 ** (nodes - 1))) + torch.testing.assert_close(estimates, contributions) + assert torch.isfinite(variances).all() + assert torch.all(variances >= 0) + + +def test_atp_star_helpers_reject_invalid_inputs(): + with pytest.raises(ValueError, match="at least two residual layers"): + _aggregate_graddrop(torch.ones((1, 3))) + + with pytest.raises(ValueError, match="Expected masks"): + _subset_statistics(torch.ones((2, 2), dtype=torch.bool), torch.ones((2, 1))) + + masks = torch.tensor([[True, False], [True, False]]) + with pytest.raises(ValueError, match="two included and two excluded"): + _subset_statistics(masks, torch.ones(2)) + + unsupported_attention = type("UnsupportedAttention", (), {"source": object()})() + with pytest.raises(ValueError, match="Unsupported pythia attention source layout"): + _attention_source_call(unsupported_attention, "pythia") From ae3c031a1bef00894fa1712acd0631d23fdbf8de Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 29 Aug 2026 15:18:33 +0530 Subject: [PATCH 4/6] docs: add tested AtP* building blocks --- docs/patterns/atp-star.md | 277 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 277 insertions(+) create mode 100644 docs/patterns/atp-star.md diff --git a/docs/patterns/atp-star.md b/docs/patterns/atp-star.md new file mode 100644 index 00000000..48ac574d --- /dev/null +++ b/docs/patterns/atp-star.md @@ -0,0 +1,277 @@ +--- +title: AtP* Building Blocks +one_liner: Build attribution patching with local Q/K corrections, GradDrop cancellation checks, and exact verification. +tags: [pattern, interpretability, gradients, attribution, patching, attention] +related: [docs/patterns/attribution-patching.md, docs/patterns/activation-patching.md, docs/usage/backward-and-grad.md, docs/usage/source.md] +sources: [tests/test_atp_star.py] +--- + +# AtP* Building Blocks + +> **Scope:** This page provides individually tested building blocks, not a +> drop-in end-to-end AtP* runner. The executable GPT-2/Pythia benchmark and +> Student-t confidence calculation are tracked as follow-up work in issue #704. +> Use the pieces below only after defining a model-specific component set, +> clean/noise direction, and scalar metric. + +## What this is for + +Attribution patching (AtP) cheaply ranks components using a first-order +approximation. AtP* addresses two ways that approximation can hide important +components: + +- **Saturated attention softmax:** gradients through Q/K can be tiny even when + patching Q or K would change the attention pattern substantially. +- **Residual cancellation:** direct and indirect gradient paths can cancel at a + component even when either path has a large effect. + +AtP* remains a screening method. Verify top candidates with exact activation +patching before making a causal claim. + +## Cost summary + +- Basic AtP needs a clean forward, a noise forward, and one backward pass. +- QK correction locally recomputes attention effects without another model pass. +- GradDrop uses one modified backward pass per residual layer instead of one + ordinary backward pass. +- Exact verification still needs one patched forward per tested component. + +## 1. Establish exact and AtP baselines + +Capture clean activations, then noise activations and their metric gradients: + +```python +with model.trace(clean_prompt): + clean_activation = component.output.save() + +with model.trace(noise_prompt): + noise_ref = component.output + noise_ref.requires_grad_(True) + noise_activation = noise_ref.save() + metric = metric_fn(model.lm_head.output[:, -1]) + noise_metric = metric.save() + with metric.backward(): + noise_gradient = noise_ref.grad.clone().save() + +atp = (clean_activation - noise_activation) * noise_gradient +``` + +Sum only over the dimensions inside one candidate component. For example, a +per-token residual score sums over hidden width but not sequence position. + +For exact verification, repeat the noise trace and patch only one candidate: + +```python +with model.trace(noise_prompt): + component.output[:, token] = clean_activation[:, token] + patched_metric = metric_fn(model.lm_head.output[:, -1]).save() + +exact_effect = patched_metric - noise_metric +``` + +Use exact effects to compute Recall@K, rank correlation, and recovered causal +effect for the approximate ranking. + +## 2. Access attention Q, K, V, and probabilities + +With eager attention, GPT-2 and Pythia expose post-transform tensors through +`.source`. Transformers 4.x and 5.x use different operation names, argument +offsets, and output layouts, so resolve and normalize them explicitly: + +```python +import torch + +def attention_source_call(attention, family): + names = { + "gpt2": ("attention_interface_0",), + "pythia": ("attention_interface_0", "unknown_0"), + } + if family not in names: + raise ValueError(f"Unsupported AtP* model family: {family}") + for name in names[family]: + try: + return getattr(attention.source, name) + except AttributeError: + pass + raise ValueError(f"Unsupported {family} attention source layout") + +def unpack_attention_call(call): + args, _ = call.inputs + offset = 0 if torch.is_tensor(args[0]) else 1 + query, key, value = args[offset : offset + 3] + output, probabilities = call.output + if output.shape == query.shape: + normalized_output = output + elif output.shape[-3:] == ( + query.shape[-2], query.shape[-3], query.shape[-1] + ): + normalized_output = output.transpose(-3, -2) + else: + raise ValueError("Unsupported per-head attention output layout") + return query, key, value, normalized_output, probabilities + +with model.trace(prompt): + call = attention_source_call(attention, family) + query, key, value, attention_output, probabilities = unpack_attention_call(call) +``` + +Use `model.transformer.h[layer].attn` for GPT-2 and +`model.gpt_neox.layers[layer].attention` for Pythia. The test suite covers both +legacy Transformers 4.48 and the current unified attention interface. Inspect +`attention.source` and fail explicitly for any unrecognized implementation. + +## 3. Correct query attribution + +Patch each clean query locally, recompute its softmax row against noise keys, and +compare the resulting per-head output with the noise output: + +```python +def attention_probabilities(query, key, mask=None, scale=None): + scores = torch.einsum("...qd,...kd->...qk", query, key) + if scale is None: + scale = query.shape[-1] ** -0.5 + scores = scores * scale + if mask is not None: + scores = scores + mask + return scores.softmax(dim=-1) + +def query_output_delta( + clean_query, noise_key, noise_value, noise_output, mask, scale=None +): + probability = attention_probabilities(clean_query, noise_key, mask, scale) + patched_output = torch.einsum("...qk,...kd->...qd", probability, noise_value) + return patched_output - noise_output +``` + +Dot each local output delta with the noise-run gradient at the corresponding +per-head attention output. This keeps the exact local softmax change while +linearizing only the downstream network. + +## 4. Correct key attribution in O(T²D) + +Changing one key changes one logit in every applicable query row. The exact +local output delta can be computed without materializing every patched T-by-T +attention matrix: + +```python +def key_output_delta( + q, clean_k, noise_k, value, probability, output, scale=None +): + delta = torch.einsum("...qd,...kd->...qk", q, clean_k - noise_k) + if scale is None: + scale = q.shape[-1] ** -0.5 + delta = delta * scale + + log_odds = torch.log(probability) - torch.log1p(-probability) + patched_p = torch.sigmoid(log_odds + delta) + denominator = torch.where( + probability == 1, torch.ones_like(probability), 1 - probability + ) + probability_scale = (patched_p - probability) / denominator + probability_scale = torch.where( + probability == 1, torch.zeros_like(probability_scale), probability_scale + ) + return probability_scale.unsqueeze(-1) * ( + value.unsqueeze(-3) - output.unsqueeze(-2) + ) +``` + +The result has shape `[batch, head, query, patched_key, head_dim]`. Dot it with +the per-head output gradient and sum over query and head-dimension to score each +key. `tests/test_atp_star.py` checks this formula against explicit single-key +patches, including a causal mask. + +## 5. Detect residual cancellation with GradDrop + +Run backward repeatedly from the same retained graph. On pass L, replace the +gradient entering residual contribution L with zero before reading upstream +component gradients: + +```python +with metric.backward(retain_graph=True): + ordinary_gradient = component_ref.grad.clone().save() + +with metric.backward(): + residual_contribution.grad = torch.zeros_like(residual_contribution.grad) + dropped_gradient = component_ref.grad.clone().save() +``` + +The residual contribution must execute downstream of the component being +scored so its gradient is encountered first during backward. Repeat for each +layer and aggregate exactly as specified by the AtP* paper. The test suite +includes a direct-plus-indirect cancellation model where ordinary AtP is zero +and the dropped gradient is nonzero. + +For L residual layers, Equation 11 aggregates the per-layer estimates as: + +```python +graddrop_score = drop_estimates.abs().sum(dim=0) / (num_layers - 1) +``` + +Compute this per clean/noise pair before averaging over the prompt distribution. +The L/(L-1) scaling preserves the direct path's expected contribution. + +## 6. Verify and diagnose + +1. Rank candidates with AtP*. +2. Patch the top K candidates individually and record exact effects. +3. On the unverified remainder, sample complementary Bernoulli subsets and + patch each subset jointly. +4. Use the paper's paired subset statistics and Welch bound to report an upper + confidence bound on a missed component's effect. + +For each node, Algorithm 1's point estimate is: + +```python +subset_estimate = abs(mean_effect_when_included - mean_effect_when_excluded) +``` + +Track included/excluded count, mean, and unbiased sample variance online for +every node. At least two samples are required in both groups. The confidence +diagnostic additionally needs a Student-t CDF; keep that dependency in the +benchmark environment rather than adding it to nnsight core. + +Do not substitute a generic bootstrap or independent-node assumption for the +paper's diagnostic: interactions between jointly patched components are a stated +limitation. + +## MVP benchmark protocol + +Use the same prompt pairs, metric, component definitions, and verification +budget for every method: + +- Models: `EleutherAI/pythia-70m-deduped` and `openai-community/gpt2`. +- Runtime: eager attention; float32 on CPU; record exact model revisions. +- Tasks: IOI-style name pairs and factual city completions. +- Methods: exact activation patching, AtP, AtP+QK, and AtP* with GradDrop. +- Quality: Recall@K, Spearman rank correlation, and exact effect recovered by + the top-K ranking. +- Cost: model forward/backward equivalents, wall time, and peak memory. +- Reproducibility: seed, prompt pairs, scores, exact effects, and sampled masks. + +Treat model-family comparisons as separate results if tokenization or component +counts differ. Do not pool ranks across incompatible node sets. + +## Gotchas + +- Force float32 for Pythia on CPU; checkpoint-default float16 can produce NaNs. +- Put every model in evaluation mode. Attention dropout invalidates the local + probability identities used above. +- Q/K correction is exact only for the local attention-softmax change. The + downstream score is still first-order. +- Treat masked probabilities of exactly zero and one explicitly without + clamping valid small nonzero probabilities. +- Q, K, and V must use the same RoPE state, mask, and scaling as the model + forward. Pass a non-default `scale` explicitly when the architecture does. +- Fused attention may not expose probabilities; use eager attention for the + reference workflow rather than silently changing semantics. +- Access gradients in reverse forward order inside each backward session. + +## Related + +- [attribution-patching](attribution-patching.md) +- [activation-patching](activation-patching.md) +- [attention-patterns](attention-patterns.md) +- [backward-and-grad](../usage/backward-and-grad.md) +- [source](../usage/source.md) +- Kramár et al. (2024), [AtP*](https://arxiv.org/abs/2403.00745) From a8aed7cc25c0e22a86ab2126b051add5364e4aec Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 29 Aug 2026 15:18:43 +0530 Subject: [PATCH 5/6] docs: link the AtP* guide from related resources --- CLAUDE.md | 1 + docs/patterns/attribution-patching.md | 3 ++- docs/patterns/index.md | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index 3862974b..27df028d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,6 +46,7 @@ If you're new to nnsight, read [docs/concepts/index.md](docs/concepts/index.md) - [docs/usage/backward-and-grad.md](docs/usage/backward-and-grad.md) — `with tensor.backward():` - [docs/patterns/gradient-based-attribution.md](docs/patterns/gradient-based-attribution.md) - [docs/patterns/attribution-patching.md](docs/patterns/attribution-patching.md) +- [docs/patterns/atp-star.md](docs/patterns/atp-star.md) — Q/K correction, GradDrop, and exact-verification building blocks ### "My model is too big for one GPU" - [docs/models/tensor-parallel.md](docs/models/tensor-parallel.md) — `transformers` tensor parallelism: `distributed_config=DistributedConfig(tp_size=N)` under `torchrun`, sharded activations gathered so the trace reads as it would on one GPU diff --git a/docs/patterns/attribution-patching.md b/docs/patterns/attribution-patching.md index b4517c8d..34610dca 100644 --- a/docs/patterns/attribution-patching.md +++ b/docs/patterns/attribution-patching.md @@ -2,7 +2,7 @@ title: Attribution Patching one_liner: Linear approximation of activation patching - one clean forward, one corrupt forward+backward, then `(act_clean - act_corrupt) * grad_corrupt` per component. tags: [pattern, interpretability, gradients, attribution, patching] -related: [docs/usage/backward-and-grad.md, docs/patterns/activation-patching.md, docs/patterns/gradient-based-attribution.md] +related: [docs/usage/backward-and-grad.md, docs/patterns/activation-patching.md, docs/patterns/gradient-based-attribution.md, docs/patterns/atp-star.md] sources: [src/nnsight/intervention/backward.py, src/nnsight/intervention/envoy.py] --- @@ -184,6 +184,7 @@ See `docs/usage/session.md`. ## Related - [activation-patching](activation-patching.md) — the exact (slower) operation this approximates. +- [atp-star](atp-star.md) — Q/K and residual-cancellation corrections for reducing false negatives. - [gradient-based-attribution](gradient-based-attribution.md) - `docs/usage/backward-and-grad.md` - https://nnsight.net/notebooks/tutorials/attribution_patching/ diff --git a/docs/patterns/index.md b/docs/patterns/index.md index 7a000e34..26750c97 100644 --- a/docs/patterns/index.md +++ b/docs/patterns/index.md @@ -39,6 +39,7 @@ Multiple prompts, multiple invokes, attribution in one batch. - [multi-prompt-comparison](multi-prompt-comparison.md) — Multiple `tracer.invoke(...)` in one trace, empty invokes for batch-wide ops, and `tracer.barrier(n)` for cross-invoke sharing. - [attribution-patching](attribution-patching.md) — Linear approximation of activation patching from corrupt-run gradients times clean-vs-corrupt activation differences. +- [atp-star](atp-star.md) — Tested building blocks for Q/K correction, GradDrop, and exact verification of attribution-patching rankings. ## Gradients From cfbb35695fa40efd7bc7e5223d3b0730202aeaf4 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Fri, 4 Sep 2026 14:50:58 +0530 Subject: [PATCH 6/6] fix: update .source attention adapter for 0.8 naming On 0.8's transformers, GPT2Attention.forward and GPTNeoXAttention.forward resolve attention_interface through an extra default-then-override step before invoking it, shifting the real call from .source.attention_interface_0 to .source.attention_interface_2 for both families. Try attention_interface_2 first, falling back to the old name. --- docs/patterns/atp-star.md | 4 ++-- tests/test_atp_star.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/patterns/atp-star.md b/docs/patterns/atp-star.md index 48ac574d..8e79c316 100644 --- a/docs/patterns/atp-star.md +++ b/docs/patterns/atp-star.md @@ -83,8 +83,8 @@ import torch def attention_source_call(attention, family): names = { - "gpt2": ("attention_interface_0",), - "pythia": ("attention_interface_0", "unknown_0"), + "gpt2": ("attention_interface_2", "attention_interface_0"), + "pythia": ("attention_interface_2", "attention_interface_0", "unknown_0"), } if family not in names: raise ValueError(f"Unsupported AtP* model family: {family}") diff --git a/tests/test_atp_star.py b/tests/test_atp_star.py index d2e74ec9..354e44f5 100644 --- a/tests/test_atp_star.py +++ b/tests/test_atp_star.py @@ -85,8 +85,8 @@ def _key_output_delta( def _attention_source_call(attention, family): operation_names = { - "gpt2": ("attention_interface_0",), - "pythia": ("attention_interface_0", "unknown_0"), + "gpt2": ("attention_interface_2", "attention_interface_0"), + "pythia": ("attention_interface_2", "attention_interface_0", "unknown_0"), } if family not in operation_names: raise ValueError(f"Unsupported AtP* model family: {family}")