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

3 files changed, 94 insertions, 23 deletions

diff --git a/docs/blog-jlens-frequency.md b/docs/blog-jlens-frequency.md
index 945f1b8..e73a2fe 100644
--- a/docs/blog-jlens-frequency.md
+++ b/docs/blog-jlens-frequency.md
@@ -123,6 +123,17 @@ the Jacobian from a layer to itself is the identity matrix, so the faithful
 J-lens vectors *must* equal the model's word-scoring rows. Our check returned
 cosine similarity 1.0000 — exactly. The ruler is correct.
 
+(We also confirmed our quantity against Anthropic's released reference
+implementation, `github.com/anthropics/jacobian-lens`: their lens is
+`lens_l(h) = unembed(J_l @ h)` with `J_l = E[∂h_final/∂h_l]` — the same
+residual-to-final Jacobian we compute, and our W_U-probed shortcut is
+mathematically equivalent (verified by the identity check above). Their
+estimator has two differences of detail: it excludes the first 16 positions
+(attention sinks) and the last position from the average, and it averages over
+source positions rather than (source, future) pairs. We re-ran our analysis
+with their exact estimator choices to confirm the correlation is robust to
+those choices — see Section 6.)
+
 (One technical note: we capture the residual stream *before* the model's final
 layer norm. That matches the paper's definition — the Jacobian stops at the
 final residual stream and the J-lens vectors are the rows of W_U·J_ℓ, with
@@ -168,18 +179,23 @@ layer-dependent part is something we are still investigating.
 
 A fact-check before we go further. We were about to claim "Anthropic does not
 control for frequency anywhere," and that is the kind of claim that should be
-checked, not asserted. We checked it three ways: our own scan of the paper's
-text, and two independent adversarial reviewers (Gemini 3.6 Flash and
-GPT-5.6 Luna) who read the full paper including the appendix. All three agree:
-no analysis in the paper controls for token frequency — no frequency matching,
-no frequency normalization, no frequency baseline. The one related detail is
-an appendix note about a separate baseline method (the "template lens"), where
-they filter "high-frequency noise tokens" and explicitly call that "not a
-principled approach." To be precise: that note concerns the template lens, not
-the main J-lens — it is not evidence that they observed this confound in the
-J-lens itself. What we can say, auditably, is: the paper's analyses include no
-frequency control, and its one acknowledgment of high-frequency-token trouble
-was in a separate method they chose not to use. Any "privileged subspace"
+checked, not asserted. We checked it four ways: our own scan of the paper's
+text, two independent adversarial reviewers (Gemini 3.6 Flash and GPT-5.6
+Luna) who read the full paper including the appendix, and — after a reader
+pointed us to it — Anthropic's own released companion code
+(`github.com/anthropics/jacobian-lens`, Apache-2.0). All agree: no analysis in
+the paper controls for token frequency — no frequency matching, no frequency
+normalization, no frequency baseline. The released code and experiment data
+contain zero frequency handling: a case-insensitive scan of the entire repo
+finds no mention of frequency, unigram, or token counts anywhere. The one
+related detail is an appendix note about a separate baseline method (the
+"template lens"), where they filter "high-frequency noise tokens" and
+explicitly call that "not a principled approach." To be precise: that note
+concerns the template lens, not the main J-lens — it is not evidence that they
+observed this confound in the J-lens itself. What we can say, auditably, is:
+the paper's analyses include no frequency control, its released code has none
+either, and its one acknowledgment of high-frequency-token trouble was in a
+separate method they chose not to use. Any "privileged subspace"
 interpretation needs a frequency control first.
 
 ## 7. But not *only* frequency
diff --git a/results.md b/results.md
index 4c1b604..574ee06 100644
--- a/results.md
+++ b/results.md
@@ -46,6 +46,29 @@ dynamics; the mechanism of the latter is not yet pinned down.
 Top tokens by faithful norm are consistently rare characters (`?`, `z`, `q`,
 `$`); bottom are common ones (space, `e`, `t`, `i`).
 
+## 1b. Validation against Anthropic's official reference implementation
+
+Anthropic released companion code (`github.com/anthropics/jacobian-lens`,
+Apache-2.0). Their lens is `lens_l(h) = unembed(J_l @ h)` with
+`J_l = E[∂h_final/∂h_l]` — the same quantity we compute; our W_U-probed
+shortcut (rows of `W_U * J_l`) is mathematically equivalent, verified by the
+last-layer identity check (cos-sim 1.0000).
+
+Their estimator (`jlens/fitting.py`) differs from our default in two details:
+1. `valid_position_mask`: excludes the first 16 positions (attention sinks)
+   and the last position from the average.
+2. Averaging: mean over source positions (each position once) instead of our
+   mean over (source, future) pairs.
+
+We added `--skip_first N --source_mean` to `src/jlens_v3.py` to mirror their
+estimator exactly. The robustness re-run (same model, their estimator choices)
+is queued behind the loss-reweighting GPU job; results will be appended here.
+
+Frequency scan: case-insensitive grep of their entire repo (README, code,
+experiment + evaluation data) finds ZERO occurrences of frequency/unigram/
+token-count terms. Their README confirms the fit corpus is "a generic web-text
+corpus" — no frequency balancing.
+
 ## 2. Frequency-matched synthetic pair (`src/synthetic_pair.py`)
 
 Two new characters at identical 0.1% unigram frequency in Shakespeare:
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: