diff options
| -rw-r--r-- | README.md | 21 | ||||
| -rw-r--r-- | docs/blog-jlens-frequency.md | 34 | ||||
| -rwxr-xr-x | scripts/test.sh | 3 | ||||
| -rw-r--r-- | src/loss_reweight.py | 3 | ||||
| -rw-r--r-- | src/synthetic_pair.py | 3 | ||||
| -rw-r--r-- | tests/test_jlens_v3.py | 69 |
6 files changed, 114 insertions, 19 deletions
@@ -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})") |