diff options
Diffstat (limited to 'docs/willkn-jlens-greaterwrong-analysis.md')
| -rw-r--r-- | docs/willkn-jlens-greaterwrong-analysis.md | 260 |
1 files changed, 0 insertions, 260 deletions
diff --git a/docs/willkn-jlens-greaterwrong-analysis.md b/docs/willkn-jlens-greaterwrong-analysis.md deleted file mode 100644 index 7a12cbf..0000000 --- a/docs/willkn-jlens-greaterwrong-analysis.md +++ /dev/null @@ -1,260 +0,0 @@ -# Anthropic's J-Lens: A Research Engineer's Analysis - -**Author:** willkn -**Source:** https://www.greaterwrong.com/posts/vHxGD5HKsFuBStirq/anthropic-s-j-lens-a-research-engineer-s-analysis -**Date:** 24 Jul 2026 -**Notebooks:** https://github.com/willkn/jlens_re -**Note:** Reproduced here (with attribution) as reference material for this -repo's analysis. Model: GPT-2-medium (355M, d=1024, L=24, V=50,257), fp32, HF -transformers, NVIDIA L4. "This is research engineering, not mechanistic -interpretability" — author's own epistemic-status note. - ---- - -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. - -**Headline Result: Lens Monitoring is nearly free at decode time with small -dictionary size** - -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. - -## The J-Space — an overview - -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. - -## Part 1: Computing the Jacobian - -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_l = (1/N) Σ_i A_i - -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. - -### The Engineering of Computing a Jacobian - -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' - -### Compute required for a Jacobian - -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). - -(Compute formula, from the author's analysis — see source for full derivation.) - -Let ε be the target relative estimator error at layer l (the "performance p" in -matrix-space terms). c_l is the layer's noise coefficient — the constant in the -measured 1/√n convergence law, where rel_F(J_n) ≈ c_l/√n. Empirically, c_l -grows toward early layers, increasing roughly 4× from layer 20 to layer 4 on -GPT-2-medium. a + b·(L − l) captures the measured per-backward cost, which -scales linearly with depth spanned. On an L4 GPU at dim_batch = 8, we measured -a ≈ 5 ms and b ≈ 9 ms/layer. d/dim_batch gives the number of backward passes -required per prompt, with the forward pass cost t_fwd negligible by comparison. - -Two important notes come with this formula: - -1. **You pay for the earliest layer, the rest are free:** When computing an - earlier layer l, all layers where l' > l are 'free' due to the necessity of - computing them during backpropagation for layer l. -2. **The formula is only valid for ε above the layer's bias floor:** If a layer - cannot represent the non-linearity of future layers above bias floor in a - linear manner, no corpus size will fit a good Jacobian. This is a fundamental - problem with linear functions being unable to approximate non-linear - functions as they grow more complex, not a sample size issue. - -### What Quality looks like — and how to fix it - -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. - -## Part 2: Jacobian at Inference - -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. Softmax poses a problem here: if we -monitor for violent content and the user asks about the weather, softmax gives -a distribution of all the monitored words rather than their magnitude — no -signal, false positive. The vocabulary term dominates: full vocab readout of -one layer = 17% of a forward pass; 5 layers = 90%; all 24 layers of GPT-2 = 4x -a forward pass. Markedly lower in a frontier-shaped model (D=8192, L=80, -V=128k) due to V/d shrinking — ~45% of a forward pass for 24 layers, still too -expensive for an always-on monitor. - -The deployment method Anthropic cover in the paper — a fixed length dictionary -to reduce vocabulary size. Everything before the softmax is linear (through the -Jacobian approximation), so this reduces the calculation into 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**. - -As expected, wall clock overhead and FLOP share grows (nearly) linearly as -layer size increases. Kernel launch floor ≈ 1.2% (c=100 gives FLOP share -0.17%). Percentages relative to a baseline serving speed of 18.3ms/tok. - -### Can the linear transport be compressed? - -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. - -### Putting it all together - -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 but we do claim that middle layers are optimal for engineering -and computational efficiency with the J-Lens — a synergistic result. Earlier -layers are more expensive to compute (backprop through every layer that follows) -and less useful ('setting the stage'; linear approximation struggles with many -nonlinear layers). Even where logit-lens is a better tool in later layers, the -difference in compute for a single layer, single token between the two methods -is small: logit lens = 2dV; J-lens = 2d² + 2dV; difference = 2d² (~2% more -FLOPs on gpt2-medium given vocab dominance). - -### So What? - -Anthropic's J-Lens monitoring is deployment viable on small models with a -dictionary of around 1000 concepts at 2% compute overhead per token. No claims -on extrapolation to larger models, but the mathematics suggest likely more -efficient on larger models (V/d shrinks). Key engineering takeaways for smaller -models: (1) Regularise the raw Jacobian; (2) Target middle layers for best -utility/cost tradeoff; (3) Compute Jacobians for each layer simultaneously. - -### Conclusion - -The J-Lens is a technique that is certainly feasible at runtime with reasonably -sized dictionaries. It has the potential to transform how models are monitored, -understood and finetuned (Goodfire, 2026). 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 results on smaller models, where certain directions dominate without -suppression of the dominant directions. At least in this 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 Models. 24 Jul 2026, GreaterWrong. |