R-lens: Making J-lens More Faithful on Early Layers

We introduce R-lens: a drop-in replacement for J-lens that produces clearer readouts on earlier layers. R-lens is identical to J-lens, except that we make minor and low-overhead changes to the backwards pass following layer-wise relevance propagation (LRP), allowing us to reduce the propagation of errors. This method allows us to surface important intermediate variables more consistently and more saliently, reduce the frequency of semantically-irrelevant readouts, and even detect relevant concepts that J-lens misses entirely. We open-source our R-lenses and accompanying J-lenses here.

The J-lens is a powerful tool for surfacing intermediate variables in the workspace layers of a model, but we find readouts in early layers to often be noisy and largely uninterpretable. Plausibly, either J-lens is degenerate at these depths and fails to resolve content that is in fact present, or the early residual stream genuinely carries no linearly accessible, causally relevant verbalizable content.

We suspect that this is a structural issue with J-lens. J-lens is fit by backpropagating from the final-layer residual stream down to the residual stream at the readout layer. Errors are likely to accumulate over the course of layers, suggesting it may be less accurate at early layers. In this post, we ask whether we can improve on J-lens to minimize such errors and achieve cleaner, causally relevant readouts at early layers.

What is RelP and how do we apply it?

We take inspiration from Relevance Patching (RelP), a method for efficiently approximating attribution patching, a gradient-based approximation to activation patching. It keeps the structure of attribution patching, which takes the dot product of an activation difference against a backward pass, but swaps the local gradient for a propagation coefficient derived from the XAI technique of Layer-wise Relevance Propagation (LRP). LRP adds several stop gradients to the backward pass, e.g., to the variance term in LayerNorm, designed to reduce the accumulation of errors/high curvature gradient terms. As this is just adding stop gradients, the cost is negligible.

We therefore use LRP as a drop-in replacement inside J-lens. We fit the R-lens exactly as before but with LRP rules installed in the backward pass, so what gets transported and averaged is a relevance coefficient rather than a raw gradient.

R-lens shows a substantial quantitative advantage over J-lens that increases as models scale, measured across a variety of evaluation categories.

R-lens produces qualitatively different readouts at earlier layers, significantly reducing the amount of incoherent tokens compared to J-lens.

Through ablation studies, we show that R-lens directions for intermediate variables are more causally important than J-lens directions.

R-lens can sometimes capture concepts that J-lens never does, especially if these concepts appear exclusively in early layers (however, because concepts typically appear in later layers as well, R-lens is primarily a tool for more faithfully tracking computation over layers).

The Jacobian lens characterizes an activation by its first-order causal effect on the model's outputs, averaged over contexts and token positions. Concretely, it replaces every layer downstream of the readout point with a single linear map, followed by the model's own unembedding, producing a ranked list of vocabulary tokens for that activation.

LRP assigns a relevance score to each component and propagates it backward through the model, redistributing relevance from a layer's outputs to its inputs according to rules specific to each component type. LRP preserves 'total relevance' from one layer to the next and constructs specialized rules to patch relevance in places where ordinary gradients would break that conservation. Notably, LRP just adds stop gradients, so it is cheap and doesn't change the output of the forward pass.

LN-rule: treat the normalization denominator as constant, making the norm linear and preventing relevance collapse. Identity-rule: detach the nonlinear factor of GELU/SiLU, so the activation's backward pass becomes a per-element linear map; i.e., replace with stop_grad(GELU). Half-rule: split relevance evenly across a multiplicative gate's two branches instead of double-counting through the product.

We use a similar setup to dense models, with a few modifications to account for the architectural differences: We extend the LRP rules to all routed experts (i.e., we treat all experts as if they were normal MLP layers). We freeze the expert routing weights (+softmax/sigmoid scoring, router logits, and gate projection). Some models have a shared expert, which is always on. We multiply the shared expert's output by a constant and sweep over a few values of the constant, which amplifies the shared expert's importance. Note: DeepSeek-v4-flash's residual stream uses manifold-constrained hyper-connections (mHC), which means that the state carried between sublayers is a matrix. For this model, we additionally freeze the mHC residual-mixing coefficients.

Our evals are similar in content to those described in Appendix A.6 of the J-lens paper. We probe on the bolded token position: Association: a short passage that evokes an unnamed concept. Example: Prompt: "She couldn't buy groceries anymore without strangers whispering, pointing, and holding up their phones." Intermediate: "fame".

For each model we filtered for the questions the model was capable of answering correctly (for multihop and multilingual). Mean pass@10 over layers and categories (i.e., in layer x, does the intermediate appear in the top 10 J-lens readout at the target position).

R-lens shows a substantial quantitative advantage over J-lens as models scale. There is no R-lens advantage for both the smallest dense and MoE models we tested. But the R-lens for every other model shows an improvement over J-lens, for both pass@10 averaged over the first half of layers and averaged over all layers. This advantage increases with model size, with the R-lens for DeepSeek-V4-Flash (a 284B-parameter/13B-active MoE model) showing the largest advantage. Eval performance across layers and for individual eval categories shown below for each model.

In early layers, J-lens tends to contain what we refer to as "trash tokens": tokens that are seemingly non-semantic, incoherent, or unrelated to the prompt. E.g., “锁定” ("locking"), “尷" (half of "awkward"), “＊＊＊＊＊＊＊＊”, “ ...\n\n", “......”, “euw”, “tav”, “zinho”, etc.