Add faithful J-lens (jlens_v3): W_U-probed residual Jacobian per paper; both-ways comparison vs log-softmax proxy; 3-model adversarial reviews

6 files changed, 550 insertions, 0 deletions

diff --git a/check_proxy.py b/check_proxy.py
new file mode 100644
index 0000000..83d1602
--- /dev/null
+++ b/check_proxy.py
@@ -0,0 +1,41 @@
+"""Verify batched proxy == per-token proxy on the same model (numerical agreement)."""
+import sys, os
+sys.path.insert(0, '.')
+sys.path.insert(0, 'src')
+import torch, numpy as np
+from model import GPT, GPTConfig
+import jlens_v3
+
+torch.manual_seed(0)
+cfg = GPTConfig(n_layer=3, n_head=4, n_embd=16, block_size=32, bias=False,
+                vocab_size=65, dropout=0.0)
+model = GPT(cfg).eval()
+data = np.random.randint(0, 65, 2000).astype(np.uint16)
+batches = jlens_v3.make_batches(data, 32, 4, 3, 'cpu')
+
+# batched (new) proxy
+batched = jlens_v3.compute_proxy_norms(model, 1, batches, 'cpu', 65)
+
+# per-token (old) proxy, inline
+d = model.config.n_embd
+accum = torch.zeros(65, d)
+counts = torch.zeros(65)
+for (x, y) in batches:
+    B, T = x.shape
+    rc = {}
+    def hook(m, i, o): rc['val'] = o
+    h = model.transformer.h[1].register_forward_hook(hook)
+    logits, loss = model(x, y)
+    h.remove()
+    lp = torch.nn.functional.log_softmax(logits, dim=-1)
+    for k in range(65):
+        g = torch.autograd.grad(lp[:, :, k].sum(), rc['val'],
+                                retain_graph=(k < 64))[0]
+        accum[k] += g.detach().cpu().reshape(-1, d).sum(dim=0)
+        counts[k] += B * T
+per_token = {k: (accum[k] / counts[k]).norm().item() for k in range(65)}
+
+maxdiff = max(abs(batched[k] - per_token[k]) for k in range(65))
+print(f"max |batched - per_token| over 65 tokens: {maxdiff:.6e}")
+assert maxdiff < 1e-4, "proxy implementations disagree!"
+print("PROXY VARIANT AGREEMENT: PASSED")
diff --git a/reviews/2026-07-31-claude-opus-4.6.md b/reviews/2026-07-31-claude-opus-4.6.md
new file mode 100644
index 0000000..2ff8227
--- /dev/null
+++ b/reviews/2026-07-31-claude-opus-4.6.md
@@ -0,0 +1,38 @@
+# Review: Claude Opus 4.6 (agy, 2026-07-31)
+
+Adversarial review of the J-space replication project. Full brief: /tmp/jspace_brief.md.
+
+## Headline
+
+**The implementation computes the wrong quantity.**
+
+Anthropic's J-lens vectors are rows of `W_U · E[∂h_final/∂h_ℓ]` — the Jacobian stops at
+the final *residual stream*, before softmax. Our code (jlens_v2.py) differentiates through
+`log_softmax`, which folds in a `(1 − p(k))` factor that mechanically anticorrelates norm
+with frequency. The r = −0.65 may be an artefact of this difference, not a property of the
+model's representations.
+
+## Q1
+
+The 67% causal drop and the min(V,d) rank law both follow from the softmax gradient
+mechanics and linear algebra respectively — neither requires a "workspace" explanation.
+The causal experiment doesn't control for the softmax saturation confound.
+
+## Q2
+
+The cheapest kill-or-save experiment: compute J-lens BOTH ways (our `∇ log p` vs
+Anthropic-faithful `∇ h_final` composed with `W_U`) on the *existing* trained model.
+If the correlation vanishes with the faithful method, the thesis is dead.
+~30 min of compute, zero retraining.
+
+## Q3
+
+gpt2_jlens.py also has a norm-averaging bug: it accumulates `E[‖∇‖]` (average of norms)
+rather than `‖E[∇]‖` (norm of average), which are different quantities by Jensen's
+inequality.
+
+## Q4
+
+Frame as "open confounds to control," not "refutation." We haven't faithfully replicated
+their method, and we haven't addressed any of their functional experiments (steering,
+verbal report, reasoning ablation).
diff --git a/reviews/2026-07-31-gemini-3.1-pro.md b/reviews/2026-07-31-gemini-3.1-pro.md
new file mode 100644
index 0000000..20ad69d
--- /dev/null
+++ b/reviews/2026-07-31-gemini-3.1-pro.md
@@ -0,0 +1,97 @@
+# Review: Gemini 3.1 Pro (agy, 2026-07-31)
+
+Adversarial review of the J-space replication project. Full brief: /tmp/jspace_brief.md.
+
+## Q1. Falsification & alternative explanations
+
+The core claim — that Anthropic's "limited capacity workspace" is merely an artifact of
+unigram token frequency and the min(V, d) linear algebra rank bound — is falsified by a
+fundamental misunderstanding of both the gradient mechanics and space dimensionality.
+
+### Why the 67% causal drop occurs (softmax gradient mechanics, not workspace falsification)
+
+The J-lens vector at the output layer L for token k is defined as the expectation over
+contexts x:
+
+    J_L(k) = E_x[ grad_{h_L} log p(k|x) ] = W_{U,k}^T - sum_j E_x[p(j|x)] W_{U,j}^T
+           = W_{U,k}^T - W_avg
+
+where W_avg is the unigram probability-weighted average unembedding vector across the corpus.
+
+When 'q' frequency is doubled by inserting it at random positions (freq_experiment.py):
+1. Unigram baseline elevation: p_q increases across all context positions. In the
+   subtraction term -sum_j p(j|x) W_{U,j}, the component -p(q|x) W_{U,q} grows larger
+   across every single non-'q' position in the batch, directly canceling out W_{U,q}.
+2. Logit bias inflation: to minimize cross-entropy loss on uniform random insertions,
+   the model increases the scalar output bias b_q, raising p(q|x) globally and driving
+   (1 - p(q|x)) -> 0 faster during evaluation.
+
+Furthermore, freq_ablation.py compares a newly trained model against a pre-existing
+checkpoint (out-shakespeare-char/ckpt.pt) trained by a different script with unmatched
+seeds and training iterations.
+
+### The min(V, d) rank fallacy
+
+A matrix V in R^(V x d) has an absolute mathematical rank ceiling of min(V, d).
+
+In our setup (V=65):
+- d=16: rank 14-15 (~94% of the min(65,16)=16 ceiling)
+- d=32: rank 26-31 (~90% of 32)
+- d=128 and d=384: rank 65 (~100% of 65)
+
+Our results show nanoGPT J-vectors occupy almost 100% of the ambient dimension available.
+There is ZERO low-rank workspace compression in the model. In contrast, Anthropic
+evaluated models where V=50,257 and d=768 or 4,096. If Anthropic's finding were a
+min(V,d) linear algebra triviality, their J-space effective rank would be
+min(50257, 768) = 768. Instead they observed a rank of 10-50 << 768 << V.
+
+Our experiment proved that nanoGPT fails to form a compressed J-space, not that
+Anthropic's compression finding is a linear algebra illusion.
+
+## Q2. Most informative single experiment
+
+Frequency-Matched Synthetic Pair Test (Contextual Predictability vs Unigram Frequency).
+
+Design: train nanoGPT (or evaluate GPT-2) on a corpus containing two synthetic tokens,
+T_struct and T_noise, with IDENTICAL unigram frequencies (e.g. exactly 0.1% each):
+- T_struct: appears strictly in specific structured syntactic templates (e.g. after a
+  fixed 3-token trigger sequence A B C -> T_struct)
+- T_noise: injected at uniform random positions
+
+Expected outcomes:
+- Under frequency-only hypothesis: identical J-lens norms across all layers.
+- Under workspace/verbalizable-representation hypothesis: T_struct maintains high
+  J-lens norm in intermediate layers; T_noise collapses to near zero.
+
+## Q3. Implementation & methodological audit
+
+1. NORM OF EXPECTATION vs EXPECTATION OF NORM (critical bug):
+   - jlens_v2.py computes ||E_x[grad]||: averages gradient vectors first, then takes
+     L2 norm. At the final layer this reduces to ||W_{U,k}^T - W_avg||, stripping all
+     context-dependent dynamic activation variance.
+   - gpt2_jlens.py computes E_x[||grad||]: takes the norm on each batch step before
+     accumulating.
+   - Comparing nanoGPT to GPT-2 compares two mathematically distinct quantities.
+
+2. Severe underpowering & silent exception suppression:
+   - gpt2_jlens.py: n_batches=3 with seq_len=32 -> only 96 token positions to estimate
+     gradients over a 50,257-token vocabulary.
+   - Lines 93-94: bare `except: pass` silently discards failed backward passes.
+
+3. Residual capture point:
+   - jlens_v2.py registers a forward hook on model.transformer.h[layer_idx]; in nanoGPT
+     this captures the block output AFTER both attention and MLP residual additions.
+     Verify layer indices match Anthropic's definition (pre-block vs post-block).
+
+## Q4. Scrutiny-surviving framing
+
+"When evaluating the Jacobian Lens (J-lens) on small character-level transformers
+(V=65), J-lens vector magnitudes exhibit a strong inverse correlation with unigram token
+frequency (r = -0.65), and gradient norm reductions can be induced by artificially
+inflating token priors. This highlights that raw J-lens norms in small-scale models are
+heavily confounded by static unembedding geometry (W_{U,k} - W_avg) and baseline unigram
+predictability. However, we do NOT claim to refute Anthropic's Global Workspace hypothesis
+or their low-rank J-space findings (10-50 active dimensions). Because V < d in our
+character-level baseline, J-vectors span the full ambient rank (~min(V, d)), demonstrating
+that toy models fail to exhibit the severe subspace compression (10 << d << V) observed in
+large language models rather than disproving its existence."
diff --git a/reviews/2026-07-31-gpt-5.6-terra.md b/reviews/2026-07-31-gpt-5.6-terra.md
new file mode 100644
index 0000000..46b5e39
--- /dev/null
+++ b/reviews/2026-07-31-gpt-5.6-terra.md
@@ -0,0 +1,104 @@
+# Review: GPT-5.6-Terra (codex, 2026-07-31)
+
+Adversarial review of the J-space replication project. Full brief: /tmp/jspace_brief.md.
+Codex read the repo and the paper (43,704 tokens used, no files changed).
+
+## Q1. Strongest alternative explanation
+
+The central comparison is not currently about Anthropic's J-lens. jlens_v2.py estimates
+
+    E[ grad_{h_l,t} log p(k | x, t') ]
+
+summed over future output positions, whereas the paper first estimates the token-
+independent matrix E[ ∂h_final,t' / ∂h_l,t ], then applies final normalization and the
+unembedding row for token k. That distinction is fatal for the frequency claim.
+
+The gradient of log probability contains a softmax/calibration term:
+
+    ∇ log p_k = ∇ z_k − Σ_j p_j ∇ z_j
+
+Thus its norm measures the local sensitivity of the *log probability* of token k,
+including output-head geometry, prediction confidence, final LayerNorm's state-dependent
+Jacobian, and cancellation across contexts. It is not a token's "amount in J-space."
+A token can have a small mean gradient because its effects vary in direction by context
+and cancel — not because it is less verbalizable or displaced by a capacity limit.
+
+The `q` intervention is especially confounded. It does not merely double a sufficient
+statistic while holding the data-generating problem fixed: it inserts `q` at random
+character positions, changing sequence length, every downstream absolute position, local
+n-grams, and the conditional distribution of both `q` and its neighbors. In Shakespeare,
+`q` is unusually structured; random inserted instances largely destroy that structure.
+The 67% drop could therefore reflect a different learned conditional-prediction circuit or
+gradient alignment, plus an unmatched training run — not frequency per se.
+
+The rank result is even less probative. A matrix of 65 token-indexed vectors in d
+dimensions necessarily has rank at most min(65,d). Observing near-maximal numerical rank
+under a 1%-of-top-singular-value cutoff shows that these particular gradient vectors are
+reasonably nondegenerate; it does not show that the model's workspace capacity equals that
+bound. Anthropic's capacity claim is about sparse nonnegative decomposition of
+*activations*, occupancy above random-direction controls, and explained variance at
+individual positions — not the global linear rank of the token-vector dictionary. The paper
+explicitly notes that the token vectors may span all of residual space; J-space is defined
+by sparse use, not a low-rank span.
+
+A further comparability problem: nanoGPT computes norm of the mean gradient, while
+gpt2_jlens.py averages norms of per-batch full tensors. Those answer different questions,
+so the two correlations cannot jointly support one mechanism.
+
+## Q2. Best single experiment
+
+Run a matched, multi-seed **loss-reweighting plus faithful-lens** experiment.
+
+Train paired models from identical initializations and identical minibatch order on the
+unchanged Shakespeare sequences. In one member of each pair, multiply cross-entropy terms
+whose target is `q` by 2; in the other, use ordinary loss. This changes the effective
+target frequency/importance without injecting malformed `q` contexts or shifting all later
+positions. Use at least 5-10 paired seeds. Also include a same-total-loss control that
+upweights randomly chosen non-`q` target positions.
+
+For each model, compute both:
+1. the present score, ||E ∇ log p(q)||; and
+2. a faithful token vector from the averaged final-residual Jacobian, followed by the
+   unembedding as in the paper.
+
+Report q probability, conditional entropy, unembedding-row norm, vector cosine similarity,
+and bootstrap confidence intervals. Cheap because it reuses the small-model setup and
+directly separates "measurement artifact" from learned representation.
+
+## Q3. Faithfulness of jlens_v2.py
+
+No: it is directionally related to a Jacobian method, but it is not faithful enough to
+validate numerical comparisons with the paper.
+
+Good news: the forward hook captures the output of the selected nanoGPT block (the
+post-attention/post-MLP residual stream), a reasonable source capture point. Summing
+gradients over source positions and all output positions also includes the causal
+future-position dependence the paper intends — masked-impossible pairs have zero gradient.
+
+The methodological error is the target. The paper defines one d_model x d_model average
+Jacobian from intermediate residual stream to FINAL residual stream, then reads it through
+the model's normal output operations. Our code differentiates log_softmax(logits) directly.
+That folds the final LayerNorm, unembedding, and token-dependent softmax subtraction into
+the object being averaged. Averaging after these nonlinear/token-dependent operations is
+not equivalent to averaging the residual Jacobian and then reading it out.
+
+Two lesser issues: dividing by B*T rather than the number of valid source-future pairs
+changes scale (though not within-run rankings at fixed sequence length); 10-20 random
+batches is a noisier approximation than the paper's corpus-scale averaging (precision, not
+core invalidation).
+
+## Q4. Defensible blog framing
+
+"We found that a simple gradient-of-log-probability proxy on a 10.6M character transformer
+is strongly associated with token frequency, and that its token-indexed gradient dictionary
+has the expected rank ceiling min(V,d_model). A random-insertion intervention is consistent
+with frequency or conditional-prediction structure affecting this proxy, but it does not
+isolate frequency, and our current estimator differs materially from Anthropic's
+residual-Jacobian J-lens. These results motivate a controlled replication using the
+faithful lens and activation-level sparse-occupancy tests."
+
+Do NOT claim to have falsified Anthropic's workspace evidence, shown that its capacity
+result is "just linear algebra," demonstrated a consciousness-relevant conclusion is wrong,
+or shifted any burden of proof. Anthropic's headline rests on functional interventions,
+sparse occupancy, variance controls, and broadcast/generalization tests in addition to
+rank; our current experiments test none of those.
diff --git a/smoke_test_v3.py b/smoke_test_v3.py
new file mode 100644
index 0000000..7e83e1c
--- /dev/null
+++ b/smoke_test_v3.py
@@ -0,0 +1,32 @@
+"""CPU smoke test for jlens_v3 (faithful J-lens machinery). Random tiny model."""
+import sys, os
+sys.path.insert(0, '.')
+sys.path.insert(0, 'src')
+import torch, numpy as np
+from model import GPT, GPTConfig
+import jlens_v3
+
+torch.manual_seed(0)
+
+cfg = GPTConfig(n_layer=3, n_head=4, n_embd=16, block_size=32, bias=False,
+                vocab_size=65, dropout=0.0)
+model = GPT(cfg).eval()
+
+data = np.random.randint(0, 65, 2000).astype(np.uint16)
+batches = jlens_v3.make_batches(data, 32, 4, 3, 'cpu')
+
+J = jlens_v3.compute_faithful_jlens(model, 1, batches, 'cpu', chunk=8)
+print("faithful vecs shape:", J.shape, "finite:", torch.isfinite(J).all().item())
+print("row norms (first 5):", [round(v, 4) for v in J.norm(dim=1)[:5].tolist()])
+
+pn = jlens_v3.compute_proxy_norms(model, 1, batches, 'cpu', 65)
+print("proxy norms finite:", all(np.isfinite(v) for v in pn.values()))
+
+fn = {k: J[k].norm().item() for k in range(65)}
+freq = np.bincount(data, minlength=65).astype(float)
+freq = freq / freq.sum() * 100
+r_p = np.corrcoef(np.array([pn[k] for k in range(65)]), freq)[0, 1]
+r_f = np.corrcoef(np.array([fn[k] for k in range(65)]), freq)[0, 1]
+print(f"random-init sanity: proxy r={r_p:+.3f} faithful r={r_f:+.3f} (should be ~0, model is random)")
+assert J.shape == (65, 16) and torch.isfinite(J).all()
+print("SMOKE TEST PASSED")
diff --git a/src/jlens_v3.py b/src/jlens_v3.py
new file mode 100644
index 0000000..7112cb1
--- /dev/null
+++ b/src/jlens_v3.py
@@ -0,0 +1,238 @@
+"""
+J-lens v3: FAITHFUL implementation of Anthropic's Jacobian lens.
+
+Paper: "Verbalizable Representations Form a Global Workspace in Language Models"
+https://transformer-circuits.pub/2026/workspace/index.html
+
+Definition (from the paper's Methods):
+    J_l = E_{t, t' >= t, prompt}[ d h_final,t' / d h_l,t ]      (d_model x d_model per layer)
+    J-lens vectors at layer l = rows of W_U * J_l
+    lens(h_l) = softmax(W_U * norm(J_l * h_l))
+
+The average is over source positions t, all future positions t' >= t, and a corpus
+of prompts. The result is a single d x d matrix per layer that maps the residual
+stream at layer l to the FINAL residual stream, read out through the model's own
+unembedding (W_U = lm_head.weight in nanoGPT).
+
+This is different from the earlier proxy (jlens_v2) which differentiated
+log_softmax(logits) directly, folding the token-dependent softmax calibration
+into the averaged object.
+
+Implementation: basis-chunked vector-Jacobian products. For a chunk C of basis
+vectors {e_c}, compute
+    S_c = sum_{t'} (h_final,t' . e_c)
+and backpropagate S_c to h_l (the captured residual stream at layer l).
+Because h_final,t' does not depend on h_l,t for t' < t (causality), summing over
+all t' yields exactly sum over valid pairs. One batched autograd call per chunk
+(is_grads_batched=True) gives all rows in the chunk.
+
+Also implements the BOTH-WAYS comparison: old proxy ||E[grad log p(k|x)]|| vs
+faithful ||(W_U J_l)[k]||, each correlated against token frequency, per layer.
+
+Usage (inside Docker on meru):
+    python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \
+        --data_dir data/shakespeare_char --n_prompts 20 --layers 2,3,4
+"""
+
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import argparse
+import torch
+import numpy as np
+import pickle
+
+import jlens_v2  # for load_model
+
+
+def make_batches(train_data, block_size, batch_size, n_prompts, device, seed=1337):
+    """Random contiguous prompts from train.bin (n_prompts prompts of batch_size)."""
+    g = torch.Generator().manual_seed(seed)
+    batches = []
+    for _ in range(n_prompts):
+        ix = torch.randint(len(train_data) - block_size, (batch_size,), generator=g)
+        x = torch.stack([torch.from_numpy(
+            train_data[i:i + block_size].astype(np.int64)) for i in ix])
+        y = torch.stack([torch.from_numpy(
+            train_data[i + 1:i + 1 + block_size].astype(np.int64)) for i in ix])
+        batches.append((x.to(device), y.to(device)))
+    return batches
+
+
+def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
+    """
+    Compute the faithful J-lens VECTORS (rows of W_U * J_l) directly.
+
+    (W_U J_l)[k] = E_{t, t'>=t, prompt}[ d(W_U[k] . h_final,t') / d h_l,t ]
+                 = E[ grad of the k-th RAW LOGIT (pre-softmax) w.r.t. h_l ]
+
+    Instead of materializing the d x d residual Jacobian (384 basis VJPs), we
+    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.
+    """
+    d = model.config.n_embd
+    V = model.config.vocab_size
+    n_layer = model.config.n_layer
+    last_block = model.transformer.h[n_layer - 1]
+    W_U = model.lm_head.weight.detach().float()  # (V, d), nanoGPT weight tying
+
+    accum = torch.zeros(V, d, device='cpu')
+    n_pairs = 0
+
+    for (x, y) in batches:
+        B, T = x.shape
+        resid_captured = {}
+
+        def hook(module, inp, out):
+            resid_captured['h_l'] = out
+
+        def hook_final(module, inp, out):
+            resid_captured['h_final'] = out
+
+        h1 = model.transformer.h[layer_idx].register_forward_hook(hook)
+        h2 = last_block.register_forward_hook(hook_final)
+        logits, loss = model(x, y)           # graph connected
+        h1.remove()
+        h2.remove()
+
+        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)
+
+        # grad_outputs[v, b, :] = W_U[v] -> batched VJPs for all V tokens
+        grad_outputs = W_U.to(device).unsqueeze(1).expand(V, B, d).contiguous()
+        grads = torch.autograd.grad(
+            h_final_sum, h_l, grad_outputs=grad_outputs,
+            is_grads_batched=True)[0]  # (V, B, T, d)
+        accum += grads.detach().cpu().reshape(V, -1, d).sum(dim=1)
+
+        n_pairs += B * T * (T + 1) // 2      # valid (t, t' >= t) pairs
+        del logits, loss, h_l, h_final, h_final_sum, grads, grad_outputs
+        torch.cuda.empty_cache()
+
+    accum /= n_pairs
+    return accum  # (V, d): faithful J-lens vectors, rows of W_U * J_l
+
+
+def compute_proxy_norms(model, layer_idx, batches, device, vocab_size):
+    """
+    Old proxy (jlens_v2 quantity): || E_x[ grad log_softmax(logits)[:, :, k] ] ||
+    averaged over positions — norm of the mean gradient, per token.
+    """
+    d = model.config.n_embd
+    accum = torch.zeros(vocab_size, d, device='cpu')
+    counts = torch.zeros(vocab_size, device='cpu')
+    for (x, y) in batches:
+        B, T = x.shape
+        resid_captured = {}
+
+        def hook(module, inp, out):
+            resid_captured['val'] = out
+
+        handle = model.transformer.h[layer_idx].register_forward_hook(hook)
+        logits, loss = model(x, y)
+        handle.remove()
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+        # sum over positions -> (V,); batched VJP over all vocab at once
+        lp_sum = log_probs.sum(dim=(0, 1))
+        grad_outputs = torch.eye(vocab_size, device=device)
+        grad = torch.autograd.grad(
+            lp_sum, resid_captured['val'], grad_outputs=grad_outputs,
+            is_grads_batched=True)[0]  # (V, B, T, d)
+        accum += grad.detach().cpu().reshape(vocab_size, -1, d).sum(dim=1)
+        counts += B * T
+        del logits, loss, log_probs, lp_sum, grad
+        torch.cuda.empty_cache()
+    norms = {}
+    for token_id in range(vocab_size):
+        norms[token_id] = (accum[token_id] / counts[token_id]).norm().item()
+    return norms
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--checkpoint', required=True)
+    ap.add_argument('--data_dir', default='data/shakespeare_char')
+    ap.add_argument('--n_prompts', type=int, default=20)
+    ap.add_argument('--batch_size', type=int, default=16)
+    ap.add_argument('--layers', default='2,3,4', help='comma-separated layer indices')
+    ap.add_argument('--device', default='cuda')
+    ap.add_argument('--chunk', type=int, default=32)
+    args = ap.parse_args()
+
+    device = args.device
+    model, config = jlens_v2.load_model(args.checkpoint, device)
+    model.eval()
+    d = config.n_embd
+    V = config.vocab_size
+    n_layer = config.n_layer
+    print(f"Model: {n_layer} layers, d={d}, V={V}, V/d={V/d:.2f}x")
+
+    train_data = np.memmap(f'{args.data_dir}/train.bin', dtype=np.uint16, mode='r')
+    with open(f'{args.data_dir}/meta.pkl', 'rb') as f:
+        meta = pickle.load(f)
+    itos = meta['itos']
+
+    # token frequencies
+    counts = np.bincount(train_data, minlength=V).astype(float)
+    freq = counts / counts.sum() * 100  # percent
+
+    batches = make_batches(train_data, config.block_size, args.batch_size,
+                           args.n_prompts, device)
+
+    layers = [int(l) for l in args.layers.split(',')]
+    W_U = model.lm_head.weight.detach().float()  # (V, d) — nanoGPT weight tying
+
+    print(f"\n{'='*78}")
+    print("BOTH-WAYS COMPARISON: old proxy vs faithful J-lens (rows of W_U * J_l)")
+    print(f"{'='*78}")
+
+    for layer_idx in layers:
+        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)
+        faithful_norms = {k: faithful_vecs[k].norm().item() for k in range(V)}
+
+        if layer_idx == n_layer - 1:
+            # Validation: J_{L-1} should be identity, so faithful vectors == W_U rows
+            sims = torch.nn.functional.cosine_similarity(
+                faithful_vecs.float(), W_U.float(), dim=1)
+            print(f"  [validation] last layer: mean cos-sim(faithful, W_U rows) = "
+                  f"{sims.mean().item():.4f} (expect ~1.0 if J=identity)")
+
+        print("  Computing old proxy norms...")
+        proxy_norms = compute_proxy_norms(model, layer_idx, batches, device, V)
+
+        # correlations with frequency
+        f_arr = np.array([freq[k] for k in range(V)])
+        p_arr = np.array([proxy_norms[k] for k in range(V)])
+        fa_arr = np.array([faithful_norms[k] for k in range(V)])
+        r_proxy = np.corrcoef(p_arr, f_arr)[0, 1]
+        r_faith = np.corrcoef(fa_arr, f_arr)[0, 1]
+
+        print(f"  freq corr | old proxy:    r = {r_proxy:+.3f}")
+        print(f"  freq corr | faithful:     r = {r_faith:+.3f}")
+
+        # top/bottom by faithful norm
+        srt = sorted(faithful_norms.items(), key=lambda kv: kv[1], reverse=True)
+        def esc(s):
+            return s.replace(chr(10), '\\n')
+        print("  Top 5 by faithful norm:")
+        for tid, n in srt[:5]:
+            print(f"    '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}")
+        print("  Bottom 5 by faithful norm:")
+        for tid, n in srt[-5:]:
+            print(f"    '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}")
+
+        torch.save({'faithful_vecs': faithful_vecs, 'faithful_norms': faithful_norms,
+                    'proxy_norms': proxy_norms},
+                   f'outputs/jlens_v3_layer{layer_idx}.pt')
+        print(f"  Saved outputs/jlens_v3_layer{layer_idx}.pt")
+
+    print("\nDONE.")
+
+
+if __name__ == '__main__':
+    main()