Verify against Anthropic's official jacobian-lens repo: estimator match, zero frequency in their code+data, add --skip_first/--source_mean mirror mode

1 files changed, 43 insertions, 11 deletions

diff --git a/src/jlens_v3.py b/src/jlens_v3.py
index 69299dc..96af86c 100644
--- a/src/jlens_v3.py
+++ b/src/jlens_v3.py
@@ -58,7 +58,8 @@ def make_batches(train_data, block_size, batch_size, n_prompts, device, seed=133
     return batches
 
 
-def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
+def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32,
+                           skip_first=0, source_mean=False):
     """
     Compute the faithful J-lens VECTORS (rows of W_U * J_l) directly.
 
@@ -69,6 +70,14 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
     probe with the unembedding rows W_U[k] (V=65 probes) — mathematically
     identical, ~6x cheaper. One batched autograd call per prompt per layer
     (is_grads_batched=True) computes all V token vectors at once.
+
+    Two estimator modes (mirrors Anthropic's official jacobian-lens repo):
+      source_mean=False (default): mean over (source, future) PAIRS, all
+        positions included — the paper-formula estimator.
+      source_mean=True + skip_first=N: cotangent placed only at target
+        positions in [N, T-1) and the mean is over VALID SOURCE positions,
+        excluding the first N (attention sinks) and the last (no next-token
+        target) — exactly the official repo's valid_position_mask / fit().
     """
     d = model.config.n_embd
     V = model.config.vocab_size
@@ -77,7 +86,7 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
     W_U = model.lm_head.weight.detach().float()  # (V, d), nanoGPT weight tying
 
     accum = torch.zeros(V, d, device='cpu')
-    n_pairs = 0
+    n_positions = 0
 
     for (x, y) in batches:
         B, T = x.shape
@@ -97,8 +106,16 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
 
         h_l = resid_captured['h_l']          # (B, T, d)
         h_final = resid_captured['h_final']  # (B, T, d)
-        # sum over all future positions t' (causal: t' < t contributes zero grad)
-        h_final_sum = h_final.sum(dim=1)     # (B, d)
+        if source_mean:
+            # official-repo valid positions: [skip_first, T-1)
+            pos = torch.arange(T, device=device)
+            valid = (pos >= skip_first) & (pos < T - 1)
+            n_valid = int(valid.sum().item())
+            h_final_sum = h_final[:, valid, :].sum(dim=1)  # (B, d)
+        else:
+            valid = None
+            n_valid = 0
+            h_final_sum = h_final.sum(dim=1)  # (B, d); t' < t gives zero grad
 
         # grad_outputs[v, b, :] = W_U[v] -> batched VJPs, chunked over vocab
         # (is_grads_batched with all V=65 probes at once OOMs the 4GB K2200)
@@ -110,18 +127,26 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
                 h_final_sum, h_l, grad_outputs=grad_outputs,
                 is_grads_batched=True,
                 retain_graph=(v0 + C < V))[0]  # (C, B, T, d)
-            accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1)
+            if source_mean:
+                assert valid is not None
+                grads_v = grads[:, :, valid, :]  # (C, B, n_valid, d)
+                accum[v0:v0 + C] += grads_v.detach().cpu().sum(dim=(1, 2))
+            else:
+                accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1)
             del grads, grad_outputs
             torch.cuda.empty_cache()
 
-        n_pairs += B * T * (T + 1) // 2      # valid (t, t' >= t) pairs
-        # NOTE: count = T(T+1)/2 per sequence (all source positions t and all
-        # futures t' >= t). Pairs with t' < t contribute zero gradient by
-        # causality, so summing over all t' and dividing by this count is exact.
+        if source_mean:
+            n_positions += B * n_valid
+        else:
+            # NOTE: count = T(T+1)/2 per sequence (all source positions t and all
+            # futures t' >= t). Pairs with t' < t contribute zero gradient by
+            # causality, so summing over all t' and dividing by this count is exact.
+            n_positions += B * T * (T + 1) // 2
         del logits, loss, h_l, h_final, h_final_sum
         torch.cuda.empty_cache()
 
-    accum /= n_pairs
+    accum /= n_positions
     return accum  # (V, d): faithful J-lens vectors, rows of W_U * J_l
 
 
@@ -176,6 +201,12 @@ def main():
                     help='directory for per-layer result artifacts')
     ap.add_argument('--device', default='cuda')
     ap.add_argument('--chunk', type=int, default=32)
+    ap.add_argument('--skip_first', type=int, default=0,
+                    help='exclude first N source/target positions (mirrors '
+                         'Anthropic official repo skip=16 attention sinks)')
+    ap.add_argument('--source_mean', action='store_true',
+                    help='mean over valid source positions (official-repo '
+                         'estimator) instead of (source, future) pairs')
     args = ap.parse_args()
 
     device = args.device
@@ -209,7 +240,8 @@ def main():
         print(f"\n--- Layer {layer_idx} ---")
         print("  Computing faithful J-lens vectors (W_U-probed VJPs)...")
         faithful_vecs = compute_faithful_jlens(model, layer_idx, batches, device,
-                                               args.chunk)  # (V, d)
+                                               args.chunk, args.skip_first,
+                                               args.source_mean)  # (V, d)
         faithful_norms = {k: faithful_vecs[k].norm().item() for k in range(V)}
 
         if layer_idx == n_layer - 1: