diff options
| -rw-r--r-- | blog/assets/fig1_layer_correlation.png | bin | 0 -> 145506 bytes | |||
| -rw-r--r-- | blog/assets/fig2_wu_decomposition.png | bin | 0 -> 277177 bytes | |||
| -rw-r--r-- | blog/assets/fig3_synthetic_pair.png | bin | 0 -> 176898 bytes | |||
| -rw-r--r-- | blog/assets/fig4_loss_reweighting.png | bin | 0 -> 155678 bytes | |||
| -rw-r--r-- | blog/index.qmd | 2 | ||||
| -rw-r--r-- | blog/jacobian-lens-frequency.qmd | 466 | ||||
| -rw-r--r-- | terminal.css | 22 |
7 files changed, 490 insertions, 0 deletions
diff --git a/blog/assets/fig1_layer_correlation.png b/blog/assets/fig1_layer_correlation.png Binary files differnew file mode 100644 index 0000000..ae81814 --- /dev/null +++ b/blog/assets/fig1_layer_correlation.png diff --git a/blog/assets/fig2_wu_decomposition.png b/blog/assets/fig2_wu_decomposition.png Binary files differnew file mode 100644 index 0000000..99eb496 --- /dev/null +++ b/blog/assets/fig2_wu_decomposition.png diff --git a/blog/assets/fig3_synthetic_pair.png b/blog/assets/fig3_synthetic_pair.png Binary files differnew file mode 100644 index 0000000..ff7ed00 --- /dev/null +++ b/blog/assets/fig3_synthetic_pair.png diff --git a/blog/assets/fig4_loss_reweighting.png b/blog/assets/fig4_loss_reweighting.png Binary files differnew file mode 100644 index 0000000..524fc8f --- /dev/null +++ b/blog/assets/fig4_loss_reweighting.png diff --git a/blog/index.qmd b/blog/index.qmd index be1fbd5..69b65b5 100644 --- a/blog/index.qmd +++ b/blog/index.qmd @@ -7,6 +7,8 @@ RANDOM THOUGHTS, EXPERIMENTS, AND SYSTEMS NOTES - How this site works — 2026-08-05 curl [jayrup.me/blog/how-this-site-works](https://jayrup.me/blog/how-this-site-works) +- What the Jacobian Lens Measures — 2026-08-07 + curl [jayrup.me/blog/jacobian-lens-frequency](https://jayrup.me/blog/jacobian-lens-frequency) -------------------------------------------------- diff --git a/blog/jacobian-lens-frequency.qmd b/blog/jacobian-lens-frequency.qmd new file mode 100644 index 0000000..6a81809 --- /dev/null +++ b/blog/jacobian-lens-frequency.qmd @@ -0,0 +1,466 @@ +--- +title: "What the Jacobian Lens Measures" +date: "2026-08-07" +--- + +*A small replication of Anthropic's J-lens, the token-frequency confound I found, and the bug I almost published.* + +*This is a story about trying to look inside a language model. I found something +Anthropic didn't mention in their paper — and then I found that I'd made a +mistake, fixed it, and the thing was still there. That second part is the +stronger result.* + +--- + +**The short version.** I rebuilt Anthropic's Jacobian lens and checked it +against their released code. In my model, its token ranking has a large and +simple bias: rare tokens get big scores; common tokens get small ones +(r ≈ -0.6 to -0.7 at every layer). Anthropic's paper and released code do not +control for token frequency. + +That is not the whole story. I can split the effect in two. Most of it is in +the model's built-in word-scoring table: training gives rare tokens bigger +rows there, and the lens necessarily reads through those rows. A smaller part +comes from the layers themselves. The first effect also appears in GPT-2 at a +50,257-token vocabulary (r ≈ -0.45/-0.49 against log-frequency). + +Nor is the lens *only* measuring frequency. When I gave two invented tokens +exactly the same frequency, the one the model could predict in context still +scored about 1.3x higher. So this is not a refutation of J-space. It is a +more modest claim: before treating a J-lens ranking as evidence for a +privileged concept workspace, control for frequency first. + +--- + +## 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. I 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 my 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 I 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. + +My first thought was about gradients. A gradient through the model's final +probability calculation has a built-in quirk: the less likely a word is, the +larger one of its raw terms can be. That would make rare words look important +before the model had said anything interesting about them. + +That intuition applies directly to the simpler measurement I tried first, +which differentiated through the final softmax. It does *not* by itself +explain Anthropic's faithful lens. As I later found, the faithful version has +a different source of bias: part of it is sitting in the geometry of the +word-scoring matrix. But it gave me the itch worth checking. + +In other words: **a "privileged workspace" might just be a frequency effect +wearing a fancy hat.** + +## 4. My first attempt — and the bug three reviewers found + +I set out to test this on a small model I could train myself: 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. + +My first implementation looked reasonable. I 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). I was excited. +I was also wrong. + +Before publishing anything, I did something slightly unusual: I asked three +independent AI reviewers to try to tear the work apart. I gave them my code +and results and asked them to find the flaws. All three, independently, +found the same one: + +**My 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. My code instead differentiated through the +softmax directly. That folds a frequency-dependent calibration factor — the +(1 - p) term — into the thing being averaged. My beautiful correlation might +have been an artifact of my own measurement. + +This is the part of the story I like best, because it's the part that's easy +to skip: I had built a measurement that *looked* like the paper's and wasn't. +The reviewers caught it, I fixed it, and the honest result got stronger. + +## 5. The right way + +I rebuilt the lens to match the paper's definition exactly. In plain English, +I ask: if I nudge this layer a little, what average change reaches the final +residual stream? Only after averaging those changes do I use the model's own +word-scoring table to turn them into token directions. Formally: + +> 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. + +I verified it the way you verify a ruler. At the last layer, the map from the +layer to itself must do nothing at all. The faithful J-lens vectors must +therefore be exactly the model's own word-scoring rows. My check returned +cosine similarity 1.0000. The ruler is correct. + +(I also confirmed my 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 I compute, and my 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. I re-ran my analysis +with their exact estimator choices: the frequency correlation is essentially +identical at every layer (max delta 0.008, see results.md section 1b), so the +result is robust to those choices.) + +One technical note before moving on, because it matters for the +interpretation: I 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 +normalization applied only when *reading* the lens, i.e. +softmax(W_U·norm(J_ℓ·h_ℓ)). Under that definition the last-layer identity +check is exact by construction. + +## 6. What I 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 +``` + +Figure 1 plots the same numbers: per layer, the faithful lens (blue), the +buggy proxy (purple), and the Spearman rank correlation (cyan). The Spearman +bars are the most negative — the effect is not carried by a few extreme +tokens. + + + + +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. The +statistics are not subtle: Pearson r ≈ -0.61 to -0.69 (p ~ 10^-9 to 10^-13), +and Spearman rank correlation is even stronger (-0.69 to -0.85), so the +result is not an artifact of a few extreme common tokens. + +### Where the correlation comes from + +The faithful lens vector for token k is W_U[k]·J_ℓ: one row of the model's +word-scoring table, passed through the layer map. Think of W_U as the ruler I +use to read the model. A token-frequency effect could be in the ruler, in the +layers, or in both. It is in both. + +1. **Static geometry — the ruler.** The unembedding row norms ||W_U[k]|| + themselves anti-correlate with frequency — most strongly against + log-frequency, the natural scale for Zipfian data: r(||W_U[k]||, log10 f) = + -0.69 (raw frequency -0.61, Spearman -0.81). + Rare tokens get bigger rows in the word-scoring matrix. Since the lens + reads through W_U by definition, the raw norm ranking inherits that bias + automatically. Looking at different prompts cannot remove a bias already + built into the ruler. + + Where does that geometry come from? Not from initialization: fresh models + show no frequency correlation in their row norms (r ≈ +0.00 to -0.19 + across three seeds). After training it is -0.61. The model *learns* to + push rare tokens' rows outward as it learns the corpus, and nothing in the + objective or the readout ever corrects it. The lens inherits learned + geometry, not a fixed one. +2. **Layer dynamics — the layers.** Regress out the W_U component and an + anti-correlation still survives in layers 0-4 (partial r ≈ -0.24 to + -0.36) — something about what the layers themselves do keeps boosting + rare tokens — and it vanishes at the last layer (+0.06). The mechanism + of that layer-dependent part is still under investigation. + +The left panel of Figure 2 plots that residual directly; the right panel +shows the ruler's bias appearing during training. + + + + +So the lens carries a frequency signal from both the ruler it reads with and +from what the layers do — and the two are separable. That is the finding I +would most want someone to test at scale. + +A fact-check before I go further. "Anthropic does not control for frequency" +is a claim worth checking, not simply making. I searched the paper and its +appendix, asked two independent reviewers to do the same, and checked +Anthropic's released companion code (`github.com/anthropics/jacobian-lens`, +Apache-2.0). I found no frequency matching, normalization, or baseline in +the J-lens analyses, and no frequency, unigram, or token-count handling in the +released code or experiment data. Their appendix does mention filtering +high-frequency noise tokens for a *different* method, the template lens, and +calls that move unprincipled. It is not a control for the main J-lens. + +## 7. But not *only* frequency + +Now the twist. Correlation is not causation, so I ran a cleaner test. I 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 +``` +*(Norm ranges are across layers 0-5 within each seed; ratios are per-layer +@/#. The middle-layer summary follows.)* + +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: at equal frequency, the two +tokens differ in norm. + +One caveat, found by a reviewer: the noise token '#' was inserted at random +character positions, which slices through the middle of a word 58% of the time +(th#e, ki#ng — letter on both sides), while '@' always sits at a clean word +boundary after "the ". That confounds predictability with n-gram corruption — +so I ran the control that isolates them: '#' inserted at random *word +boundaries* (0% word-slicing, still unpredictable), same 0.0998% frequency, +three fresh seeds. + +The control is done, and it is the honest kind of result — partly confirming, +partly correcting: + +``` + placement of '#' middle-layer ratio @/# bootstrap 95% CI + random (58% slicing) 1.47 ± 0.09 [1.368, 1.529] + clean boundary (0%) 1.31 ± 0.08 [1.258, 1.402] +``` + +Figure 3 shows both conditions against the ratio=1 baseline (dashed red line, +the no-signal case): the clean-boundary bar still clears it, CI included. + + + + +The corruption confound was real: it inflated the estimate by about 12%. But +it was not the whole story. At identical frequency, with clean boundaries and +nothing sliced, the predictable token still scores ~1.3x higher than the +unpredictable one, and the CI stays entirely above 1 in every seed. The +conditional-predictability signal — the thing "verbalizable" should mean — +survives the control, modestly smaller than my first estimate. + +## 8. The causal test: what actually happened + +The last experiment was the one designed to make the frequency story causal. +Train three models per seed from the *identical* starting weights and the +*identical* minibatch order — the only difference is the loss: one model gives +the letter 'q' twice the learning pressure (2x CE weight on 'q' targets, which +raises its effective frequency without corrupting the text), one is a plain +control, and one upweights the same number of random *other* letters (to check +that "any reweighting" isn't the thing doing the work). Three seeds, three +models each. If doubling 'q's effective frequency causally shrinks its J-lens +norm below both controls, the frequency story is causal, not just +correlational. The faithful lens norm of 'q' (layers 2-4, mean per seed): + +``` + seed | q(2x) control ctrl_random | q/control q/ctrl_random + 0 | 0.0163 0.0174 0.0152 | 0.934 1.069 + 1 | 0.0150 0.0171 0.0161 | 0.881 0.932 + 2 | 0.0161 0.0158 0.0169 | 1.019 0.952 +``` + +Figure 4 shows the three arms per seed — q upweighted (red), plain control +(blue), random-upweight (green) — plus the cross-seed mean with the annotation +q/control = 0.944, CI [0.881, 1.019]. The CI straddling 1 is the picture of +"too small to settle it." + + + + +Cross-seed: q/control mean = 0.944 (bootstrap 95% CI [0.881, 1.019]), +q/ctrl_random mean = 0.985 (CI [0.932, 1.069]). + +The answer is: perhaps, but this experiment is too small to settle it. In two +of three runs, giving `q` twice the loss weight lowered its norm; on average it +was 5.6% below the ordinary control. But the confidence interval crosses 1, +one run went the other way, and the result disappears against the +random-upweight control. With the variation I saw, a properly powered version +needs about 13 seeds per arm, not three. + +There is a useful lesson in the weak result. Doubling actual occurrences of +`q` in the training text had previously dropped its norm by 67%. Doubling the +loss weight moved it only about 6%. Changing what the model sees is a much +stronger lever than changing the size of its gradient after the fact; AdamW +appears to absorb part of the latter change. + +Meanwhile the overall frequency correlation barely moved: it stayed around +r ≈ -0.63 to -0.69 across all nine models. That is what I would expect if +most of the pattern is in the learned word-scoring geometry, rather than a +fragile effect of one token's loss weight. + +Net: the frequency confound is strongly correlational and geometrically +stable; the causal lever I could afford to test is weak. This is the honest +state of the causal evidence. (Absolute 'q' norms differ across experiments — +base model 0.011 vs these 0.015-0.018 — so only within-experiment +comparisons are meaningful.) + +## 9. What I am NOT saying + +Let me be very careful here, because it would be easy to overclaim. + +- I am **not** saying the J-space doesn't exist. I 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). +- I am **not** saying they have no controls at all. Their occupancy + analysis compares against random-direction baselines, and their probes + subtract mean concept directions. Those are real experimental controls — + but none of them is a token-frequency control, which is the specific thing + my claim is about. +- I am **not** claiming the lens is useless. The synthetic-pair result shows it + carries real structure signal. +- I am **not** claiming that ranking by lens-vector *norm* is the same as + ranking by *lens output on real activations*. My numbers rank tokens by the + norm of their faithful J-lens vector — a summary of the readout geometry — + not by how strongly, or how often, those directions actually fire in running + text. Anthropic's capacity claim is about the latter (occupancy). The norm + still matters: it is the size of a token's lever on the readout, and the W_U + geometry that drives it enters any readout-based analysis by definition. But + the gap between "geometry is frequency-confounded" and "the capacity claim + is frequency-confounded" is real, and it is the specific gap an at-scale + occupancy test has to close. +- I am **not** saying "it's just linear algebra." My 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 I **am** saying is narrower and, I think, more durable: on the paper's +own measurement, J-lens *norm-rankings* are strongly confounded by token +frequency at every scale I can test, and frequency is a variable any J-lens +analysis should control for. Whether the confound survives at Anthropic's scale +is an empirical question — one I am taking to bigger models next, and one I +already have a first, partial answer to for the geometric half (Section 10). + +## 10. What's next + +I did try bigger once already, and I owe you the number, because a reader +who opens the repo will find it either way: an early probe on GPT-2 small +(`src/gpt2_jlens.py`) returned an average correlation of only r ≈ -0.18 across +layers. I do not count it as evidence, for three concrete reasons: it sampled +96 token positions out of a 50,257-token vocabulary; it averaged over only 100 +sampled tokens per batch; and it measured a subtly different quantity +(norm-per-batch rather than norm-of-the-mean). It was a directional probe, and +it pointed weak. It is logged in `results.md`, flagged do-not-cite — but a +post that promises "bigger models next" should not pretend the attempt never +happened. + +Toy scale answers the methodological question. Scale answers the real one. I +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. + +The learned-geometry finding makes one piece of that cheaper than the probe +was: if the W_U row-norm anti-correlation is a general property of +softmax-output models trained on Zipfian data, it should appear in GPT-2's +unembedding matrix directly — no Jacobian computation at all. So I ran it: +GPT-2's unembedding row norms correlate with token log-frequency at V = 50,257 +(r ≈ -0.45 on gpt2-small, -0.49 on gpt2-medium, n = 46,887 tokens seen in +wikitext-103; Spearman -0.46 to -0.50). The decile picture is monotone in both +models — the rarest tenth of tokens have the biggest rows — and the top-norm +tokens are rare ones (SPONSORED, soDeliveryDate, Reviewer). The geometric half +of the confound is not a 65-character vocabulary artifact: it is present in +two real, independently trained language models at Anthropic's vocabulary +scale. One honest caveat: it is weaker than at toy scale (-0.45/-0.49 vs +-0.69), so the char-scale result was not a size artifact, but neither is the +at-scale effect as strong as the toy suggested. The layer-dynamics half — +what the layers themselves add beyond the ruler — still needs the faithful +lens at scale, and that remains the next post. + +## 11. How to reproduce everything + +All code, data-prep scripts, experiment scripts, tests, and this analysis live +in the repository: <https://git.jayrup.me/c/jspace-nanogpt.git/>. 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; I +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 I 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/terminal.css b/terminal.css index 22d4295..70da81e 100644 --- a/terminal.css +++ b/terminal.css @@ -52,6 +52,28 @@ nav a { box-shadow: 0 10px 30px rgba(0,0,0,0.5); } +/* Figures: constrain to the content column, keep aspect ratio, center */ +/* NOTE: with format minimal:true quarto emits .quarto-figure divs directly in + body — there is NO #quarto-content wrapper in the rendered HTML. Target the + figure classes + generic imgs, not #quarto-content. */ +.quarto-figure { + margin: 1.5rem 0; +} +.quarto-figure img, +img { + display: block; + max-width: 100%; + height: auto; + margin: 0 auto; +} +.quarto-figure figcaption, +.quarto-figure > p { + text-align: center; + font-size: 0.85rem; + color: #7f849c; + margin-top: 0.5rem; +} + /* Subtle scanline effect */ body::before { content: " "; |
