diff options
| -rw-r--r-- | docs/blog-jlens-frequency.md | 40 | ||||
| -rw-r--r-- | results.md | 23 | ||||
| -rw-r--r-- | src/jlens_v3.py | 54 |
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 @@ -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: |