diff options
| -rw-r--r-- | docs/blog-jlens-frequency.md | 290 | ||||
| -rw-r--r-- | docs/willkn-jlens-greaterwrong-analysis.md | 260 | ||||
| -rw-r--r-- | results.md | 24 |
3 files changed, 135 insertions, 439 deletions
diff --git a/docs/blog-jlens-frequency.md b/docs/blog-jlens-frequency.md index f5a459e..0efd9ce 100644 --- a/docs/blog-jlens-frequency.md +++ b/docs/blog-jlens-frequency.md @@ -1,32 +1,30 @@ # What the Jacobian Lens Measures -### A small replication of Anthropic's J-lens, the token-frequency confound we found, and the bug we almost published +### 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. We found something -Anthropic didn't mention in their paper — and then we found that we'd made a +*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.** We reimplemented Anthropic's Jacobian lens faithfully -(verified against their released code) and found that the J-space ranking is -strongly confounded by token frequency: rare tokens score high, common tokens -score low (r ≈ -0.6 to -0.7 at every layer, p ~ 10^-9 or less; against -log-frequency — the natural scale for Zipfian data — the unembedding geometry -alone hits r = -0.69). Anthropic never -controls for frequency — not in the paper, not in the released code. Digging -into *why* gave us the most interesting result: the frequency signal splits -into two separable parts. Half lives in the static geometry of the model's -word-scoring matrix — baked into the lens by definition, so any user inherits -it. A smaller, layer-dependent part lives in what the layers themselves do, and -vanishes at the final layer. The static half survives at real scale: GPT-2's -unembedding rows anti-correlate with token log-frequency too (r ≈ -0.45/-0.49, -V = 50,257, n = 46,887). A frequency-matched synthetic pair shows the lens -also carries genuine structure signal (confirmed by a clean-boundary control at -~1.3x), and a causal test found the demotion effect is small under loss -reweighting. We are **not** claiming the J-space doesn't exist. We're claiming -that any "privileged subspace" interpretation needs a frequency control first. -The full story — numbers, mistakes, and all — is below. +**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. --- @@ -40,7 +38,7 @@ 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 +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. @@ -70,102 +68,100 @@ 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 — +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 we read the paper, something felt off. Here's the thing about token +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. -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. - -(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.) +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. Our first attempt — and the bug three reviewers found +## 4. My first attempt — and the bug three reviewers found -We set out to test this on a small model we could train ourselves: a +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. -Our first implementation looked reasonable. We hooked into each layer, computed +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). We were excited. -We were also wrong. +J-lens norms, common characters had small ones (r ≈ -0.65). I was excited. +I was also wrong. -Before publishing anything, we did something slightly unusual: we asked three -independent AI reviewers to try to tear the work apart. We gave them our code -and our results and asked them to find the flaws. All three, independently, +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: -**Our implementation was not computing Anthropic's Jacobian lens.** +**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. Our code instead differentiated 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. Our beautiful correlation might -have been an artifact of our own measurement. +(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 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. +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 -We rebuilt the lens to match the paper's definition exactly. The faithful -computation is: +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. -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. +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. -(We also confirmed our quantity against Anthropic's released reference +(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 we compute, and our W_U-probed shortcut is +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. We re-ran our analysis +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: we capture the residual stream *before* the model's final +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 we found: frequency is everywhere +## 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: @@ -180,6 +176,9 @@ faithful lens, correlated with token frequency like this: 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 @@ -188,22 +187,21 @@ 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. -### The most interesting thing we found: where the correlation comes from +### Where the correlation comes from -The faithful lens vector for token k is W_U[k]·J_ℓ — the row of the -unembedding matrix times the layer Jacobian. The correlation can come from -either factor, and the two behave very differently. This decomposition is, we -think, the actual novel mechanistic contribution of this project: **the -frequency confound is not one thing.** +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, any user of the lens — including - Anthropic's capacity analysis — inherits this bias automatically. A - frequency control would have to live in the geometry, not in the prompts. + 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 @@ -217,47 +215,26 @@ frequency confound is not one thing.** rare tokens — and it vanishes at the last layer (+0.06). The mechanism of that layer-dependent part is still under investigation. + + + 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 we +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 we go further. We were about to claim "Anthropic does not -control for frequency anywhere," and that is the kind of claim that should be -checked, not asserted. We checked it four ways: our own scan of the paper's -text, two independent adversarial reviewers who read the full paper including -the appendix, and — after a reader pointed us to it — Anthropic's own released -companion code (`github.com/anthropics/jacobian-lens`, Apache-2.0). All agree: -no analysis in the paper controls for token frequency — no frequency matching, -no frequency normalization, no frequency baseline. The released code and -experiment data contain zero frequency handling: a case-insensitive scan of -the entire repo finds no mention of frequency, unigram, or token counts -anywhere. 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, its released code has none -either, 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. - -**Related work.** We are not alone in noticing the raw lens is distorted by -token statistics. An independent research-engineer analysis of the same paper -(willkn, GreaterWrong, 24 Jul 2026), working on GPT-2-medium (355M — thirty -times our model), found that the raw fitted Jacobian "misweights structural -tokens (grammar, punctuation) over semantic content" — dominant spectral -channels carry ~10x the residual pathway's gain, and structural tokens are the -high-frequency tokens. They also found the Jacobian essentially full-rank -(562-858 dimensions for 90% of spectral variance), matching our toy-scale rank -result. Their fix — a single-parameter shrinkage regularizer (J + λI) that -restores next-token faithfulness — is worth testing against our frequency -correlation; Anthropic's released fitting code applies no such regularization. -Whether shrinkage removes the correlation is an experiment we have not run -yet; it is a natural next step. +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 we ran a cleaner test. We made +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%): @@ -271,9 +248,9 @@ 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 + 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.)* @@ -288,7 +265,7 @@ 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 we ran the control that isolates them: '#' inserted at random *word +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. @@ -297,16 +274,19 @@ partly correcting: ``` placement of '#' middle-layer ratio @/# bootstrap 95% CI - random (58% slicing) 1.47 ± 0.09 [1.37, 1.53] - clean boundary (0%) 1.31 ± 0.08 [1.26, 1.40] + random (58% slicing) 1.47 ± 0.09 [1.368, 1.529] + clean boundary (0%) 1.31 ± 0.08 [1.258, 1.402] ``` + + + 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 our first estimate. +survives the control, modestly smaller than my first estimate. ## 8. The causal test: what actually happened @@ -328,53 +308,53 @@ correlational. The faithful lens norm of 'q' (layers 2-4, mean per seed): 2 | 0.0161 0.0158 0.0169 | 1.019 0.952 ``` + + + 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]). -What this shows, honestly: -1. There IS a signal in the expected direction: 2x loss pressure lowers 'q's - faithful norm in 2 of 3 seeds, ~5.6% on average below the plain control. -2. It is small — and the sharper truth is that this design was never powered - to see it. Observed sd 0.07 on the ratio against a mean deficit of 0.056: - a two-sided 80%-power test needs ~13 seeds per arm, and we ran 3. The - experiment could answer "is the effect large?" (no — CI crosses 1.0, one - seed goes the other way, and the effect vanishes against the - random-upweight control) but not "is there any effect at all?". The right - version is ~13 seeds per arm, or a manipulation that moves frequency more - than 2x. -3. The frequency correlation itself is invariant: across all nine trained - models — every mode, every seed — r ≈ -0.63 to -0.69. Training with 'q' - upweighted does not change the correlation structure at all, consistent - with the W_U-decomposition reading that most of the effect is geometric. -4. The contrast with our earlier ablation is instructive: doubling *actual - corpus occurrences* of 'q' dropped its norm by 67%; doubling its *loss - weight* drops it ~6%. The data-frequency lever is a much stronger causal - handle than the gradient lever (AdamW's adaptive per-parameter scaling - absorbs some of the signal — the reviewer who warned about this was right). +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 we could afford to test is weak. This is the honest +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 we are NOT saying +## 9. What I am NOT saying -Let us be very careful here, because it would be easy to overclaim. +Let me 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 +- 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). -- We are **not** saying they have no controls at all. Their occupancy +- 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 - our claim is about. -- We are **not** claiming the lens is useless. The synthetic-pair result shows it + my claim is about. +- I am **not** claiming the lens is useless. The synthetic-pair result shows it carries real structure signal. -- We are **not** claiming that ranking by lens-vector *norm* is the same as - ranking by *lens output on real activations*. Our numbers rank tokens by the +- 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 @@ -383,23 +363,23 @@ Let us be very careful here, because it would be easy to overclaim. 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. -- We are **not** saying "it's just linear algebra." Our toy models don't show +- 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 we **are** saying is narrower and, we think, more durable: on the paper's +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 we can test, and frequency is a variable any J-lens +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 we are taking to bigger models next, and one we +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 -We did try bigger once already, and we owe you the number, because a reader +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. We do not count it as evidence, for three concrete reasons: it sampled +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 @@ -407,7 +387,7 @@ 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. We +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 @@ -416,7 +396,7 @@ 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 we ran it: +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 @@ -435,7 +415,7 @@ lens at scale, and that remains the next post. 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; we +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: @@ -461,6 +441,6 @@ python3 src/loss_reweight.py --step summary --- -*Written in the spirit of the rule we keep trying to follow: the first +*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/docs/willkn-jlens-greaterwrong-analysis.md b/docs/willkn-jlens-greaterwrong-analysis.md deleted file mode 100644 index 7a12cbf..0000000 --- a/docs/willkn-jlens-greaterwrong-analysis.md +++ /dev/null @@ -1,260 +0,0 @@ -# Anthropic's J-Lens: A Research Engineer's Analysis - -**Author:** willkn -**Source:** https://www.greaterwrong.com/posts/vHxGD5HKsFuBStirq/anthropic-s-j-lens-a-research-engineer-s-analysis -**Date:** 24 Jul 2026 -**Notebooks:** https://github.com/willkn/jlens_re -**Note:** Reproduced here (with attribution) as reference material for this -repo's analysis. Model: GPT-2-medium (355M, d=1024, L=24, V=50,257), fp32, HF -transformers, NVIDIA L4. "This is research engineering, not mechanistic -interpretability" — author's own epistemic-status note. - ---- - -Epistemic status: this is research engineering, not mechanistic interpretability. -The compute/cost claims are measured or derived from architecture constants. The -quality claims (faithfulness comparisons, spectral channel interpretations) are -from one small base model (gpt2-medium), one metric, and in places small samples -(n=32); I make no claims about whether the J-space constitutes reasoning or a -workspace and this post is about determining what it costs to run the tool in -production environments, not the tool's outputs. - -**Headline Result: Lens Monitoring is nearly free at decode time with small -dictionary size** - -Anthropic recently released a paper on the transformer circuits platform called -'Verbalizable Representations Form a Global Workspace in Language Models'. This -paper posits that models have an internal workspace where non-verbalised -concepts, perhaps certain reasoning steps or other intermediary computations, -exist. Anthropic call this the J-Space. - -Here is an example from the paper — we see that the model holds certain values -in the J-Space when computing an output. One of the core discussions ongoing -around this paper is whether the values we observe in the J-Space are useful and -if they constitute reasoning. We will not engage with that discussion in this -post. Although, I feel that these values have a lot of potential and imagine -many practitioners will be looking to implement the techniques into their own -experiments or production systems, so this writeup serves as a first venture -into analysis of production-level engineering with the J-Space. - -Throughout the writeup we will focus our analysis on memory, compute and -inference speeds. - -## The J-Space — an overview - -Before we dive into the mathematics, at a high level, how do we access the -J-Space? The residual stream is used to determine an output in the final layer -by providing a set of logits over a vocabulary, such as the English language. -This however is not possible within middle layers as they do not share the same -understanding of the residual stream as the final layer does. The paper -therefore aims to understand these layers as we do the final layer. To do this, -we create an object, the Jacobian, that measures how much the final layer -changes (and in which direction) based on some change at an earlier layer. For -example, we may observe that an earlier layer is making it more likely that the -output will be related to a certain concept. - -The formal mathematical definition of the J-Space gives us instant intuition on -whether the techniques used are feasible for production environments. To access -the J-Space we need to firstly compute a Jacobian, and secondly, apply it at -inference time. - -## Part 1: Computing the Jacobian - -Let h_l[t] be the residual stream vector of dimension d at layer l, position t. -Everything downstream of this vector is a function mapping h_l[t] to the final -layer residuals at h_k[t'], where t' >= t since we are using a causal mask -(earlier positions can't attend to later positions to stop the model 'cheating'). - -We then define our Jacobian. For one prompt and one position pair, we run the -prompt through the transformer blocks (which can be thought of as a non-linear -function f) which produces a final output h_k[t']: - - 1.1 h_L[t'] = f(h_l[t]) - -We can represent this nonlinear function as a Jacobian (which is linear), a -collection of these outputs: - - 1.2 A_l = ∂h_L[t'] / ∂h_l[t] - -Every entry in this Jacobian, for example A[i, j], tells us how much coordinate -i of the final output vector at position t' were to shift if we made a small -change in the residual stream at layer l. Practically, this encodes the -intermediate computations that happen in a forward pass and lets us know what -exactly each token changes about the output. - -It is helpful to think of intermediate layers in terms of the final layer. In -the final layer, we output a word based on the value of the residual vector. In -intermediate layers, we have no such thing, and the value of a residual vector -in layer 5 might not mean the same thing it does in the final layer. The -Jacobian we have here allows us to link non-final layers to the final layers -and extract meaning that otherwise would not be found — analogous to a linear -map (lossy, non-reversible). However, computing the Jacobian with one prompt -gives us a biased view due to taking on individual 'characteristics' of that -prompt, so we must compute with multiple (n=1000 in the Anthropic paper) to get -a better representation: - - 1.3 A_l = (1/N) Σ_i A_i - -This is the general definition of how to obtain a Jacobian for the J-lens. -Let's go into some of the engineering tricks Anthropic used to compute theirs -in the paper. - -### The Engineering of Computing a Jacobian - -We will first look at how Anthropic determine the Jacobian in the paper. - -We call backprop once per output coordinate: inject a one-hot gradient at -coordinate i of the final-layer residual (at every valid target position at -once) and backpropagate to layer l; each backward returns row i. We then stack -these rows to create our final Jacobian. Since backprop passes through every -intermediary layer l' where l' > l, we can pick up those too and build -Jacobians per layer without having to rerun independently. We also need to -remember to take averages when computing to destroy the noise that is -accumulated by individual prompts or tokens. - -In one sentence: 'Run backprop once per output coordinate to harvest the -Jacobian row by row, with positions averaged inside each pass and prompts -averaged across passes — d backwards per prompt, and the average of it all is -the J-lens matrix' - -### Compute required for a Jacobian - -Setup: GPT-2-medium (355M, d=1024, L=24, V=50,257), fp32, HF transformers, -NVIDIA L4. - -Fitting: 128-token WikiText-2 prompts; dim_batch=8 (output dimensions per -backward, prompt replicated along the batch axis). - -(Compute formula, from the author's analysis — see source for full derivation.) - -Let ε be the target relative estimator error at layer l (the "performance p" in -matrix-space terms). c_l is the layer's noise coefficient — the constant in the -measured 1/√n convergence law, where rel_F(J_n) ≈ c_l/√n. Empirically, c_l -grows toward early layers, increasing roughly 4× from layer 20 to layer 4 on -GPT-2-medium. a + b·(L − l) captures the measured per-backward cost, which -scales linearly with depth spanned. On an L4 GPU at dim_batch = 8, we measured -a ≈ 5 ms and b ≈ 9 ms/layer. d/dim_batch gives the number of backward passes -required per prompt, with the forward pass cost t_fwd negligible by comparison. - -Two important notes come with this formula: - -1. **You pay for the earliest layer, the rest are free:** When computing an - earlier layer l, all layers where l' > l are 'free' due to the necessity of - computing them during backpropagation for layer l. -2. **The formula is only valid for ε above the layer's bias floor:** If a layer - cannot represent the non-linearity of future layers above bias floor in a - linear manner, no corpus size will fit a good Jacobian. This is a fundamental - problem with linear functions being unable to approximate non-linear - functions as they grow more complex, not a sample size issue. - -### What Quality looks like — and how to fix it - -On GPT-2-medium, the raw fitted Jacobian underperforms logit lens on -next-token faithfulness. Spectral analysis reveals the cause: the Jacobian's -dominant channels carry ~10× the gain of the residual pathway, misweighting -structural tokens (grammar, punctuation) over semantic content. This is not a -signal problem, but rather a transport weighting problem. - -Crucially, this is fixable. A single-parameter shrinkage regularizer J + λI -monotonically recovers faithfulness across all layers. At layer 12 specifically, -this simple fix doesn't just match logit lens, it exceeds it (0.294 vs. 0.275). -This reveals an important engineering requirement: practitioners using J-lens -monitoring must regularise the raw Jacobian to generalize properly. The -mechanism is straightforward and the improvement is material. - -It should be stressed that the transport weighting issue was only identified on -gpt2-medium. This does not mean that the issue does or does not exist for other -models, we merely provide a fix for gpt2-medium that may work for other models -with the same phenomenon. - -## Part 2: Jacobian at Inference - -Fitting is one off — once we have the Jacobian we do not need to recompute it. -Therefore, it is just a case of applying it at inference time. - -We assume here that we apply monitoring at every token. The lens readout at one -layer and one position is two matrix vector products (transport then decode). - - 3.1.1 lens(h) = softmax(W_U · ln_f( J_l · h )) - cost per layer per position = (2d² (transport J_l·h) + 2dV (unembed)) - d = model dimensionality, V = vocab size, h = residual stream - -Against the model's own per-token forward cost of 24·L·d² monitoring K layers -with a full-vocabulary readout at every generated token adds: - - 3.1.2 overhead = K·(d + V) / (24·L·d) - -Practical use note: Anthropic use a method where only certain words are -considered in monitoring for the J-Lens. Softmax poses a problem here: if we -monitor for violent content and the user asks about the weather, softmax gives -a distribution of all the monitored words rather than their magnitude — no -signal, false positive. The vocabulary term dominates: full vocab readout of -one layer = 17% of a forward pass; 5 layers = 90%; all 24 layers of GPT-2 = 4x -a forward pass. Markedly lower in a frontier-shaped model (D=8192, L=80, -V=128k) due to V/d shrinking — ~45% of a forward pass for 24 layers, still too -expensive for an always-on monitor. - -The deployment method Anthropic cover in the paper — a fixed length dictionary -to reduce vocabulary size. Everything before the softmax is linear (through the -Jacobian approximation), so this reduces the calculation into dot products of -the residual against precomputed J-Lens vectors, one per watched token: - - 3.2 score(c) = (J_lᵀ u_c) · h - cost = 2·C·d per layer per token - -For C=1000 concepts over five layers on gpt2-medium this is **less than 2% of -a forward pass**. - -As expected, wall clock overhead and FLOP share grows (nearly) linearly as -layer size increases. Kernel launch floor ≈ 1.2% (c=100 gives FLOP share -0.17%). Percentages relative to a baseline serving speed of 18.3ms/tok. - -### Can the linear transport be compressed? - -A natural idea is to try and reduce the rank of the Jacobian to reduce the -computation required. However, on this model this does not seem to be a -particularly viable. To capture 90% of spectral variance requires 562-858 -dimensions; the remaining 10% spreads across the last 166-462 dimensions -dependent on layer. **The Jacobian is essentially full-rank, leaving little -room for compression.** I posit that similar size models likely have the same -problem due to being overcomplete with features, but it is yet to be seen if -there is an opportunity to reduce the rank with larger models. - -### Putting it all together - -Anthropic claim in the paper that middle layers are where the J-Lens works -best. We make no claim about the quality of the concepts verbalised in these -middle layers but we do claim that middle layers are optimal for engineering -and computational efficiency with the J-Lens — a synergistic result. Earlier -layers are more expensive to compute (backprop through every layer that follows) -and less useful ('setting the stage'; linear approximation struggles with many -nonlinear layers). Even where logit-lens is a better tool in later layers, the -difference in compute for a single layer, single token between the two methods -is small: logit lens = 2dV; J-lens = 2d² + 2dV; difference = 2d² (~2% more -FLOPs on gpt2-medium given vocab dominance). - -### So What? - -Anthropic's J-Lens monitoring is deployment viable on small models with a -dictionary of around 1000 concepts at 2% compute overhead per token. No claims -on extrapolation to larger models, but the mathematics suggest likely more -efficient on larger models (V/d shrinks). Key engineering takeaways for smaller -models: (1) Regularise the raw Jacobian; (2) Target middle layers for best -utility/cost tradeoff; (3) Compute Jacobians for each layer simultaneously. - -### Conclusion - -The J-Lens is a technique that is certainly feasible at runtime with reasonably -sized dictionaries. It has the potential to transform how models are monitored, -understood and finetuned (Goodfire, 2026). All with the added benefit of taking -relatively few FLOPs as a percentage of total FLOPs, and being inexpensive to -train, even at frontier model levels. However, this must be read with the caveat -of results on smaller models, where certain directions dominate without -suppression of the dominant directions. At least in this implementation, this -was a real emergent phenomenon that needed addressing to get the J-Lens to work -in smaller models. It is implied that this was not a problem in Claude's 4.5/4.6 -generation of models through Anthropic's work, but it is yet to be seen if this -happens with other models. - -References: Gurnee, W., et al, 2026. Verbalizable Representations Form a Global -Workspace in Language Models. 24 Jul 2026, GreaterWrong. @@ -89,30 +89,6 @@ experiment + evaluation data) finds ZERO occurrences of frequency/unigram/ token-count terms. Their README confirms the fit corpus is "a generic web-text corpus" — no frequency balancing. -## 1c. Independent corroboration (willkn, GreaterWrong, GPT-2-medium) - -An independent engineering analysis of the paper (willkn, GreaterWrong, 24 Jul -2026; notebooks: github.com/willkn/jlens_re; full text: docs/willkn-jlens- -greaterwrong-analysis.md), on GPT-2-medium (355M, d=1024, L=24, V=50,257), -found: -- the raw fitted Jacobian "misweights structural tokens (grammar, punctuation) - over semantic content" — dominant spectral channels carry ~10x the gain of - the residual pathway. Structural tokens ARE the high-frequency tokens; this - is the same distortion we measure as frequency anti-correlation, at 30x our - model scale and via a different metric (next-token faithfulness of the - readout). -- the Jacobian is essentially full-rank: 562-858 dimensions capture 90% of - spectral variance, "leaving little room for compression" — corroborates our - dimensional-starvation / full-rank result at scale. -- a single-parameter shrinkage regularizer (J + lambda*I) monotonically - recovers next-token faithfulness (beats logit lens at layer 12: 0.294 vs - 0.275). Anthropic's released fitting code applies no such regularization. -- measured noise law rel_F(J_n) ~ c_l/sqrt(n), c_l growing ~4x from layer 20 - to layer 4; per-backward cost ~5ms + 9ms/layer on L4 at dim_batch=8. - -Untested (natural next experiment): does shrinkage also remove our frequency -correlation? - ## 2. Frequency-matched synthetic pair (`src/synthetic_pair.py`) Two new characters at identical 0.1% unigram frequency in Shakespeare: |