From 2a2ba68bc6e27cec27f6cab24ee6ee8000468f97 Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Mon, 14 Sep 2026 08:48:19 -0500 Subject: [PATCH] improved direct path patching test --- tests/unit/test_direct_path_patching.py | 43 +++++++++---------------- 1 file changed, 15 insertions(+), 28 deletions(-) diff --git a/tests/unit/test_direct_path_patching.py b/tests/unit/test_direct_path_patching.py index 2224a14b0..8f7b739b6 100644 --- a/tests/unit/test_direct_path_patching.py +++ b/tests/unit/test_direct_path_patching.py @@ -274,16 +274,9 @@ def test_clean_equals_corrupted_gives_zero_delta(self, tiny_model): class TestCorrectness: - def test_correctness_against_actual_ln_forward(self, tiny_model, tokens_and_caches): - """Logit-diff metric: linear-LN approximation should match actual LN within 1e-3. - - Folding LN does not make the approximation exact: LayerNormPre still divides - by an input-dependent norm, while the approximation reuses the corrupted - run's scale. The error depends on the weights, so the tolerance holds for - the seeded tiny_model but is not guaranteed for arbitrary initialisations. - Using logit diff (correct_tok - incorrect_tok) cancels the centering offset - introduced by process_weights_(). - """ + @pytest.mark.parametrize("component", ["q", "k", "v"]) + def test_matches_frozen_ln_scale_reference(self, tiny_model, tokens_and_caches, component): + """Matches a hand-written hook dividing delta_resid by the corrupted run's ln1 scale.""" _, corrupted_tokens, clean_cache, corrupted_cache = tokens_and_caches src_layer, src_head = 0, 0 dst_layer, dst_head = 2, 1 @@ -300,28 +293,22 @@ def logit_diff(logits): corrupted_z = corrupted_cache[f"blocks.{src_layer}.attn.hook_z"][:, :, src_head, :] delta_resid = (clean_z @ W_O[src_head]) - (corrupted_z @ W_O[src_head]) # type: ignore[index] - # Independent reference: patch through actual LayerNorm forward - corrupted_resid = corrupted_cache[f"blocks.{dst_layer}.hook_resid_pre"] - patched_resid = corrupted_resid + delta_resid - - with torch.no_grad(): - ln1 = tiny_model.blocks[dst_layer].ln1 # type: ignore[index] - patched_normed = ln1(patched_resid) - corrupted_normed = ln1(corrupted_resid) + # Not the true LN forward: its gap to the approximation is weight-dependent and, on + # tiny_model, about as large as the patch effect, so it can't catch scale bugs. + ln_scale = corrupted_cache[f"blocks.{dst_layer}.ln1.hook_scale"] + W_comp = getattr(tiny_model.blocks[dst_layer].attn, f"W_{component.upper()}")[dst_head] + ref_delta = (delta_resid / ln_scale) @ W_comp - W_Q_dst = tiny_model.blocks[dst_layer].attn.W_Q[dst_head] # type: ignore[index,union-attr] - true_delta_q = (patched_normed - corrupted_normed) @ W_Q_dst - - def true_hook(value, hook): + def ref_hook(value, hook): if value.requires_grad: value = value.clone() - value[:, :, dst_head, :] = value[:, :, dst_head, :] + true_delta_q + value[:, :, dst_head, :] = value[:, :, dst_head, :] + ref_delta return value with torch.no_grad(): ref_logits = tiny_model.run_with_hooks( corrupted_tokens, - fwd_hooks=[(f"blocks.{dst_layer}.attn.hook_q", true_hook)], + fwd_hooks=[(f"blocks.{dst_layer}.attn.hook_{component}", ref_hook)], ) ref_metric = logit_diff(ref_logits).item() @@ -334,14 +321,14 @@ def true_hook(value, hook): patching_metric=logit_diff, src_layer=src_layer, src_head=src_head, - component="q", + component=component, verbose=False, ) our_metric = results[dst_layer, dst_head].item() - assert abs(our_metric - ref_metric) < 1e-3, ( - f"Linear-LN approx {our_metric:.6f} disagrees with actual-LN ref {ref_metric:.6f} " - f"(diff={abs(our_metric - ref_metric):.2e})." + assert abs(our_metric - ref_metric) < 1e-5, ( + f"Direct-path patch {our_metric:.6f} disagrees with frozen-scale reference " + f"{ref_metric:.6f} (diff={abs(our_metric - ref_metric):.2e})." ) def test_all_sources_consistent_with_single(self, tiny_model, tokens_and_caches):