# J-space on nanoGPT

A small-scale replication and critique of Anthropic's **Jacobian Lens**
("J-lens") technique from *"Verbalizable Representations Form a Global
Workspace in Language Models"* (2026,
[transformer-circuits.pub/2026/workspace](https://transformer-circuits.pub/2026/workspace/index.html)).

**Short version of the findings:**

1. The faithful J-lens (rows of `W_U * J_l`, exactly per the paper's Methods)
   has J-lens norms that are strongly anti-correlated with token frequency at
   every layer of a 10.65M-parameter character-level transformer
   (r ≈ -0.61 to -0.69). Anthropic does not control for frequency anywhere in
   their analysis.
2. The correlation is NOT a measurement artifact: it survives the paper's exact
   quantity, verified by a last-layer identity check (cosine similarity 1.0000).
3. The lens is not *only* a frequency meter: two synthetic tokens at identical
   unigram frequency get different faithful J-lens norms when one is predictable
   in context (~1.4-1.5x higher for the structured token, 3 seeds).
4. A causal loss-reweighting test (2x loss weight on 'q' targets vs two
   controls) tests whether effective frequency causally demotes a token's
   J-lens norm. See `results.md` for the latest numbers.

See `docs/blog-jlens-frequency.md` for the write-up and `results.md` for the
numbers. The three independent adversarial reviews that shaped the project
(and caught a real bug in the first implementation) are in `docs/reviews/`.

## What we are NOT claiming

- That Anthropic's J-space doesn't exist. Their headline capacity claim is
  about activation *occupancy* per position, which this repo does not test.
- That the J-lens is useless — it carries genuine conditional-predictability
  signal.
- That toy-model results refute large-model findings. The claim is narrower:
  J-lens *rankings* are frequency-confounded, so a frequency control is
  required before any "privileged subspace" interpretation.

## Repository layout

```
model.py, train.py        nanoGPT (Karpathy) with Maxwell-GPU fixes
src/jlens.py, jlens_v2.py first (buggy) J-lens implementations — superseded
src/jlens_v3.py           FAITHFUL J-lens: rows of W_U * J_l (canonical)
src/synthetic_pair.py     frequency-matched synthetic pair experiment
src/loss_reweight.py      causal loss-reweighting experiment
src/gpt2_jlens.py         GPT-2 scale test (under-powered; see results.md)
tests/                    unit tests (see scripts/test.sh)
scripts/test.sh           canonical test command
docs/blog-jlens-frequency.md  write-up (Feynman-style)
docs/reviews/             three adversarial model reviews
results.md                committed experiment numbers
```

## Reproducing everything

### Environment

- Any Linux box with Docker and an NVIDIA GPU (we used a 4GB Quadro K2200).
- Container image: `pytorch/pytorch:2.4.1-cuda11.8` (last CUDA for Maxwell).
- Sync this repo into the container, e.g. `/workspace/code`.

The K2200 has 4GB VRAM, so batched VJP probes must be chunked (`--chunk 16`);
on a modern GPU you can raise it. Everything runs fp32 (no bf16 on Maxwell).

### 1. Tests

```sh
sh scripts/test.sh
```

Skips gracefully where torch is unavailable.

### 2. Train the base model

Train nanoGPT on `data/shakespeare_char` (10.65M params, 6 layers, d=384,
block 128) and keep the checkpoint at `out-shakespeare-char/ckpt.pt`:

```sh
python3 train.py config/train_shakespeare_char.py
```

### 3. Faithful J-lens, both-ways comparison

```sh
python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \
    --data_dir data/shakespeare_char --layers 0,1,2,3,4,5
```

Prints per-layer frequency correlations for the faithful lens and the old
proxy, plus the last-layer identity validation. Artifacts land in
`outputs/jlens_v3/`.

### 4. Synthetic frequency-matched pair

```sh
python3 src/synthetic_pair.py --step prep      # builds data/synth_pair
python3 src/synthetic_pair.py --step train --seed 0
python3 src/synthetic_pair.py --step jlens --seed 0
python3 src/synthetic_pair.py --step summary
```

### 5. Loss-reweighting causal test

```sh
python3 src/loss_reweight.py --step train --mode q --seed 0
python3 src/loss_reweight.py --step train --mode control --seed 0
python3 src/loss_reweight.py --step train --mode ctrl_random --seed 0
python3 src/loss_reweight.py --step summary --layers 2,3,4
```

## References

- [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) — Anthropic, 2026
- [nanoGPT](https://github.com/karpathy/nanoGPT) — Andrej Karpathy