J-Lens: A Failed Replication on GPT-2

TLDR: We tested Anthropic's new(ish) J-Lens against the classic Logit Lens on GPT-2 small and medium. J-Lens loses at every layer (0/11 on small, 1/23 on medium). We ran 5 stress tests (more data, frequency checks, sparsity, tuning, and scaling) and it still loses every time. We draw two sharp methodological lessons for those doing probe-based interpretability.

Repo: https://github.com/nelithb/mech-interp-lab/

1: Result

J-Lens does not work on GPT-2 small or medium.

Across 11 layers on GPT-2 small, J-Lens gets beaten by Logit Lens on every single one. On GPT-2 medium, it wins 1 of 23 layers, and even this ‘win’ is a rounding error (rank 221 vs 220).

This matters because the J-Lens paper makes a specific claim: that Logit Lens fails because early layer representations aren’t geometrically aligned with the final layer’s coordinate space, and that J-Lens corrects this. Our experiment directly tests this claim.

We ran five tests to make sure this wasn’t anomalous: refitting on more data; checking against confounding token frequencies; sparsity thresholding; tuning ‘skip_first’ parameter and scaling up from 124M to 255M parameters. J-Lens lost each time. The last scaling result is interesting in that the gap widened on the future token prediction metric, which is the opposite of J-Lens’s purpose.

The real value in this post is two methodological lessons that we extracted. If you are doing probe-based interpretability work, you can internalise these:

1. Rank metrics are gameable by token frequency. A lens can actually score well by just predicting common tokens (e.g. commas and spaces). This was confirmed with 0.656 correlation between the shared bias term and token log frequency.
2. Magnitude thresholding surface outlier dimensions, but not structure. When we thresholded J-Lens to sparsity, it ‘won’ at 7/11 layers. But upon inspection, there were only 9 unique top-1 tokens across all prompts. All of them were generic filler. These tracked to GPT-2’s known ‘massive activation’ outlier dimensions.

We made these mistakes so that you don’t have to.

2: Technical background

Logit Lens

Transformers work by passing a single vector – the residual stream – through the network. Each layer adds to this vector without replacing it.. So, an earlier layer's residual stream sits in the exact same dimensional space as the final one. Logit Lens (nostalgebraist, 2020) exploits this directly. It takes the model's final LayerNorm and unembedding matrix (the same weights normally reserved for the last layer) and applies them to an earlier layer's residual stream, to see what token the model would guess if it stopped there.

J-Lens

But Logit Lens rests on a faulty assumption: that the early layer's representation lives within the geometric space that the final layer expects. J-Lens, in theory, corrects for this. As the paper puts it: 'While the logit lens assumes that representations use the same coordinates in all layers, the Jacobian lens corrects for representational changes that take place across layers.' If this claim is true, J-Lens becomes a cheap lens into the model at every layer. This experiment is a direct test of this.

3: Setup

For this project I only had access to an Intel Mac. PyTorch dropped x86 MacOS builds. Rather than moving everything to Colab, I ported the core into my existing CPU vent, using the official repo purely as a reference to check against.

To verify correct implementation, I conducted two checks:

1. The paper states that setting J to the identity matrix should make J-Lens collapse to exactly Logit Lens. This was tested, returning a max difference of 0.00e+00, i.e. byte-identical.
2. I ran a finite difference gradcheck against the actual perturbed network, at float64 (GPT-2's large late-layer residual norms lose gradient signal to floating-point cancellation at float32). This returned a max relative error of 8.4e-07: essentially machine precision.

The ported CPU implementation used for these experiments is available at https://github.com/nelithb/mech-interp-lab/. This is a clean re-implementation. The official Anthropic repo was used only as a reference for verification checks.

4: Headline result

Firstly, we ask: 'does the model's own top-1 next token get recovered, at 60 held-out prompts, per layer' (Metric A)? J-Lens loses at all 11 layers. J-Lens's median rank for the correct answer is in the thousands at early layers, reaching ~5-38 in the last two layers. Conversely, Logit Lens stays in the 1-43 range throughout.

Secondly, we look at future tokens at position p+1 to p+5 (Metric B). We measure at layer 9. We'd expect J-Lens to perform better here, given its whole purpose is to capture what the model is about to say, and not just what it's currently saying. Again, J-Lens is worse at every offset, with this gap persisting as we look further into the future (flat around -8 to -9 nats). J-Lens loses even on the task its own theory says it should win.

In truth, we use only 30 sequences (vs the paper's ~1000 sequences). So next, we test whether this undermines our result.

5: Stress testing

We firstly refit on 150 sequences, but the results remained essentially unchanged. J Lens still loses in all 11 layers on Metric A. J-Lens's rank for the correct answer is 11184, 12717, 9493, 3596, 4351, 1756, 2035, 457, 405... at successive layers, only reaching single digits at the very last one. Logit Lens at this point is near-perfect (rank 0). On Metric B, we run flat at -8.6 to -9.1 nats – practically matching the original pattern.

I ran into this comment from LessWrong user 'phoenix' which claimed that J-Lens's output is partially explained by raw token frequency, i.e. not genuine prediction. To test this, we measured the correlation between the shared final bias term (added identically in both lenses as they share the same final unembed step) and token log frequency across a 300k-token sample. We got r=0.656 — pretty close to phoenix's estimate (r=0.67). Although subtracting this bias out severely alters the absolute rank of both lenses (at layer 0: Logit Lens 43 -> 910; J-Lens 2759 -> 40992), the relative comparison remains unchanged. J-Lens still scores worse at all 11 layers.

Inspired by sparsity-based approaches, we tested whether a hard threshold towards sparsity on the fitted J would reveal a cleaner structure.. The first pass looked like a win: at 1% keep (the most aggressive threshold), J Lens won in 7/11 layers (!!), up from 0/11. But three signs pointed to this 'win' being the metric rewarding a content-agnostic filler predictor: only generic filler (e.g. commas or spaces) appeared in the top 5; only nine unique top-1 tokens ever appeared across the 60 held-out prompts; and the surviving weights traced to two of GPT-2's three known 'massive activation' outlier dimensions, independently documented by Timkey and van Schijndel (2021). So, common tokens score well on Metric A, regardless of whether anything real is being extracted.

Our fourth swing was to refit J at skip_first values of 0, 16 and 32. This resulted in up to a 20x improvement: mean J Lens median rank at early layers went from 8320 to 1937 to 408 as skip_first increases. Despite this, J Lens still goes 0/11 at every value tested. Even the best tuned configuration (~400 median rank) is roughly 100x worse than Logit Lens's single-to-low double digit range. Although this was methodologically useful, it just didn't save the result. One more roll of the dice...

6: Scale sweep

Counterintuitively, when we increase scale from gpt2-small (124M) to gpt2-medium (355M), the gap widens. On Metric A, J-Lens wins in 1/23 layers, with this victory at layer 0. And yet, even this win is rounding-error level: J-Lens ranks 221; Logit Lens at 220.

Metric B's (future token prediction) result is striking not because the gap got bigger, but because its shape changed. On gpt2-small it was flat across all five offsets. On gpt2-medium it's no longer flat, it widens the further out you look. This is the wrong direction if J-Lens's future-token advantage was supposed to grow with scale. Recall that this is the more important metric for J-Len’s central claim. On gpt2-small, the logit-minus-J gap was flat across all five future offsets, roughly -8 to -9 nats, regardless of how far ahead you looked. On gpt2-medium, it's not flat, but monotonically widens: -4.970, -6.753, -7.239, -7.179, -7.277 at offsets p+1 through p+5. If J-Lens’s advantages were to emerge with scale, we’d expect this gap to shrink.

This result must be caveated that this is one data point, bridging 124M to 355M.

7: Open questions

Our result holds within the scope we tested. Here are three open opportunities to expand:

(1) We max out at 355M parameters. gpt2-large (774M) and gpt2-xl (1.5B) are natural next steps. The trend so far (albeit between two data points), i.e. the gap widening, argues against expecting a reversal.

(2) Expand beyond GPT-2. Belrose et al.'s Tuned Lens paper (arXiv:2303.08112) states that 'this simple form of the Logit Lens works reasonably well for GPT-2', with their own appendix testing, across four ~125M models, an improved lens against Logit Lens. They found that the improved lens clearly won only on OPT-125m, not GPT-2 medium, stating that 'these results did not generalize to other models tested.' On GPT-2, Logit Lens is unusually strong.

(3) This project actually started as a test of deception under instruction, and we didn't test that. The natural next step would be to use an instruction tuned model (e.g. Qwen 0.5-1.5B which is CPU feasible) and a localise-then-patch design instead of a probe-only readout.

8: Takeaways

As a beginner, I wanted to include this section, so that if you’re doing lens or probe style interpretability work you can just internalise these lessons. I made the mistakes so that you don’t have to!

Firstly, rank metrics are gameable by frequency. We saw in reference to phoenix’s comments that a rank based metric can be moved a lot by oddities that are unrelated to genuine prediction.

Secondly, the magnitude thresholding test that resulted in the 7/11 J-Lens win surfaced the weights that were numerically the largest, which is not the same as the most meaningful. Always check the contents of the returned tokens to verify if this is the case.

I wouldn’t have found these obstacles had I not got stuck in.

Note on AI use

I used AI to help write all the code in this project. It was also used to provide explanations for the results and theoretical concepts, and find and summarise key existing research. The visualiser was also made with AI. I'm not an AI researcher, or even a developer, but I can't wait to continue digging deeper with AI's help.

Sources

Gurnee, W., Sofroniew, N., Pearce, A., Piotrowski, M., Kauvar, I., et al. (2026). Verbalizable Representations Form a Global Workspace in Language Models. Transformer Circuits Thread, Anthropic.

• Available at: https://transformer-circuits.pub/2026/workspace

Anthropic. (2026). jacobian-lens [Companion code repository]. GitHub.

• Available at: https://github.com/anthropics/jacobian-lens

nostalgebraist. (2020). Interpreting GPT: The Logit Lens. LessWrong.

• Available at: https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens

Belrose, N., Furman, Z., Smith, L., Halawi, D., Ostrovsky, I., McKinney, L., Biderman, S., & Steinhardt, J. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens. arXiv:2303.08112.

• Available at: https://arxiv.org/abs/2303.08112

Timkey, W., & van Schijndel, M. (2021). All Bark and No Bite: Rogue Dimensions in Transformer Language Models Obscure Representational Quality. In Proceedings of EMNLP 2021, pp. 4527–4546.

• Available at: https://aclanthology.org/2021.emnlp-main.370/

Sun, M., Chen, X., Kolter, J. Z., & Liu, Z. (2024). Massive Activations in Large Language Models. arXiv:2402.17762.

• Available at: https://arxiv.org/abs/2402.17762

Ran-Milo, Y. (2026). Attention Sinks Are Provably Necessary in Softmax Transformers: Evidence from Trigger-Conditional Tasks. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 2).

• Available at: https://aclanthology.org/2026.acl-short.8/

Nanda, N., & Bloom, J. (2022). TransformerLens [Software library]. GitHub.

• Available at: https://github.com/TransformerLensOrg/TransformerLens

phoenix. (2026). Comment on "A global workspace in language models". LessWrong.

• Available at: https://www.lesswrong.com/posts/3PaLrzxagpbnNtPLT?commentId=c4dNnEwARCxLBm9YG [Referenced in comments section]