summaryrefslogtreecommitdiff
path: root/blog/jacobian-lens-frequency.qmd
blob: 6a818097b7ccf3b9844c2d4dbc430c830305d86c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
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.

![Figure 1: Token Frequency Anti-Correlation Across Layers](assets/fig1_layer_correlation.png)


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.

![Figure 2: The frequency confound splits into two separable parts. Left: what remains after the word-scoring-table effect is regressed out — the layer dynamics, strongest in layers 0-4 and gone at the last layer. Right: the table's frequency bias is learned during training — absent at initialization (r ~ 0), strong after (r ≈ -0.61).](assets/fig2_wu_decomposition.png)


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.

![Figure 3: Synthetic Frequency-Matched Pair Control (@ vs # at 0.1% Frequency). Error bars are bootstrap 95% CIs; the dashed line marks ratio 1 (no predictability signal).](assets/fig3_synthetic_pair.png)


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."

![Figure 4: Loss-Reweighting Causal Test Across 3 Seeds. Per-seed bars: q(2x), control, random-upweight; annotation = cross-seed ratio + 95% CI.](assets/fig4_loss_reweighting.png)


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.*