There’s two types of interpretability techniques: variable and algorithm interpretability, ones that find the features vs the circuits. J-Lens is very much about variable interpretability. We shouldn’t expect it to tell us much about how the variables are being computed, except what we can infer by viewing and intervening on the variables themselves.
There is an alternative view of the Jacobian for certain nonlinear networks where it can capture exactly how the variable or representation is being computed (for a particular input). A deep network that only uses zero-bias linear (or convolutional) layers and ReLU activation functions has a Jacobian that exactly reconstructs the ouptut computation with an input-dependent linear operator:
For interpretation of image denoising or generation, this has some very nice properties where the exact input-dependent linear filters that generate the output can be visualized.
I extended this approach to LLMs including Qwen 3 14B and Gemma 3 12B. By detaching the gradient op at inference for particular nonlinear components of the forward computation, inference can be mapped to an exact input-dependent linear transform (more specifically, one “detached” Jacobian matrix per embedded token in the prompt).
This detached Jacobian captures the whole nonlinear forward op in a set of input-dependent linear transforms on the input embedding vectors. This can also be computed for intermediate layers, or only the attention module in a layer. We can view the overall transform as the matrix product of each layer’s set of detached Jacobians. This has lots of interesting properties, like the largest right singular vectors decode to tokens like the input tokens, while the left singular vectors are candidate output tokens. We can exactly compute the dimensionality of the representation because there are no higher order terms. Since these detached Jacobian operators are exact and not approximations, they are also effective for steering concepts in responses.
There is an alternative view of the Jacobian for certain nonlinear networks where it can capture exactly how the variable or representation is being computed (for a particular input). A deep network that only uses zero-bias linear (or convolutional) layers and ReLU activation functions has a Jacobian that exactly reconstructs the ouptut computation with an input-dependent linear operator:
https://arxiv.org/abs/2310.02557 [Kadkhodaie, Guth, Simoncelli, Mallat, ICLR outstanding paper]
For interpretation of image denoising or generation, this has some very nice properties where the exact input-dependent linear filters that generate the output can be visualized.
I extended this approach to LLMs including Qwen 3 14B and Gemma 3 12B. By detaching the gradient op at inference for particular nonlinear components of the forward computation, inference can be mapped to an exact input-dependent linear transform (more specifically, one “detached” Jacobian matrix per embedded token in the prompt).
This detached Jacobian captures the whole nonlinear forward op in a set of input-dependent linear transforms on the input embedding vectors. This can also be computed for intermediate layers, or only the attention module in a layer. We can view the overall transform as the matrix product of each layer’s set of detached Jacobians. This has lots of interesting properties, like the largest right singular vectors decode to tokens like the input tokens, while the left singular vectors are candidate output tokens. We can exactly compute the dimensionality of the representation because there are no higher order terms. Since these detached Jacobian operators are exact and not approximations, they are also effective for steering concepts in responses.
https://arxiv.org/abs/2505.24293