The aim of this post is to share a quick attempt at grokking the conceptual ideas that lie behind the notion of J-space and how it is calculated in the paper Verbalizable Representations Form a Global Workspace in Language Models. Specifically, I am trying to understand Section 2.1 of the paper, rather than the abundant empirical findings. I also share a few quick takes.
I make no claim for originality or insight, and only a light claim on correctness. Nevertheless, I believe it will help researchers understand the core ideas.
Pre-requisites: knowledge of transformer architecture.
I try to re-construct the methodology by imagining how one might have come up with the ideas.
Starting question: Do LLMs have an analogue of a global workspace, in which the LLM is 'verbally thinking'?
Simplifying assumption 1: Use tokens as a proxy for what the LLM can feasibly verbally think about. (In Appendix A.9 they describe early attempts at finding multi-token concepts.)
New question: In each layer and for each token in the vocabulary, which directions in residual space correspond to 'thinking about token '
Simplifying assumption 2: As a proxy for 'thinking about token ', we look for directions in residual space which correspond to 'increases the odds of token being output in the near future'
Key methodological idea: For a given prompt , at token position and layer , we have the residual vector . To identify the directions of layer 's residual space which increase the odds of particular tokens being output in the near future, first fix some future position , then ask 'if i add a small perturbation to , what effect does that have on ?'. This is exactly what the Jacobian tells you, to a first approximation. (Talk with an LLM if you're not familiar with Jacobians.)
Simplifying assumption 3: There is a useful notion of a 'global Jacobian'. By default, the Jacobian depends on the prompt , the layer , the token positions and and the residual vector . We simplify by averaging across prompts and token positions, so we get one Jacobian per layer, .
I wonder if this assumption is one way to formalize or one ingredient in the Linear Representation Hypothesis.
Simplifying assumption 4: We treat as a small perturbation. So, in some sense, we treat the zero vector to be the 'baseline vector', so that in turn we treat as the result of perturbing the baseline vector. Thus, we interpret as an estimate of what impact has on the final layer for future tokens.
I find this assumption counter-intuitive and unjustified, and would recommend the authors of the original paper explicitly mention this in Section 2.1. If I have made a mistake or if you have a conceptual way to justify this assumption, please let me know!
Hypothesis / empirical finding: The transformer only uses a small part of the residual space to do 'verbal thinking'. Hence, the rank of should be significantly smaller than . The resulting subspace union of cones spanned by a sparse subframe is called the -space.