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.
Reverse-engineering their thinking process
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'?
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 subspaceunion of conesspanned by a sparse subframe is called the -space.
Quick takes
My gut instinct is this is a neat idea but not a big deal. The underlying questions and motivations are probably more impactful as seeds of future research, rather than the technique itself. As a comparison, I consider the Activation Oracles paper to be a big deal.
There is lots of low-hanging fruit available, especially for early-career researchers who want to upskill and/or prove their skills to other people.
On the easier side: replicating the work, creating high quality open-source code to find and use the J-space, seeing how this technique compares or can be combined with other interp techniques.
On the harder side: find clear empirical examples where the method fails, find methodological improvements that require fewer or weaker assumptions, create a 'J-space benchmark' consisting of toy non-LLM networks and LLMs, somehow using this fresh area of research as a way to test the ability of LLM agents to do novel research (I wouldn't be surprised if a significant percentage of the Anthropic paper was just Opus or Mythos being let loose.)
A vague speculation is that as deception detection techniques get used more and more and hence implicitly used as evolutionary pressure on which AIs get to exist, the 'true' reasoning of the AI will be hidden even to itself. Compare with idea of the 'consciousness as a press secretary' from The Elephant in the Brain.
Introduction
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.
Reverse-engineering their thinking process
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
Simplifying assumption 2: As a proxy for 'thinking about token
Key methodological idea: For a given prompt
Simplifying assumption 3: There is a useful notion of a 'global Jacobian'. By default, the Jacobian depends on the prompt
I wonder if this assumption is one way to formalize or one ingredient in the Linear Representation Hypothesis.
Simplifying assumption 4: We treat
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
subspaceunion of cones spanned by a sparse subframe is called theQuick takes
On the easier side: replicating the work, creating high quality open-source code to find and use the J-space, seeing how this technique compares or can be combined with other interp techniques.
On the harder side: find clear empirical examples where the method fails, find methodological improvements that require fewer or weaker assumptions, create a 'J-space benchmark' consisting of toy non-LLM networks and LLMs, somehow using this fresh area of research as a way to test the ability of LLM agents to do novel research (I wouldn't be surprised if a significant percentage of the Anthropic paper was just Opus or Mythos being let loose.)