The details seem to be present in the paper (section 2.1). I'm still trying to understand, but it seems instead of computing gradients with respect to cross-entropy loss for the 1-hot "next word" vs output logits, you compute the gradient for the last hidden layer with respect to some middle layer L. This gives you a `hidden x hidden` jacobian matrix, hence the "J-lens". They don't just do this for the last hidden layer of the current token, but the last hidden layers of all subsequent tokens too, and average them. And then repeat for a bunch of documents like in pre-training.
It's still not clear to me intuitively what this represents though. I get that it somehow encodes a link between future words the model says and the current activation, but the confusing thing is that I always think of derivatives and gradients as basically a "sensitivity" between output & input, i.e. if you nudge the input x by h, the output changes by h * f'(x). So then on the face of it applying the J-lens matrix directly to a given activation rather than a small perterbation seems like a "type" issue.
Maybe we're somehow treating f(0) = 0 so that you can apply it directly? Or is there some shift invariance somehow? ignoring that, I do see how it's like selecting a linear combination of the directions, and then it can maybe be represented as "possible continuations" in the same way the gradient is usually thought of as tangent space. Maybe that's what the other commenters meant by information geometric approach.
Other things i'm not clear about is how this is related to two other interprability methods:
* SAE (sparse auto encoder) they showed a few months back, where you train an autoencoder directly off of the hidden states/residual stream to convert it into words. The doc only mentions it briefly, but it seems that j-lens is sensitive to things that SAE are not. They're both working off of the same residual stream so clearly the the inputs must be there, but for some reason SAEs can't detect it while J-lens can (they seem to hint at some explanation but it's over my head)
* Logit lens. This was a more primitive technique that simply applies the unembedding matrix directly to the residual stream. I do like that they mention it:
>The J-lens can be understood as a principled refinement of the logit lens. While the logit lens assumes that representations use the same coordinates in all layers, the Jacobian lens corrects for representational changes that take place across layers, allowing it to uncover meaningful information in earlier layers where the logit lens produces uninterpretable readouts... The J-lens can be understood as the principled correction: J_l is precisely the average linear map that relates layer-l directions to their final-layer counterparts.