improved direct path patching test

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):