*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. *
All notebooks available here: https://github.com/willkn/jlens_re Anthropic 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:
Here 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.
Throughout the writeup we will focus our analysis on memory, compute and inference speeds.
Before 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.
The 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.
Let *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’). We 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’]:
1.1* *h_L[t’] = f(h_l[t]) We can represent this nonlinear function as a Jacobian (which is linear), a collection of these outputs:
**1.2 **A_l = ∂h_L[t’] / ∂h_l[t] Every 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.
It 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:
1.3** ****A**ₗ = (1/N) Σᵢ Aᵢ This 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.
We will first look at how Anthropic determine the Jacobian in the paper.
We 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.
In one sentence:
‘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’
Setup
GPT-2-medium (355M, d=1024, L=24, V=50,257), fp32, HF transformers, NVIDIA L4. 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.
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
to <10⁻⁶ on a toy model.
**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:
2.1 C = 2· N · T · d_model · n_prompts
(where N ≈ 0.95 is an empirically fitted overhead constant and T is sequence length)
to 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
Fitting 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:
2.2
JA=J*+EA, JB=J*+EB
JA−JB=EA−EB
Where 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.
Computing 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.
We 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:
Combined, 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:
2.3
C(l, ε) ≈ n(l, ε) · t_prompt(l)
n(l, ε) = (c_l / ε)² — prompts needed
t_prompt(l) = t_fwd + (d / dim_batch) · (a + b·(L − l)) — seconds per prompt
where:
Two important notes come with this formula:
On 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.
Crucially, 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.
It 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.
Fitting 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.
We 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).
3.1.1
lens(h) = softmax(W_U · ln_f( J_l · h ))
cost per layer per position = (2d² (transport J_l·h) + 2dV (unembed))
d = model dimensionality
v = vocab size
h = residual stream
Against 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:
3.1.2
overhead = K·(d + V) / (24·L·d) Practical 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.*
The 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.
The 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.
3.2
score(c) = (J_lᵀ u_c) · h cost = 2·C·d per layer per token
For c=1000 concepts over five layers on gpt2-medium this is less than 2% of a forward pass. Below is a graph illustrating the effects of running different dictionary and layer sizes on wall-clock measurements and FLOP share.
As 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%.
Full 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.
A 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).
Anthropic 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.
Shown 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:
3.3
logit lens: 2dV (unembed only, transport is free)
J-lens: 2d² + 2dV (transport, then the same unembed)
difference: 2d²
On gpt2-medium, *2d² *would be equivalent to ~2% more FLOPs given the dominance of vocab (note that we assume a full sized vocab here).
Anthropic'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:
-
Regularise the raw Jacobian
-
Target middle layers for best utility/cost tradeoff
-
Compute Jacobian's for each layer simultaneously
In 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 (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.
References:
Gurnee, W., et al, 2026 Verbalizable Representations Form a Global Workspace in Language Model