A research engineering analysis of the Jacobian lens ("J-lens") from Verbalizable Representations Form a Global Workspace in Language Models (Gurnee et al., Anthropic, 2026), reproduced on GPT-2-medium. The question: what does it actually cost to fit the lens, and what does it cost to run it as a monitor at decode time?
Companion code for the writeup. Reference implementation of the method itself: anthropics/jacobian-lens.
Fitting is cheap and the cost model is simple. Fitting cost follows
C ≈ 2·N·T·d_model·n_prompts (validated against a profiler to ~5%). The
backward:forward FLOP ratio is ~0.95 — input-gradients only, not training's 2x —
though wall-clock backward runs at 1.3x a forward. On an L4, GPT-2-medium fits at
22.4 s/prompt (12% MFU); the paper's 1000-prompt recipe costs ~6.7 h, but lens
quality saturates around 100 prompts (~40 min), confirming the reference
implementation's claim.
Pricing is set by the earliest layer. One backward pass serves every fitted
layer, so cost scales with depth spanned, not layer count. Convergence follows a
1/sqrt(n) law and is depth-ordered: layer 4 carries ~4x the relative error of
layer 20 at every corpus size, so early layers need up to ~16x the prompts for
matched error. The pricing model: C(l, e) = (c_l/e)^2 * t_prompt(l).
The raw lens underperforms the logit lens on GPT-2-medium. On next-token
faithfulness (top-5 overlap with the model's final logits), the converged J-lens
loses to the logit lens at every depth. Spectral analysis shows why: the averaged
Jacobian's dominant channels carry ~10x the gain of the residual pathway and
transport attention-sink structure and document register rather than token
content. A one-parameter shrinkage J + lambda*I monotonically recovers
faithfulness and exceeds the logit lens at layer 12 (0.294 vs 0.275, n=32).
The transport cannot be low-rank compressed. The fitted Jacobians need 562-858 of 1024 singular directions (by layer) to capture 90% of spectral mass.
Monitoring is nearly free with concept dictionaries; full-vocabulary readout
is not. Analytically, a full-vocab readout costs ~17% of a forward pass per
layer (4x a forward for all layers); a precomputed concept dictionary
(W_U @ J_l rows) costs under 2% for 1000 concepts over five layers. Measured
during KV-cached batch-1 decoding on an L4: full-vocab at five layers adds +5.3%
wall-clock (its FLOPs hide in single-stream idle slack; at serving batch sizes
costs converge to the analytic ratios), dictionary mode adds +1.2% (a fixed
kernel-launch floor that batching amortizes away).
Open a notebook in Google Colab (or VS Code's Colab extension), select a GPU
runtime, and Run All. Each notebook is self-contained: it installs missing
packages, writes jlens_mini.py into the runtime, and saves/loads all results
through Google Drive (MyDrive/jlens/). Run locally instead and results go to
./jlens_out/. Fitted Jacobians are computed once (notebook 02) and reused by
everything else.
Run in order:
| # | notebook | measures | runtime |
|---|---|---|---|
| 01 | 01_fitting_cost.ipynb |
FLOPs and wall-clock of fitting: cost-model validation, backward/forward ratio, dim_batch sweep, MFU, scaling probes, extrapolation to larger models | minutes |
| 02 | 02_convergence_frontier.ipynb |
The real fit, checkpointed: 1/sqrt(n) convergence and noise floor, functional knee in GPU-minutes, J-lens vs logit-lens faithfulness, variance across disjoint thirds | ~1 h (T4) |
| 03 | 03_decode_overhead.ipynb |
Measured per-token monitoring cost during KV-cached decoding: full-vocab (K=1/3/5 layers) and concept-dictionary (C=100/1000) vs analytic FLOP ratios | ~10 min |
Files produced in Drive: fit_timing.json (01), jlens_checkpoints.pt and
frontier_metrics.json and figures (02), decode_overhead.json (03).
- Model: GPT-2-medium (355M parameters, d=1024, 24 layers, vocab 50,257), fp32
- Corpus: WikiText-2 (
Salesforce/wikitext), 128-token prompts - Hardware: NVIDIA L4 (Colab); notebooks also run on MPS/CPU with reduced defaults
- Stack: PyTorch, HF transformers, HF datasets, matplotlib
- Estimator:
jlens_mini.py, a port of the reference implementation's reduction (one forward ondim_batchprompt replicas, one-hot cotangents at every valid target position per backward, mean over source positions, fp32 accumulation). Verified against brute-forcetorch.autograd.functional.jacobianto <1e-6 through an identical forward path. - Timing: CUDA-synchronized, warm-up passes discarded, best-of-3 for decode
- transformers >=5:
GPT2Blockreturns a bare tensor rather than a tuple; the activation hook handles both. - Target positions: cotangents are injected only at valid positions (final position excluded as a target, first 16 positions excluded as sources, matching the reference implementation). A naive sum over all t' >= t differs by ~1% and will not match the reference estimator.
- The lens matrices are fitted per layer as full d x d matrices; the deployed
dictionary form folds
W_U @ J_loffline so monitoring is dot products only.
All results are GPT-2-medium: a small base model with learned absolute position embeddings, outside the model family the paper studied (Claude 4.5/4.6 generation). The faithfulness metric is next-token prediction, which is not the paper's concept-readout use case. Measured decode overheads are batch-1 numbers; serving batch sizes converge toward the analytic FLOP ratios.