{"slug": "anthropic-s-j-lens-a-research-engineer-s-analysis", "title": "Anthropic's J-Lens: A Research Engineer's Analysis", "summary": "Anthropic's J-Space technique, which accesses a model's internal workspace via Jacobian computation, requires significant memory and compute overhead for production deployment, according to a research engineer's analysis of the gpt2-medium model. The analysis measured memory, compute, and inference speeds, finding that the Jacobian computation adds substantial cost per forward pass, though exact figures depend on model size and batch configuration.", "body_md": "*Epistemic status: this is research engineering, not mechanistic interpretability . The compute/cost claims are measured or derived from architecture constants. The quality claims (faithfulness comparisons, spectral channel interpretations) are from one small base model (gpt2-medium), one metric, and in places small samples (n=32); I make no claims about whether the J-space constitutes reasoning or a workspace and this post is about determining what it costs to run the tool in production environments, not the tool's outputs. *\n\nAll notebooks available here: [https://github.com/willkn/jlens_re](https://github.com/willkn/jlens_re)\n\nAnthropic recently released a paper on the transformer circuits platform called ‘Verbalizable Representations Form a Global Workspace in Language Models’. This paper posits that models have an internal workspace where non-verbalised concepts, perhaps certain reasoning steps or other intermediary computations, exist. Anthropic call this the *J-Space*:\n\nHere is an example from the paper — we see that the model holds certain values in the J-Space when computing an output. One of the core discussions ongoing around this paper is whether the values we observe in the J-Space are useful and if they constitute reasoning. We will not engage with that discussion in this post. Although, I feel that these values have a lot of potential and imagine many practitioners will be looking to implement the techniques into their own experiments or production systems, so this writeup serves as a first venture into analysis of production-level engineering with the J-Space.\n\nThroughout the writeup we will focus our analysis on memory, compute and inference speeds.\n\nBefore we dive into the mathematics, at a high level, how do we access the J-Space? The residual stream is used to determine an output in the final layer by providing a set of logits over a vocabulary, such as the English language. This however is not possible within middle layers as they do not share the same understanding of the residual stream as the final layer does. The paper therefore aims to understand these layers as we do the final layer. To do this, we create an object, *the Jacobian*, that measures how much the final layer changes (and in which direction) based on some change at an earlier layer. For example, we may observe that an earlier layer is making it more likely that the output will be related to a certain concept.\n\nThe formal mathematical definition of the J-Space gives us instant intuition on whether the techniques used are feasible for production environments. To access the J-Space we need to firstly compute a Jacobian, and secondly, apply it at inference time.\n\nLet *h_l[t] *be the residual stream vector of dimension *d* at layer *l*, position *t* Everything downstream of this vector is a function mapping *h_l[t]* to the final layer residuals at *h_k[t’], *where *t’ *≥ *t* since we are using a causal mask (earlier positions can’t attend to later positions to stop the model ‘cheating’).\n\nWe then define our Jacobian. For one prompt and one position pair, we run the prompt through the transformer blocks (which can be thought of as a non-linear function *f) *which produces a final output *h_k[t’]:*\n\n**1.1*** *h_L[t’] = f(h_l[t])\n\nWe can represent this nonlinear function as a Jacobian (which is linear), a collection of these outputs:\n\n**1.2 **A_l = ∂h_L[t’] / ∂h_l[t]\n\nEvery entry in this Jacobian, for example *A[i, j], *tells us how much coordinate *i *of the final output vector at position *t’ *were to shift if we made a small change in the residual stream at layer l. Practically, this encodes the intermediate computations that happen in a forward pass and lets us know what exactly each token changes about the output.\n\nIt is helpful to think of intermediate layers in terms of the final layer. In the final layer, we output a word based on the value of the residual vector. In intermediate layers, we have no such thing, and the value of a residual vector in layer 5 might not mean the same thing it does in the final layer. The Jacobian we have here allows us to link non-final layers to the final layers and extract meaning that otherwise would not be found — analogous to a linear map (lossy, non-reversible). However, computing the Jacobian with one prompt gives us a biased view due to taking on individual ‘characteristics’ of that prompt, so we must compute with multiple (n=1000 in the Anthropic paper) to get a better representation:\n\n**1.3**** ****A**ₗ = (1/N) Σᵢ Aᵢ\n\nThis is the general definition of how to obtain a Jacobian for the J-lens. Let's go into some of the engineering tricks Anthropic used to compute theirs in the paper.\n\nWe will first look at how Anthropic determine the Jacobian in the paper.\n\nWe call backprop once per output coordinate: inject a one-hot gradient at coordinate i of the final-layer residual (at every valid target position at once) and backpropagate to layer l; each backward returns row i. We then stack these rows to create our final Jacobian. Since backprop passes through every intermediary layer *l’ *where *l’ > l*, we can pick up those too and build Jacobians per layer without having to rerun independently. We also need to remember to take averages when computing to destroy the noise that is accumulated by individual prompts or tokens.\n\nIn one sentence:\n\n‘**Run backprop once per output coordinate to harvest the Jacobian row by row, with positions averaged inside each pass and prompts averaged across passes — d backwards per prompt, and the average of it all is the J-lens matrix’**\n\n**Setup**\n\nGPT-2-medium (355M, d=1024, L=24, V=50,257), fp32, HF transformers, NVIDIA L4.\n\n*Fitting:* 128-token WikiText-2 prompts; dim_batch=8 (output dimensions per backward, prompt replicated along the batch axis); first 16 positions excluded (attention sinks); fp32 accumulation; timings synchronised, warm-up passes discarded.\n\n*Inference:* greedy KV-cached decode, batch 1, 200 new tokens, best of 3; monitor = precomputed W_U x J_L applied in a forward hook. Estimator verified against `torch.autograd.functional.jacobian`\n\nto <10⁻⁶ on a toy model.\n\n**Fitting Cost: **The Jacobian requires *d_model* backward passes per prompt (one vector-Jacobian product per output dimension/coordinate, all layers computed in one backward). Measured on gpt2-medium (*d_model*=1024, 355M parameters) on an L4 GPU the observed fitting time followed a cost model of:\n\n2.1 C = 2· N · T · d_model · n_prompts\n\n(where N ≈ 0.95 is an empirically fitted overhead constant and T is sequence length)\n\nto 5% accuracy, with noise due to overhead. We observed a backward:forward FLOP ratio of 0.95 (we don’t compute parameter gradients since we aren’t updating weights, hence we don’t see the typical 2:1 ratio) but a wall clock ratio of 1.30 (backwards kernels are less efficient). We identified the optimal dim_batch of 8, fitting costs 22.4 s/prompt at 12% MFU, i.e. 8.2×10¹⁶ FLOPs / ~6.7 h for the paper's 1000-prompt recipe\n\nFitting with checkpoints allows us to test the quality/cost tradeoff in one run. We used wikitext-2 as our corpus and saw that convergence followed a 1/√n samples law and also beat cross-half floor noise at n=100 prompts, concurring with Anthropic’s claim that 100 prompts is the saturation point for determining a Jacobian. Cross-half floor noise is defined as such:\n\n2.2\n\nJA=J*+EA,\n\nJB=J*+EB\n\nJA−JB=EA−EB\n\nWhere *Jᵢ* is a Jacobian estimated with dataset *i *and *Jⱼ* is a dataset* j*, and there is no overlap between sets *Sᵢ ∩ Sⱼ = ∅. *This leaves us purely with noise, but of course the complement is the useful component of the Jacobian. This lets us know roughly how similar our Jacobian is to our true optimal Jacobian, without computing the true optimal Jacobian.\n\nComputing the Jacobian over a corpus of 100 prompts took 40 minutes on an L4 GPU, but not all layers are made equally. Firstly, we observe that convergence is depth ordered. layer 4 has ~4x the error of layer 20. We are yet to see whether earlier layers plateau due to the Jacobian’s inability to model non-linear processes to the required extent. This likely contributes significantly to early-layer convergence time and hints at the fact that there may be an optimal stopping algorithm available, but we leave this to future work.\n\nWe also observed that layers that are earlier in the model generally take longer to converge on agreement with other layers, we give two possible reasons for this:\n\nCombined, these give us something to consider when training Jacobians for early layers. The pricing model for training a Jacobian is set by the earliest layer that we use. If we want to achieve performance *p *on a layer *l,* we must compute intermediary layers through backpropagation to get to layer *l* in the first place. Therefore, we have a pricing model that looks approximately like this:\n\n2.3\n\nC(l, ε) ≈ n(l, ε) · t_prompt(l)\n\nn(l, ε) = (c_l / ε)² — prompts needed\n\nt_prompt(l) = t_fwd + (d / dim_batch) · (a + b·(L − l)) — seconds per prompt\n\nwhere:\n\nTwo important notes come with this formula:\n\nOn GPT-2-medium, the raw fitted Jacobian underperforms logit lens on next-token faithfulness. Spectral analysis reveals the cause: the Jacobian's dominant channels carry ~10× the gain of the residual pathway, misweighting structural tokens (grammar, punctuation) over semantic content. This is not a signal problem, but rather a transport weighting problem.\n\nCrucially, this is fixable. A single-parameter shrinkage regularizer J + λI monotonically recovers faithfulness across all layers. At layer 12 specifically, this simple fix doesn't just match logit lens, it exceeds it (0.294 vs. 0.275). This reveals an important engineering requirement: practitioners using J-lens monitoring must regularise the raw Jacobian to generalize properly. The mechanism is straightforward and the improvement is material.\n\nIt should be stressed that the transport weighting issue was only identified on gpt2-medium. This does not mean that the issue does or does not exist for other models, we merely provide a fix for gpt2-medium that may work for other models with the same phenomenon.\n\nFitting is one off — once we have the Jacobian we do not need to recompute it. Therefore, it is just a case of applying it at inference time.\n\nWe assume here that we apply monitoring at every token. The lens readout at one layer and one position is two matrix vector products (transport then decode).\n\n3.1.1\n\nlens(h) = softmax(W_U · ln_f( J_l · h ))\n\ncost per layer per position = (2d² (transport J_l·h) + 2dV (unembed))\n\nd = model dimensionality\n\nv = vocab size\n\nh = residual stream\n\nAgainst the model's own per-token forward cost of 24·L·d²monitoring K layers with a full-vocabulary readout at every generated token adds:\n\n3.1.2\n\noverhead = K·(d + V) / (24·L·d)\n\nPractical use note:* Anthropic use a method where only certain words are considered in monitoring for the J-Lens, for example if we only care about monitoring for violent use cases we might monitor for that. Softmax poses a problem here. Lets imagine we are monitoring for violent content in the lens. A user might ask what the weather is like in London today — no violent content here — but using softmax here would give us a distribution of all the words we are monitoring rather than their magnitude, which in this case contains no signal and we have a false positive on our hands.*\n\nThe vocabulary term dominates here. If we were to run this over a complete vocab, the computation for per token monitoring of one layer would cost 17% of a forward pass. 5 layers, 90%. The full 24 layers of gpt2, 4x the compute of a single forward pass. It is worth noting that this would be markedly lower in a frontier shaped model (e.g. D=8192, L=80, V=128k) due to V/d shrinking, we’d be looking at something more like 45% of a forward pass for 24 layers, which is still far too expensive for an always on monitor.\n\nThe deployment method is one that Anthropic cover in the paper — creating a fixed length dictionary to reduce the vocabulary size. Everything before the softmax is linear (through our Jacobian approximation), so this reduces the calculation into a dot products of the residual against precomputed J-Lens vectors, one per watched token.\n\n3.2\n\nscore(c) = (J_lᵀ u_c) · h\n\ncost = 2·C·d per layer per token\n\nFor c=1000 concepts over five layers on gpt2-medium this is *less than 2% of a forward pass*.\n\nBelow is a graph illustrating the effects of running different dictionary and layer sizes on wall-clock measurements and FLOP share.\n\nAs expected, wall clock overhead and FLOP share grows (nearly) linearly as layer size increases. We also observe a kernel launch floor of approximately 1.2%, given that when c=100 the FLOP share is 0.17%.\n\nFull vocab readouts actually cost much less than their FLOP share implies, but this is likely the batch size of 1 leaving the GPU idle enough to absorb them (I will not pretend to be an expert on GPUs and kernels — this is outside the scope of this analysis, but a good question that should be answered). Percentages are relative to a baseline serving speed of 18.3ms/tok.\n\nA natural idea is to try and reduce the rank of the Jacobian to reduce the computation required. However, on this model this does not seem to be a particularly viable. To capture 90% of spectral variance requires 562-858 dimensions; the remaining 10% spreads across the last 166-462 dimensions dependent on layer. The Jacobian is essentially full-rank, leaving little room for compression. I posit that similar size models likely have the same problem due to being overcomplete with features, but it is yet to be seen if there is an opportunity to reduce the rank with larger models (I'd be interested in results on larger models, as those results could change the feasibility calculus significantly).\n\nAnthropic claim in the paper that middle layers are where the J-Lens works best. We make no claim about the quality of the concepts verbalised in these middle layers (though Anthropic give ample evidence that these are the prime layers for the tool) but we do claim that middle layers are optimal for engineering and computational efficiency with the J-Lens — a synergistic result.\n\nShown earlier in the piece, earlier layers are more expensive to compute due to backpropagation needing to go through every layer that follows it. However, Anthropic claim that they aren’t of as much use as middle layers (they are ‘setting the stage’ for future computations and the linear approximation struggles to approximate that many nonlinear layers), so there is less practical application here. Even if the logit-lens is a better tool in later layers, the utility of knowing the state of a model in final layers (when research suggests it is not materially changing the answer as much as final layers) is relatively lower. Furthermore, the difference in compute for a single layer, single token between the two methods is small:\n\n3.3\n\nlogit lens: 2dV (unembed only, transport is free)\n\nJ-lens: 2d² + 2dV (transport, then the same unembed)\n\ndifference: 2d²\n\nOn gpt2-medium, *2d² *would be equivalent to ~2% more FLOPs given the dominance of vocab (note that we assume a full sized vocab here).\n\nAnthropic's J-Lens monitoring is deployment viable on small models with a dictionary of around 1000 concepts at 2% compute overhead per token. We make no claims on the extrapolation of the post's equations to larger models, but the mathematics suggest we are likely to see more efficient performance on larger models (V/d shrinks). The key engineering takeaways for working with smaller models are:\n\n1) Regularise the raw Jacobian\n\n2) Target middle layers for best utility/cost tradeoff\n\n3) Compute Jacobian's for each layer simultaneously\n\nIn conclusion, the J-Lens is a technique that is certainly feasible at runtime with reasonably sized dictionaries. This technique has the potential to transform how models are monitored, understood and [finetuned](https://x.com/GoodfireAI/status/2077073005088501780) (Goodfire, 2026). This is all with the added benefit of taking relatively few FLOPs as a percentage of total FLOPs, and being inexpensive to train, even at frontier model levels. However, this must be read with the caveat of our results on smaller models, where certain directions dominate without suppression of the dominant directions. At least in our implementation, this was a real emergent phenomenon that needed addressing to get the J-Lens to work in smaller models. It is implied that this was not a problem in Claude’s 4.5/4.6 generation of models through Anthropic’s work, but it is yet to be seen if this happens with other models.\n\nReferences:\n\nGurnee, W., et al, 2026 Verbalizable Representations Form a Global Workspace in Language Model", "url": "https://wpnews.pro/news/anthropic-s-j-lens-a-research-engineer-s-analysis", "canonical_source": "https://www.lesswrong.com/posts/vHxGD5HKsFuBStirq/anthropic-s-j-lens-a-research-engineer-s-analysis", "published_at": "2026-07-24 00:57:19+00:00", "updated_at": "2026-07-24 01:29:23.114035+00:00", "lang": "en", "topics": ["artificial-intelligence", "ai-research", "ai-infrastructure", "large-language-models"], "entities": ["Anthropic", "gpt2-medium", "J-Space", "Jacobian"], "alternates": {"html": "https://wpnews.pro/news/anthropic-s-j-lens-a-research-engineer-s-analysis", "markdown": "https://wpnews.pro/news/anthropic-s-j-lens-a-research-engineer-s-analysis.md", "text": "https://wpnews.pro/news/anthropic-s-j-lens-a-research-engineer-s-analysis.txt", "jsonld": "https://wpnews.pro/news/anthropic-s-j-lens-a-research-engineer-s-analysis.jsonld"}}