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

Summary

TL;DR: We introduce the 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 layerwise-relevance propagation, 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. Introduction Motivation 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. Contributions 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). Background J-lens 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 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. In particular, we make use of three LRP rules: LN-rule: treat the normalization denominator mjx-math { display: inline-block; text-align: left; line-height: 0; text-indent: 0; font-style: normal; font-weight: normal; font-size: 100%; font-size-adjust: none; letter-spacing: normal; border-collapse: collapse; word-wrap: normal; word-spacing: normal; white-space: nowrap; direction: ltr; padding: 1px 0; } mjx-container[jax="CHTML"][display="true"] { display: block; text-align: center; margin: 1em 0; } mjx-container[jax="CHTML"][display="true"][width="full"] { display: flex; } mjx-container[jax="CHTML"][display="true"] mjx-math { padding: 0; } mjx-container[jax="CHTML"][justify="left"] { text-align: left; } mjx-container[jax="CHTML"][justify="right"] { text-align: right; } mjx-msqrt { display: inline-block; text-align: left; } mjx-root { display: inline-block; white-space: nowrap; } mjx-surd { display: inline-block; vertical-align: top; } mjx-sqrt { display: inline-block; padding-top: .07em; } mj