Address Luna repo review: dynamic repo root (no hardcoded cwd), J-lens correctness test (last-layer identity), auditable fact-check, mechanism narrative fix, softened conditional-predictability claim, README data-prep + artifacts note

6 files changed, 114 insertions, 19 deletions

diff --git a/README.md b/README.md
index 155adde..ae9ba86 100644
--- a/README.md
+++ b/README.md
@@ -70,10 +70,17 @@ sh scripts/test.sh
 
 Skips gracefully where torch is unavailable.
 
-### 2. Train the base model
+### 2. Prepare data + train the base model
 
-Train nanoGPT on `data/shakespeare_char` (10.65M params, 6 layers, d=384,
-block 128) and keep the checkpoint at `out-shakespeare-char/ckpt.pt`:
+Prepare the character-level Shakespeare dataset (downloads Shakespeare and
+builds `data/shakespeare_char/{train,val}.bin` + `meta.pkl`):
+
+```sh
+python3 data/shakespeare_char/prepare.py
+```
+
+Then train nanoGPT (10.65M params, 6 layers, d=384, block 128) and keep the
+checkpoint at `out-shakespeare-char/ckpt.pt`:
 
 ```sh
 python3 train.py config/train_shakespeare_char.py
@@ -108,6 +115,14 @@ python3 src/loss_reweight.py --step train --mode ctrl_random --seed 0
 python3 src/loss_reweight.py --step summary --layers 2,3,4
 ```
 
+## Data & artifacts
+
+Raw data files (`.bin`/`.pkl`), checkpoints, and experiment outputs are NOT
+committed (gitignored — they are regenerable and large). Committed instead:
+`results.md` (the numbers) and `src/stats_decomp.py` (the statistics that
+reproduce them from the saved artifacts). To regenerate everything, follow the
+reproduction steps above; the full pipeline takes a few hours on a 4GB GPU.
+
 ## References
 
 - [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) — Anthropic, 2026
diff --git a/docs/blog-jlens-frequency.md b/docs/blog-jlens-frequency.md
index d45b183..945f1b8 100644
--- a/docs/blog-jlens-frequency.md
+++ b/docs/blog-jlens-frequency.md
@@ -69,6 +69,11 @@ lens is ranking words by the size of this gradient, the ranking is partly
 pre-written by the frequency distribution before the model even learns
 anything.
 
+(A note on that intuition: it applies directly to our first, simpler
+implementation, which differentiated through the softmax. With the faithful
+lens the mechanism is different — it turns out to live partly in the geometry
+of the word-scoring matrix itself. Section 6 has the full decomposition.)
+
 In other words: **a "privileged workspace" might just be a frequency effect
 wearing a fancy hat.**
 
@@ -167,11 +172,14 @@ 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 only place the word
-shows up is an appendix note about a separate baseline method (the "template
-lens"), where they filter "high-frequency noise tokens" — and they explicitly
-call that "not a principled approach," then never apply it to the main
-J-lens. They saw the effect. They didn't fix it. Any "privileged subspace"
+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"
 interpretation needs a frequency control first.
 
 ## 7. But not *only* frequency
@@ -198,10 +206,10 @@ training runs showed:
 The predictable token scores **~1.4-1.5x higher** than the noise token at
 identical frequency, in every layer of every seed. Middle-layer ratio across
 the three seeds: 1.47 ± 0.09, bootstrap 95% CI [1.37, 1.53] — entirely above
-1. So the lens is not a pure frequency meter. It genuinely responds to
-conditional predictability — which, honestly, is what "verbalizable" should
-mean. The J-lens measures *both*: a frequency prior that is never subtracted
-out, and a real structure signal on top of it.
+1. So the lens is not a pure frequency meter: at equal frequency, the two
+tokens differ in norm. Whether that difference is specifically *conditional
+predictability* (what "verbalizable" should mean) depends on a control that is
+still running — see the caveat below.
 
 One caveat, found by a reviewer: the noise token '#' was inserted at random
 character positions, which slices *inside* words ~95% of the time (th#e,
@@ -258,10 +266,10 @@ post.
 ## 11. How to reproduce everything
 
 All code, data-prep scripts, experiment scripts, tests, and this analysis live
-in the repository: [link to cgit]. Summary of results in `results.md`.
-Reproduction steps in the README. The only requirements are a Linux machine
-with Docker, a CUDA GPU (any modern card; we used a 4GB Quadro K2200), and the
-`pytorch/pytorch:2.4.1-cuda11.8` image.
+in the repository (URL to be added once the public git instance is live).
+Summary of results in `results.md`. Reproduction steps in the README. The only
+requirements are a Linux machine with Docker, a CUDA GPU (any modern card; we
+used a 4GB Quadro K2200), and the `pytorch/pytorch:2.4.1-cuda11.8` image.
 
 Run the test suite:
 ```
diff --git a/scripts/test.sh b/scripts/test.sh
index 0875608..480e49a 100755
--- a/scripts/test.sh
+++ b/scripts/test.sh
@@ -3,4 +3,5 @@
 # torch is not installed on bare-metal voidlaptop, so this skips gracefully
 # there; run it in the meru container (torch + CUDA) for the full pass:
 #   ssh meru "docker exec -w /workspace/code jspace-nanogpt sh scripts/test.sh"
-exec python3 tests/test_loss_reweight.py
+python3 tests/test_loss_reweight.py || exit 1
+python3 tests/test_jlens_v3.py || exit 1
diff --git a/src/loss_reweight.py b/src/loss_reweight.py
index f01d88a..08dc056 100644
--- a/src/loss_reweight.py
+++ b/src/loss_reweight.py
@@ -31,6 +31,7 @@ import torch.nn.functional as F
 DATA_DIR = 'data/shakespeare_char'
 OUT_ROOT = 'out-loss-reweight'
 JLENS_OUT = 'outputs/loss_reweight'
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 TARGET = 'q'
 WEIGHT = 2.0
 MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=128,
@@ -142,7 +143,7 @@ def jlens(mode, seed, layers, n_prompts):
            "--chunk", "16",
            "--output_dir", out_dir]
     print("running:", " ".join(cmd))
-    r = subprocess.run(cmd, cwd='/workspace/code')
+    r = subprocess.run(cmd, cwd=REPO_ROOT)
     assert r.returncode == 0, "jlens_v3 failed"
 
 
diff --git a/src/synthetic_pair.py b/src/synthetic_pair.py
index 7614fbe..6c6c722 100644
--- a/src/synthetic_pair.py
+++ b/src/synthetic_pair.py
@@ -32,6 +32,7 @@ DATA_SRC = 'data/shakespeare_char/input.txt'
 DATA_DIR = 'data/synth_pair'
 OUT_ROOT = 'out-synth-pair'
 JLENS_OUT = 'outputs/synth_pair'
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 T_STRUCT = '@'
 T_NOISE = '#'
 TRIGGER = 'the '
@@ -164,7 +165,7 @@ def jlens(seed, layers, n_prompts):
            "--chunk", "16",
            "--output_dir", out_dir]
     print("running:", " ".join(cmd))
-    r = subprocess.run(cmd, cwd='/workspace/code')
+    r = subprocess.run(cmd, cwd=REPO_ROOT)
     assert r.returncode == 0, "jlens_v3 failed"
 
 
diff --git a/tests/test_jlens_v3.py b/tests/test_jlens_v3.py
new file mode 100644
index 0000000..5ac3f67
--- /dev/null
+++ b/tests/test_jlens_v3.py
@@ -0,0 +1,69 @@
+"""Minimal tests for the faithful J-lens (jlens_v3) correctness.
+
+Run inside the meru container where torch exists:
+    python3 tests/test_jlens_v3.py
+Skips gracefully when torch is unavailable.
+
+Key assertion: at the LAST layer the residual-to-residual Jacobian is the
+identity, so the faithful J-lens vectors must equal the unembedding rows
+(cosine similarity ~1.0). This is the ruler-check that validates the
+W_U-probed VJP machinery.
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+try:
+    import numpy as np
+    import torch
+    import jlens_v3
+except ImportError:
+    print("SKIP: torch not available (run inside the meru container)")
+    sys.exit(0)
+
+from model import GPT, GPTConfig  # noqa: E402
+
+DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def _tiny_model(n_layer=2, d=16, block=32):
+    torch.manual_seed(0)
+    cfg = GPTConfig(n_layer=n_layer, n_head=2, n_embd=d, block_size=block,
+                    bias=False, vocab_size=65, dropout=0.0)
+    return GPT(cfg).eval()
+
+
+def _tiny_data(n=2000, block=32):
+    return np.random.RandomState(0).randint(0, 65, n).astype(np.uint16)
+
+
+def test_shapes_and_finite():
+    model = _tiny_model()
+    data = _tiny_data()
+    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
+    fv = jlens_v3.compute_faithful_jlens(model, 1, batches, DEV, chunk=8)
+    assert fv.shape == (65, 16), f"faithful vecs shape {fv.shape}"
+    assert torch.isfinite(fv).all(), "faithful vecs not finite"
+    pn = jlens_v3.compute_proxy_norms(model, 1, batches, DEV, 65, chunk=16)
+    assert all(np.isfinite(v) for v in pn.values()), "proxy norms not finite"
+
+
+def test_last_layer_identity():
+    """At the final layer J must be identity: faithful vectors == W_U rows."""
+    n_layer, d = 2, 16
+    model = _tiny_model(n_layer=n_layer, d=d)
+    data = _tiny_data()
+    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
+    fv = jlens_v3.compute_faithful_jlens(model, n_layer - 1, batches, DEV,
+                                         chunk=8)
+    wu = model.lm_head.weight.detach().float()
+    sims = torch.nn.functional.cosine_similarity(fv.float(), wu.cpu().float(),
+                                                 dim=1)
+    assert sims.mean().item() > 0.99, f"identity check failed: {sims.mean():.4f}"
+
+
+if __name__ == '__main__':
+    test_shapes_and_finite()
+    test_last_layer_identity()
+    print(f"PASS ({DEV})")