>Maybe we're somehow treating f(0) = 0 so that you can apply it directly
Hm thinking about it a bit more, I think what's going on is that you treat the baseline of hidden layer L at which you apply the J-lens as 0 activation, then apply the activation on top of that and see what direction your future outputs get skewed towards. Even so, you're still throwing away a constant term f(0) since the "true" logits given by the linear approximation would be Unembed{f(0) + J*h)... but I guess it doesn't matter since by linearity we have Unembed{f(0)} + Unembed{(f(h)}, and the baseline is probably just low-frequency noise (like whitespace or punctuation) which while important for actually predicting a next token matching the ground-truth data distribution, is unimportant for the purposes of interpretability of layer L activation.
And for the connection to logit-lens, as they say the J matrix is really just the change of basis matrix (or at least best linear approximation) from layer L to the final hidden layer, very similar to what you'd use in multivariable integration given change of coordinates. I guess you also need some explanation of why we can expect a linear approximation to hold even well even outside the infinitesimal regime though. I don't know enough here so I asked an LLM and it said that you can argue handwavily via the following chain.
1) Something about stein's lemma saying taking expectation of gradient gives you a global linear best-fit in the OLS sense) (this seems intuitive I guess, even if I don't know the details).
2) Because of the resnet type structure of LLMs which passes through the residual added with some delta (attention + MLP), overall residual stream doesn't undergo any "wild" nonlinearities. So it's plausible that a linear approximation might work. If you think about it, even the "naive" logit lens works fairly well.
3) Semantic meaning is encoded in angle rather than magnitude of vectors (linear representation hypothesis). I'm not sure I fully buy this outside of simple word2vec style embeddings, but assuming it holds for both the intermediary and final layers, then conversion is just a rotation, and even if the magnitudes are off it doesn't matter much for recovering the underlying concept.
Oh I guess another thing related to all of this, is prior work on steering vectors. "Manipulating the j-space" seems not too different from steering, both ultimately work on the residual stream. I think perhaps it makes more sense to think of J-space as just a coordinate system for the residual stream where each coordinate axis is a vocab direction. Compared to vector steering which was much more naive and had to derive the direction via PCA.
>The “J-space” is not a separate, hidden space. It is an alternative coordinate system for intermediate layer activations. Using a Jacobian between the last layer right before unembedding and the intermediate layer, you can “move” rows of the unembedding matrix (corresponding to distinct tokens) into the space that the intermediate activations live in. So each unembedding row/vector has a corresponding vector in the intermediate activation space this way. They use those vectors to generate coordinates for the same intermediate activations (the “J-lens”). Since each coordinate in this alternative coordinate system is now matched with a token, they can now use it to interpret and manipulate the same intermediate activations
>LLMs think in a subconscious space using tokens
narrative is completely misleading because
>(1) It's a coordinate system. Not a new/separate space.
>(2) Tokens only appear because they specifically built the coordinate system using the unembedding vectors of tokens
There is also a good companion piece by Neel Nanda [1] which answers "Why Jacobians rather than linear regression?" which was another question that came to mind
Hm thinking about it a bit more, I think what's going on is that you treat the baseline of hidden layer L at which you apply the J-lens as 0 activation, then apply the activation on top of that and see what direction your future outputs get skewed towards. Even so, you're still throwing away a constant term f(0) since the "true" logits given by the linear approximation would be Unembed{f(0) + J*h)... but I guess it doesn't matter since by linearity we have Unembed{f(0)} + Unembed{(f(h)}, and the baseline is probably just low-frequency noise (like whitespace or punctuation) which while important for actually predicting a next token matching the ground-truth data distribution, is unimportant for the purposes of interpretability of layer L activation.
And for the connection to logit-lens, as they say the J matrix is really just the change of basis matrix (or at least best linear approximation) from layer L to the final hidden layer, very similar to what you'd use in multivariable integration given change of coordinates. I guess you also need some explanation of why we can expect a linear approximation to hold even well even outside the infinitesimal regime though. I don't know enough here so I asked an LLM and it said that you can argue handwavily via the following chain. 1) Something about stein's lemma saying taking expectation of gradient gives you a global linear best-fit in the OLS sense) (this seems intuitive I guess, even if I don't know the details). 2) Because of the resnet type structure of LLMs which passes through the residual added with some delta (attention + MLP), overall residual stream doesn't undergo any "wild" nonlinearities. So it's plausible that a linear approximation might work. If you think about it, even the "naive" logit lens works fairly well. 3) Semantic meaning is encoded in angle rather than magnitude of vectors (linear representation hypothesis). I'm not sure I fully buy this outside of simple word2vec style embeddings, but assuming it holds for both the intermediary and final layers, then conversion is just a rotation, and even if the magnitudes are off it doesn't matter much for recovering the underlying concept.