Red-teaming the Jacobian Lens

This was originally my write-up for Neel Nanda's MATS stream application.

Executive Summary

Problem: red-teaming the Jacobian Lens (J-Lens) and the "Global Workspace" theory [Gurnee et al.]. Specifically, I conducted a comparison between J-Lens and logit lens "recall" and "specificity".

• Recall: ability to identify correct intermediate thoughts
• Specificity: ability to dim signal on irrelevant / potentially confusing concepts

The original paper focused on the former — comparing how well different lenses surface the true intermediate thought (via the pass@k metric), deeming J-Lens as a more refined "recall" instrument over logit lens.

In addition to recreating the "recall" advantage of J-Lens over the logit lens, I investigated its specificity as well, answering the central question:

"On top of saying 'yes' to concepts that are present, does the J-Lens do a better job at saying 'no' to irrelevant, absent concepts?"

Takeaway: J-Lens's recall advantage, validated on pass@k, does not carry over to its specificity at N=41.

For each prompt and the associated true intermediate, I constructed 3 "distractor" intermediates, sorted into two broad categories: 2 adjacent distractors and 1 near-miss distractor.

• Adjacent: same-category words, tokens carrying similar information
• Near-miss: tokens that could be caused by model confusion, given the context of the prompt (see example in Randomly Selected Samples)

Experiments below are conducted on Qwen 3.5 9B, with N=41 multihop prompts (limitation: too few prompts). I am additionally using Camila and Agam's open source J-Lens for the model.

Key experiment 1: pass@k score on true and distractor intermediates. Scored in the exact same way as the original paper, comparing J-Lens and logit lens.

The takeaway: J-lens lifts distractor pass@k AUC (near-miss: 0.27 vs 0.22, adjacent: 0.26 vs 0.19) by about the same margin as it lifts true intermediates (0.69 vs 0.65), though that difference is marginal given N=41. This means J-lens raises the whole semantic neighborhood together.

Key experiment 2: ROC-AUC curve on true positive rate (TPR) versus false positive rate (FPR) of correct versus distractor intermediates, scored on negative recovery rank. Two takeaways:

• For the top 1 token surfaced, J-Lens catches ~37% of true intermediates while flagging ~0% of distractors. This shows the high specificity of J-Lens itself.
• Logit lens scores very similarly across the board, achieving an AUC of 0.83 vs 0.84. This shows J-Lens does not achieve a specificity advantage over logit lens.

Key experiment 3: Pass@k and ROC-AUC curve on random baseline vectors, both isotropic and covariance-matched "pseudo-token" vectors. Takeaway: Both lenses beat isotropic and covariance-matched random vectors decisively within the workspace layers L18-30 (ROC-AUC ~0.99), meaning its parity with logit lens is not an artifact of both being noise.

Randomly Selected Samples

Below are some randomly selected samples of prompts and true intermediates (taken directly from the original Global Workspace paper), as well as experimental "distractor" tokens. Adjacent distractors were handpicked from a list provided by my agent. Near-miss distractors were generated by a separate reasoning model (Claude Fable 5) given only the prompt and the true intermediate, without access to the adjacent set.

• Prompt: "Fact: The ocean on the coast of the country where Carnival is most famously celebrated is the "
  • Target: Atlantic
  • True intermediate: Brazil
  • Adjacent distractors: Argentina, Peru
    • Semantically similar to Brazil; South America countries with high populations.
  • Near-miss distractor: Italy
    • Carnivals are also popular in Venice, Italy. Potential point of confusion for model.
• Prompt: "Fact: The atomic number of the gas that makes up about 78 percent of the atmosphere is "
  • Target: 7
  • True intermediate: Nitrogen
  • Adjacent distractors: Oxygen, Hydrogen
    • Semantically similar to Nitrogen; Common gases.
  • Near-miss distractor: Oxygen
    • Second most abundant gas in the atmosphere, one that life needs.
    • Notice there are sometimes overlaps between near-miss and adjacent distractors! See What I Verified on how that's handled.
• Prompt: "Fact: The first letter of the month containing Halloween is "
  • Target: O
  • True intermediate: October
  • Adjacent distractors: September, November
    • Adjacent months
  • Near-miss distractor: November
    • Day-of-the-dead, celebrated by certain cultures, is similar to Halloween but celebrated during November.
• Prompt: "Fact: The ocean east of the country that built the Great Wall is the "
  • Target: Pacific
  • True intermediate: China
  • Adjacent distractors: Korea, Japan
    • Other East Asian countries
  • Near-miss distractor: India
    • India has the 'Great Wall of India' (Kumbhalgarh).

Setup

• Model: Qwen3.5-9B. 32 layers, hidden_dim=4096.
  • Began experiments on Qwen3.5-4B, but model was getting far too many answers incorrect, so I made the switch over.
• J-Lens: Official release of Camila & Agam's open source J-Lens library as found on HuggingFace (@qwen3.5-9b/j-lens/lens.pt).
• Hardware: Remote RunPod GPU (A40 with 48gb VRAM)
• Prompts: A subset of the original paper's released multihop eval set (41/93).
  • Qwen3.5-9B was not able to answer all prompts with the correct target token for a variety of reasons (formatting, verbosity, incorrect factual recall, etc.). Only 48/93 were answered correctly.
  • Of the 48, only 41 were comparable with single-token intermediates and distractors. This resulted in a limited sample set N=41 (see Limitations below).

Main Findings

Each claim will be backed with experimental finding(s), documented in the following format:

• Question + what potential outcomes look like
• Methods
• Results
• Interpretation / implication of outcomes

Claim 1: J-Lens provides no specificity advantage over logit lens.

Experiment 1: Pass@k recall, substituting true intermediate with distractors.

Question: Beyond surfacing the true intermediate, do J-Lens surface semantically similar concepts that could lead to an incorrect factual recall? What about as compared to logit lens?

• Outcomes should determine the recall ability of J-Lens. Specifically, we should expect J-Lens to better surface true intermediates as compared to logit lens, demonstrating a more refined recall mechanism.
• If J-Lens have good specificity, we should see conceptually similar but irrelevant tokens (hereon referred to as "distractor" tokens) scoring low on pass@k as well.
• Though this is not a direct comparison, we could empirically observe how well the different lenses score distractor token as a preliminary gauge of specificity.

Methods

• Similar to the original paper, I plot pass@k (the fraction of prompts where the intermediate is recovered at the k-th position of J-Lens output) over a range of values k, and take the area under the resulting curve, plotted against log k and normalized. I refer to this plot as "pass@k AUC".
• Distractor token generation: For each {prompt, true_intermediate} pair, generate 2 additional "adjacent distractors" and 1 additional "near-miss distractor".
  • Adjacent distractors should be semantically similar to the true intermediate.
  • Near-miss distractor should reflect potential model confusion in the context of the specific factual recall.
• I plot the pass@k AUC for the true intermediate as well as the two types of distractors.

Results

J-Lens standalone pass@k AUC for the three different types of intermediates. We could see that true intermediates boast an impressive AUC of 0.69, with both types of distractors trailing behind at 0.27 and 0.26.

Pass@k AUC with logit lens additionally overlayed. The logit lens surfaces true intermediates marginally worse than J-Lens (0.65 vs 0.69), consistent with the results found in the original paper.

Similar to true intermediates, the logit lens also surfaces the two types of distractors marginally worse than J-Lens does (near-miss 0.22 vs 0.27, adjacent 0.19 vs 0.26).

Interpretations

J-Lens has good recall and specificity. At the top 1 position (pass@1), J-Lens surfaces the true intermediate in 36.6% of the prompts, while never surfacing a semantically similar distractor. And as we vary the top k position, the area-under-curve of true intermediates always outscore the distractors.

Despite marginal recall advantage, J-Lens does not have specificity advantage. Notice that although J-Lens has marginally higher AUC on true intermediates as compared to the logit lens, this effect is carried over to the distractors as well, lifting the scores of distractors by about the same margin. This is an empirical evidence that J-Lens may be better at surfacing the correct intermediate, but may also be more likely to surface distracting signals, which may confuse the model downstream. We will explore the ratio of signal given to the true intermediate versus the distractors in the next experiment.

Experiment 2: Receiver Operating Characteristic (ROC)

Question: We've seen thus far that J-Lens raises the signal on the entire semantic neighborhood, including both true and distractor intermediates. This begs the question: is the rate of true positives (true intermediates surfaced) versus false positives (distractors surfaced) consistent across the lenses?

• Outcome should signal the specificity of the two lenses. Specifically, if the proportion of true positives and false positives are similar, then both lenses would have similar specificity, despite J-Lens's recall advantage.

Methods

• I plot a Receiver Operating Characteristic (ROC) curve between the true positive rate (TPR) and false positive rate (FPR) of each lens for each type of false positive (adjacent vs near-miss), scored on negative discovered rank.
• The area-under-curve is then measured.
  • With this set-up, a curve that hugs the upper-left corner will achieve a high AUC (closer to 1.0), which will imply the lens is able to identify true positives with high accuracy while not marking many false positives.
  • Similarly, if an AUC collapses, that would imply the model is scoring false positives at similar rates as it is scoring true positives.

Results

(Left) The standalone J-Lens ROC curve, with an AUC of 0.83 with respect to adjacent distractors and 0.82 with respect to near-miss distractors. At k=1, J-Lens boasts an impressive TPR of 0.37 with an FPR of 0.00. At k=10, about 0.68/0.18 for both distractors. At k=25, 0.85/0.38.

(Right) ROC AUC of both lenses, overlayed. Logit lens achieves an AUC of 0.84 with respect to both types of distractors, a similar figure to J-Lens. At k=1, logit lens achieves a TPR of 0.27 with an FPR of 0.01.

Interpretations

The specificity of J-Lens is reinforced. With an average AUC of 0.83 for both types of distractors, this implies J-Lens is able to correctly identify true intermediates with limited false positive reports.

J-Lens's recall advantage, validated on pass@k, does not carry over to its specificity. As seen above, both J-Lens and logit lens achieve very similar AUC, meaning they surface true intermediates along with false intermediates at very similar proportions.

Along with the (marginal) recall advantage of J-Lens over logit lens in the previous experiment, this provides evidence that J-Lens is simply a more refined recall instrument, and not a better classifier of intermediate thoughts.

Claim 2: The results above (Claim 1) are not artifacts.

I address some potential skepticisms to the above claims and experimental findings. Apologies in advance for the "rougher" shape of this section – I am running out of time.

Skepticism 1: "Both lenses have similar specificity because they're both subjected to noise."

This is the trivial explanation for the claimed "tie" in specificity. To ensure this isn't the case, I measured the ROC-AUC of the J-Lens with respect to random vectors as a baseline. Random vectors are sampled as follows:

• Isotropic: a completely random direction in the residual space
• Covariance-based word-like token: fitting a Gaussian to the rows of "word-like" tokens of the unembedding (≥ 4 ascii characters), then sampling from that distribution.
• Random real-word: Randomly chosen row of the unembedding.

As we can see, both lenses achieve an AUC of ~0.99 within the workspace layers L18-30, against all types of randomly drawn vectors, implying the parity in specificity is non-trivial and not a vacuous result.

Skepticism 2: "The true-vs-distractor separation in Experiment 2 reflects a frequency gap between the token groups rather than content."

If true intermediates are simply more common tokens than the distractors, both lenses would rank them higher for reasons unrelated to what the residual stream is representing, causing an overstated specificity.

To ensure this isn't the case, for every true intermediate and distractor, I took its unigram count per million in WikiText-103 as a corpus-frequency proxy (not Qwen's actual pretraining distribution), and plotted it against the token's J-lens rank at the readout position.

As seen above, the three groups occupy about the same frequency range, ranging from $10^{0.5}$ to $10^{3.5}$ per million, meaning the distractors were not systematically rarer than the true intermediates.

Given the pooled log-log slope of −0.31, frequency has a mild effect on rank overall, but it does not produce the separation. Notice that true intermediates reach rank 1 across the whole frequency range, including ~10 per million, while distractors at 300+ per million still sit at rank 100–1000.

What I Verified

• Read all distractor sets by hand. Reviewed every near-miss token Fable generated for the 41 prompts. Counted overlap between near-miss and adjacent sets: 26 of 41 prompts share a token across the two conditions.

• Checked the top-1 operating point by counting. 15/41 true intermediates at rank 1 for J-lens, 11/41 for logit lens; 0/(41 × 3) distractors at rank 1. Matches the 37% / 27% / 0% in the ROC markers.

• Confirmed logit lens is $J = I$ in the code path. In the jlens library (jlens/lens.py), the code appears:

  residual = select(layer)
  if use_jacobian:
      residual = self.transport(residual, layer)
  lens_logits[layer] = model.unembed(residual)
  

  And my agent's code indeed sets use_jacobian = False when calculating logit lens related results.

• Reproduced the paper's own number before building on it. J-lens true-intermediate pass@k AUC of 0.69 vs logit 0.65 on multihop, consistent with the "modest" multihop margin the paper reports in A.6, before running any distractor experiment.

• Bugs caught. When reproducing the original paper's answers, Qwen3.5 answered only 8/93 prompts correctly. Agent claimed this was due to smaller model size and degraded reasoning capabilities. However, upon inspecting model outputs, I found that the model was overwhelmingly outputting blanks "__", "_____", "...", treating the prompts as a worksheet rather than factual recall. I then found that the agent cut the "Fact:" prefix out of the original prompts. Once this was adjusted, we arrived at the final accuracy of 48/93, which was then thinned down to 41/93 based on comparability.

Limitations & Next Steps

Small prompt set & single prompt family. Filtering the original multihop prompts by those that Qwen was able to correctly answer, we were left with only N=41 usable multihop prompts. This is a very small set of prompts.

Additionally, I only used multihop prompts. This may be a cause of the only marginal recall difference between J-Lens and logit lens (the paper even admits that the margin is "modest on multihop"). Given more time, I would repeat all experiments using different prompt families as introduced in the original paper (multilingual, poetry, order_ops, association, typos).

All experiments done were read-out only. The original paper's strongest evidence that J-Lens improves on logit lens is causal, which was discovered through ablating certain J-Lens directions. I did not do that, so the specificity parity I found was based on what the readout itself tells.

A natural next step would be to run the distractor comparison on the causal side. For example, I would ablate the distractor's lens direction, and check whether it moves the output less than ablating the true intermediate's, compared between J-Lens and logit lens.

Appendix

Supplemental materials, including more figures, tables, as well as code to reproduce the main results can be found at: github.com/zchen-66/jlens-redteam