diff options
| -rw-r--r-- | LICENSE | 21 | ||||
| -rw-r--r-- | README.md | 111 | ||||
| -rw-r--r-- | docs/blog-jlens-frequency.md | 242 | ||||
| -rw-r--r-- | docs/reviews/2026-07-31-claude-opus-4.6.md (renamed from reviews/2026-07-31-claude-opus-4.6.md) | 0 | ||||
| -rw-r--r-- | docs/reviews/2026-07-31-gemini-3.1-pro.md (renamed from reviews/2026-07-31-gemini-3.1-pro.md) | 0 | ||||
| -rw-r--r-- | docs/reviews/2026-07-31-gpt-5.6-terra.md (renamed from reviews/2026-07-31-gpt-5.6-terra.md) | 0 | ||||
| -rw-r--r-- | requirements.txt | 4 | ||||
| -rw-r--r-- | results.md | 62 |
8 files changed, 427 insertions, 13 deletions
@@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Jayrup Nakawala + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. @@ -1,27 +1,112 @@ # J-space on nanoGPT -Replicating Anthropic's Jacobian Lens ("J-space") technique on a small character-level transformer (nanoGPT). +A small-scale replication and critique of Anthropic's **Jacobian Lens** +("J-lens") technique from *"Verbalizable Representations Form a Global +Workspace in Language Models"* (2026, +[transformer-circuits.pub/2026/workspace](https://transformer-circuits.pub/2026/workspace/index.html)). -## Hypothesis +**Short version of the findings:** -If J-space (the subspace of representations readily available for verbal report) is an **architectural/structural property of transformers** rather than an emergent feature of advanced models, it should appear at all scales — including 10M-parameter char-level models. +1. The faithful J-lens (rows of `W_U * J_l`, exactly per the paper's Methods) + has J-lens norms that are strongly anti-correlated with token frequency at + every layer of a 10.65M-parameter character-level transformer + (r ≈ -0.61 to -0.69). Anthropic does not control for frequency anywhere in + their analysis. +2. The correlation is NOT a measurement artifact: it survives the paper's exact + quantity, verified by a last-layer identity check (cosine similarity 1.0000). +3. The lens is not *only* a frequency meter: two synthetic tokens at identical + unigram frequency get different faithful J-lens norms when one is predictable + in context (~1.4-1.5x higher for the structured token, 3 seeds). +4. A causal loss-reweighting test (2x loss weight on 'q' targets vs two + controls) tests whether effective frequency causally demotes a token's + J-lens norm. See `results.md` for the latest numbers. -## Background +See `docs/blog-jlens-frequency.md` for the write-up and `results.md` for the +numbers. The three independent adversarial reviews that shaped the project +(and caught a real bug in the first implementation) are in `docs/reviews/`. -Anthropic's 2026 paper "Verbalizable Representations Form a Global Workspace in Language Models" introduces the Jacobian Lens (J-lens), which computes the average linearized effect of activations on future token probabilities, averaged over many contexts. This reveals a privileged "J-space" of representations that the model can report on, modulate, and use for reasoning. +## What we are NOT claiming -Full paper: https://transformer-circuits.pub/2026/workspace/index.html +- That Anthropic's J-space doesn't exist. Their headline capacity claim is + about activation *occupancy* per position, which this repo does not test. +- That the J-lens is useless — it carries genuine conditional-predictability + signal. +- That toy-model results refute large-model findings. The claim is narrower: + J-lens *rankings* are frequency-confounded, so a frequency control is + required before any "privileged subspace" interpretation. -## Experiments +## Repository layout -1. **J-space visualization** — Compute J-lens vectors for all vocabulary tokens at each layer. Visualize which characters/concepts enter "verbalizable space" and when. -2. **Ablation test** — Remove J-space components vs random directions vs full activations. Measure prediction quality impact. -3. **Training dynamics** — Save checkpoints during training, compute J-space at each, track when it crystallizes. -4. **Capacity measurement** — Count active J-lens tokens per position. +``` +model.py, train.py nanoGPT (Karpathy) with Maxwell-GPU fixes +src/jlens.py, jlens_v2.py first (buggy) J-lens implementations — superseded +src/jlens_v3.py FAITHFUL J-lens: rows of W_U * J_l (canonical) +src/synthetic_pair.py frequency-matched synthetic pair experiment +src/loss_reweight.py causal loss-reweighting experiment +src/gpt2_jlens.py GPT-2 scale test (under-powered; see results.md) +tests/ unit tests (see scripts/test.sh) +scripts/test.sh canonical test command +docs/blog-jlens-frequency.md write-up (Feynman-style) +docs/reviews/ three adversarial model reviews +results.md committed experiment numbers +``` -## Setup +## Reproducing everything -Runs on meru's Quadro K2200 (4GB VRAM) via Docker with GPU passthrough. +### Environment + +- Any Linux box with Docker and an NVIDIA GPU (we used a 4GB Quadro K2200). +- Container image: `pytorch/pytorch:2.4.1-cuda11.8` (last CUDA for Maxwell). +- Sync this repo into the container, e.g. `/workspace/code`. + +The K2200 has 4GB VRAM, so batched VJP probes must be chunked (`--chunk 16`); +on a modern GPU you can raise it. Everything runs fp32 (no bf16 on Maxwell). + +### 1. Tests + +```sh +sh scripts/test.sh +``` + +Skips gracefully where torch is unavailable. + +### 2. 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`: + +```sh +python3 train.py config/train_shakespeare_char.py +``` + +### 3. Faithful J-lens, both-ways comparison + +```sh +python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \ + --data_dir data/shakespeare_char --layers 0,1,2,3,4,5 +``` + +Prints per-layer frequency correlations for the faithful lens and the old +proxy, plus the last-layer identity validation. Artifacts land in +`outputs/jlens_v3/`. + +### 4. Synthetic frequency-matched pair + +```sh +python3 src/synthetic_pair.py --step prep # builds data/synth_pair +python3 src/synthetic_pair.py --step train --seed 0 +python3 src/synthetic_pair.py --step jlens --seed 0 +python3 src/synthetic_pair.py --step summary +``` + +### 5. Loss-reweighting causal test + +```sh +python3 src/loss_reweight.py --step train --mode q --seed 0 +python3 src/loss_reweight.py --step train --mode control --seed 0 +python3 src/loss_reweight.py --step train --mode ctrl_random --seed 0 +python3 src/loss_reweight.py --step summary --layers 2,3,4 +``` ## References diff --git a/docs/blog-jlens-frequency.md b/docs/blog-jlens-frequency.md new file mode 100644 index 0000000..d0f7ffa --- /dev/null +++ b/docs/blog-jlens-frequency.md @@ -0,0 +1,242 @@ +# What the Jacobian Lens Actually Measures +### A small replication of Anthropic's J-lens, the token-frequency confound we found, and the bug we almost published + +*This is a story about trying to look inside a language model. We found something +Anthropic didn't mention in their paper — and then we found that we'd made a +mistake, fixed it, and the thing was still there. That second part is the +stronger result.* + +--- + +## 1. The machine that guesses words + +A language model is, at its heart, a machine that guesses the next word. Show it +"the cat sat on the" and it produces a list of probabilities for what comes next: +"mat" high, "chair" high, "banana" low. Everything it "knows" is wrapped up in +that guessing. + +The interesting question is: *where* does the guessing happen? A modern model +has dozens of layers, each transforming the sentence a little. Somewhere in +those layers, the model is deciding that "cat" is an animal, that "sat" is past +tense, that a location is coming. We would like to watch that happen. The +problem is that the inside of a transformer is a soup of high-dimensional +vectors, and no one has a map. + +For a long time, people used the "logit lens": at each layer, take the +representation, and ask "if the model had to guess *right now*, what would it +guess?" The trouble is that representations change coordinate systems as they +travel through the layers, so early layers give you nonsense. It's like trying +to read a letter that's been translated into a language you don't know — at the +start of the chain, the translation is too rough. + +## 2. Anthropic's idea: the Jacobian lens + +In 2026, Anthropic published a paper — "Verbalizable Representations Form a +Global Workspace in Language Models" — introducing a smarter version: the +*Jacobian lens*. Instead of asking "what would the model guess right now?", it +asks a sharper question: *"if I nudge this representation a tiny bit, how much +does the final guess move?"* + +That's what a Jacobian is: a table of "how much does each output move when each +input moves." The lens computes, for every layer, the average nudge-effect of +that layer's representation on every word in the vocabulary, averaged over a +thousand different contexts. Words whose representations are strongly "poised" +to be spoken — ready to be said, should the occasion arise — get big numbers. +Anthropic calls this collection of word-vectors the **J-space**, and they claim +it's a kind of "global workspace": a small, privileged subset of the model's +internal state that can be reported on, modulated, and used for reasoning. They +even note the resemblance to theories of consciousness, carefully, the way you +would mention a bear while making clear you are not feeding it. + +The headline claim that caught our eye: **the J-space has limited capacity — +only 10 to 50 concepts are "active" at once.** A tiny privileged workspace +inside a big model. That's a strong claim. Strong claims deserve strong tests. + +## 3. The itch + +The moment we read the paper, something felt off. Here's the thing about token +frequencies: in any language, a handful of words ("the", "of", "and") appear +all the time, and thousands of words appear almost never. In the model's +vocabulary of 50,257 tokens, the rarest are nearly invisible. + +Now, the J-lens vector for a word is a gradient — it measures how much the +model's computation tunes toward that word. And there's a mechanical quirk of +gradients through softmax: the *less* likely a word is, the *larger* the raw +gradient term can be. A gradient of log-probability contains a term that looks +like (1 - p), where p is the word's probability. Rare words have small p, so +(1 - p) is close to 1. Common words have large p, so (1 - p) is small. If the +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. + +In other words: **a "privileged workspace" might just be a frequency effect +wearing a fancy hat.** + +## 4. Our first attempt — and the bug three reviewers found + +We set out to test this on a small model we could train ourselves: a +10.65-million-parameter character-level transformer (Karpathy's nanoGPT), +trained on Shakespeare. Small enough to run on a 4GB GPU in a few hours. Big +enough to have real layers. + +Our first implementation looked reasonable. We hooked into each layer, computed +the gradient of log-probability for every character, averaged over contexts, +and — sure enough — found a strong correlation: rare characters had big +J-lens norms, common characters had small ones (r ≈ -0.65). We were excited. +We were also wrong. + +Before publishing anything, we did something slightly unusual: we asked three +large independent AI models to try to tear the work apart — Gemini 3.1 Pro, +Claude Opus 4.6, and GPT-5.6. We gave them our code and our results and asked +them to find the flaws. All three, independently, found the same one: + +**Our implementation was not computing Anthropic's Jacobian lens.** + +Anthropic's lens computes the average Jacobian from a layer to the *final +representation* — the residual stream — and *then* reads it out through the +model's word-scoring matrix. Our code instead differentiated through the +softmax directly. That folds a frequency-dependent calibration factor — the +(1 - p) term — into the thing being averaged. Our beautiful correlation might +have been an artifact of our own measurement. + +This is the part of the story we like best, because it's the part that's easy +to skip: we had built a measurement that *looked* like the paper's and wasn't. +The reviewers caught it, we fixed it, and the honest result got stronger. + +## 5. The right way + +We rebuilt the lens to match the paper's definition exactly. The faithful +computation is: + +> For each layer ℓ, compute the average Jacobian from that layer to the final +> residual stream, over all source positions, all future positions, and many +> prompts. The J-lens vector for a word is that matrix read through the +> model's own unembedding rows. + +We verified our implementation the way you verify a ruler: at the last layer, +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. + +## 6. What we found: frequency is everywhere + +On the real trained model, all six layers, both the old (buggy) proxy and the +faithful lens, correlated with token frequency like this: + +``` + Layer proxy r faithful r + L0 -0.661 -0.643 + L1 -0.673 -0.668 + L2 -0.653 -0.672 + L3 -0.648 -0.685 + L4 -0.562 -0.637 + L5 -0.665 -0.606 +``` + +The correlation survived the faithful implementation — slightly *stronger*, if +anything. The rare characters ('?', 'z', 'q', '$') sit at the top of the +J-space ranking; the common ones (space, 'e', 't', 'i') sit at the bottom. On +the paper's own quantity, the J-lens ranking is frequency-confounded. Anthropic +does not control for this anywhere in their analysis. + +## 7. But not *only* frequency + +Now the twist. Correlation is not causation, so we ran a cleaner test. We made +a new corpus with two brand-new characters, both at *exactly* the same +frequency (0.1%): + +- `@` — appears only after the trigger "the ". The model can predict it in + context. It is *poised to be said*. +- `#` — appears at random positions. Nothing predicts it. + +Same frequency. Different structure. If the J-lens were purely a frequency +meter, the two tokens would get identical norms. Here is what three separate +training runs showed: + +``` + seed @ norm (predictable) # norm (noise) ratio + 0 0.0232 - 0.0246 0.0152 - 0.0154 1.51 - 1.60 + 1 0.0224 - 0.0237 0.0148 - 0.0151 1.50 - 1.60 + 2 0.0215 - 0.0233 0.0154 - 0.0163 1.35 - 1.51 +``` + +The predictable token scores **~1.4-1.5x higher** than the noise token at +identical frequency, in every layer of every seed. 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. + +## 8. The causal test (in progress) + +We are currently running the last experiment: train three models per seed, +identical in every way, except one model gives the letter 'q' twice the +learning pressure (2x loss weight on 'q' targets — increasing its effective +frequency without corrupting the text), a control model with normal loss, and a +second control that upweights the same number of random *other* letters. If +doubling 'q's effective frequency causally shrinks its J-lens norm below both +controls, the frequency story is causal, not just correlational. Results land +within hours; this post will be updated. + +## 9. What we are NOT saying + +Let us be very careful here, because it would be easy to overclaim. + +- We are **not** saying the J-space doesn't exist. We haven't tested + Anthropic's actual capacity claim (which is about *occupancy* — how often + J-lens directions are used per position — not about the rank of the word + vectors). +- We are **not** saying the lens is useless. The synthetic-pair result shows it + carries real structure signal. +- We are **not** saying "it's just linear algebra." Our toy models don't show + the compression Anthropic sees in large models; that's a limitation of toy + models, not evidence against large ones. + +What we **are** saying is narrower and, we think, more durable: on the paper's +own measurement, J-lens *rankings* are strongly confounded by token frequency, +and any claim about a privileged subspace must control for frequency first. +Anthropic's paper does not. The burden of proof is on them — and it's a fair +one. + +## 10. What's next + +Toy scale answers the methodological question. Scale answers the real one. We +want to run the faithful lens on a real language model (V = 50K, d = 768 — the +regime where Anthropic's claims live) with proper statistical power, and to run +the occupancy test their capacity claim is actually about. That's the next +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. + +Run the test suite: +``` +sh scripts/test.sh +``` + +Rebuild the main experiment from scratch: +``` +# 1. train the character-level model on Shakespeare (10.65M params) +# 2. compute the faithful J-lens + old proxy, all layers: +python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \ + --data_dir data/shakespeare_char --layers 0,1,2,3,4,5 +# 3. synthetic frequency-matched pair: +python3 src/synthetic_pair.py --step prep +python3 src/synthetic_pair.py --step train --seed 0 +python3 src/synthetic_pair.py --step jlens --seed 0 +python3 src/synthetic_pair.py --step summary +# 4. loss-reweighting causal test: +python3 src/loss_reweight.py --step train --mode q --seed 0 +python3 src/loss_reweight.py --step summary +``` + +--- + +*Written in the spirit of the rule we keep trying to follow: the first +principle is that you must not fool yourself — and you are the easiest person +to fool.* diff --git a/reviews/2026-07-31-claude-opus-4.6.md b/docs/reviews/2026-07-31-claude-opus-4.6.md index 2ff8227..2ff8227 100644 --- a/reviews/2026-07-31-claude-opus-4.6.md +++ b/docs/reviews/2026-07-31-claude-opus-4.6.md diff --git a/reviews/2026-07-31-gemini-3.1-pro.md b/docs/reviews/2026-07-31-gemini-3.1-pro.md index 20ad69d..20ad69d 100644 --- a/reviews/2026-07-31-gemini-3.1-pro.md +++ b/docs/reviews/2026-07-31-gemini-3.1-pro.md diff --git a/reviews/2026-07-31-gpt-5.6-terra.md b/docs/reviews/2026-07-31-gpt-5.6-terra.md index 46b5e39..46b5e39 100644 --- a/reviews/2026-07-31-gpt-5.6-terra.md +++ b/docs/reviews/2026-07-31-gpt-5.6-terra.md diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d7bc501 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +# Runtime tested in the pytorch/pytorch:2.4.1-cuda11.8 image (see README). +torch==2.4.1 +numpy +transformers==4.44.0 # only needed for the GPT-2-scale experiments (gpt2_jlens.py) diff --git a/results.md b/results.md new file mode 100644 index 0000000..27b4005 --- /dev/null +++ b/results.md @@ -0,0 +1,62 @@ +# Results + +All numbers below are from the faithful J-lens (`src/jlens_v3.py`), which +computes exactly the paper's quantity: rows of `W_U * J_l` where +`J_l = E[ d h_final / d h_l ]` (average residual-to-residual Jacobian, read out +through the unembedding). Verification: at the last layer, J must be the +identity, and the check returns cosine similarity **1.0000**. + +## 1. Both-ways comparison: old proxy vs faithful lens (trained 10.65M char model) + +Pearson r between token frequency and J-lens norm, per layer. n = 65 tokens +(full char vocab). The correlation survives the faithful implementation at +every layer. + +``` + Layer proxy r faithful r + L0 -0.661 -0.643 + L1 -0.673 -0.668 + L2 -0.653 -0.672 + L3 -0.648 -0.685 + L4 -0.562 -0.637 + L5 -0.665 -0.606 +``` + +Top tokens by faithful norm are consistently rare characters (`?`, `z`, `q`, +`$`); bottom are common ones (space, `e`, `t`, `i`). + +## 2. Frequency-matched synthetic pair (`src/synthetic_pair.py`) + +Two new characters at identical 0.1% unigram frequency in Shakespeare: +`@` appears only after the trigger "the " (predictable in context); +`#` appears at uniform random positions. Faithful J-lens norm per seed, range +over layers 0-5: + +``` + seed @ norm (predictable) # norm (noise) ratio freq corr r + 0 0.0232 - 0.0246 0.0152 - 0.0154 1.51-1.60 -0.59..-0.68 + 1 0.0224 - 0.0237 0.0148 - 0.0151 1.50-1.60 -0.59..-0.66 + 2 0.0215 - 0.0233 0.0154 - 0.0163 1.35-1.51 -0.62..-0.66 +``` + +Reading: at equal frequency, the structured token scores ~1.4-1.5x higher. +The frequency anti-correlation holds, but the lens also carries genuine +conditional-predictability signal. + +## 3. Loss-reweighting causal test (`src/loss_reweight.py`) + +Three models per seed, identical init + minibatch order: 'q' targets weighted +x2 in the loss, plain control, and a same-total-loss control upweighting random +non-'q' targets. Question: does raising effective frequency causally reduce +'q's faithful J-lens norm? PENDING — run completes within hours of this file +being written; the summary table is printed by +`python3 src/loss_reweight.py --step summary --layers 2,3,4`. + +## 4. Historical / do-not-copy + +- Original proxy finding (r = -0.65, `jlens_v2`): superseded by the faithful + implementation; kept only for the both-ways comparison. +- GPT-2 correlation (r = -0.18, `gpt2_jlens.py`): UNDER-POWERED (96 token + positions, n=100 sampled tokens) and computed a different quantity + (norm-per-batch vs norm-of-mean). Directionally consistent but not + publishable evidence on its own. |