Transformers represent belief state geometry in their residual stream

Summary

This work uses a new theoretical framework from Computational Mechanics to analyze representations in transformers. Using these tools, the authors can convert between a data-generating structure, such as one modeled with a Hidden Markov Model, and a geometric structure in a probability simplex that describes the beliefs over the states of the data-generating structure given finite observations. They then use linear readout analysis to see if the geometric structure can be linearly decoded from transformer representations and find that, indeed, they can. Thus, the authors discover that transformers represent more than just the structure of the data generating process; they also represent the current belief of the state of this process. The authors reason that this is ultimately advantageous for next-token prediction, because a model that can more accurately infer the current hidden state will be able to better predict future tokens, if not the immediately next one: “even if different belief states imply the same probability distribution over the next token, any distinction in probability distributions over the entire future may eventually, after further context, imply distinct distributions over the next token at some point in the future”.

Strengths

This is a well executed piece of research. The writing and ideas are clear; as someone without any background in computational mechanics, there are adequate examples and figures to walk through the ideas, explaining concepts such as the mixed state representation and belief state geometry. Also appreciated are the high level intuitions, such as the discussion on prediction vs generation, which helps to anchor the goal of pursuing the belief simplex in transformers (e.g., “Computational Mechanics shows that prediction is substantially more complicated than generation. Informally, the reason for this is that even if one knows the underlying hidden generative structure, one must still perform computational work in order to infer which state the generative process is in given finite observations of data.”), and the discussion on why it might be useful for a model to discover these representations (e.g., “for a recurrent neural network to perform as well as possible on next token prediction, all information about the past necessary to predict the entire future as well as possible must be present in the latent state of the network.“). While it might be easy to cast the results as obvious, they have the flavor of being obvious *in hindsight*; that transformers form these kinds of representations certainly makes a lot of sense, but the work still needs to be done to analyze that this is indeed the case, and the authors here do a great job proving it. The experiments are simple, but completely adequate. The standout is figure 6, which includes a number of careful controls which show that: the belief-state geometry emerges throughout training; the regression task withstands held-out validation; and the regression mapping can be collapsed if the input-output pairings are shuffled, demonstrating that the geometry is not spurious. I cannot be certain of this work’s originality as I’m not an expert in the field of interpretability, but I have not encountered any papers that do a similar analysis.

Weaknesses

One weakness might be that it will be difficult to demonstrate similar findings in more complex tasks, especially those for which we do not have ground truth data-generating models. This could limit the significance of the findings and the usefulness of the techniques, since they would only be applicable to toy domains. I do not think this needs to be addressed in this work necessarily, and the authors do point out as much towards the end of the paper. While I do not think a full “related work” section is needed, a few more references to past work and similar types of analyses would have been appreciated to compare and contrast the findings and techniques, and to serve as a starting point for those who wish to dive deeper into the literature.

Questions

Can a network represent multiple belief simplexes (simplices?), and if so, how fluidly might it switch between them? If, up to the current token, the “correct” data structure is still ambiguous, would the transformer represent a meta-belief state over the possible structures? How can one measure similarity to the ground truth more quantitatively? How hard is this analysis to perform in higher data-structure dimensions? How much does one need to worry about false positives? How might the analysis go wrong, or be misleading? How does matching of the ground truth belief state correlate with performance on the task?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Yes

Summary

The authors present evidence that belief states are linearly represented in the residual stream of transformer-based large language models. They also argue that belief state geometry represents information beyond next-token prediction and can be spread across the residual streams of multiple layers. This work contributes to understanding the ability of the transformer structure.

Strengths

- Their results show that only next-token predictions with the transformer structure can represent belief states. This is an important empirical contribution to the interpretability field. - The authors provide nice visualizations, such as the figure for the fractal geometry.

Weaknesses

- The formalization in Section 2.2 is very ambiguous. See the Questions below. - More experiments are needed for a larger number of hidden states and vocab sizes. - It is only for the HMM model, not language models. - The belief state $\eta$ directly affects the next word prediction, so it may be very natural. Not that surprising. As they mentioned, it might be because the models use a residual stream and were pre-trained on next-token prediction. - Even if $R^2=0.31$ in Figure 7, there is a positive correlation. I don’t understand why they said RRXOR belief state geometry is not captured by next-token predictions alone. If so, what else is the geometry captured by?

Questions

- In Figure 2, $T^{(1)}$ seems wrong. - In line 63, you said, “Computational mechanics shows that prediction is substantially more complicated than generation.” What’s the difference between prediction and generation? Does “prediction” mean the inference for the hidden state given the observations? Also, even if “computational mechanics” is a large field, you need to provide some reference for such arguments. - In Equation (1), $\eta’$ is like $\mathrm{Pr}(s_i \mid x, \eta)$ based on Bayes' rule. The notation is a little ambiguous. Especially, in line 82, what does $\mathrm{Pr}(\eta’ , x\mid \eta)$ mean? There is no $\eta’$ in the RHS. Does it mean that the probability is the same for all $\eta’$? - In line 90,  $\eta_{10}$ and $\eta_{01}$ are not defined. Maybe you should define $\eta_{x_0, x_1, \dots, x_N}$. - I think you defined the belief state as a probability distribution $\eta$. Then, what is the definition of belief state geometry? - What is Z1R in Figure 3? - In Figure 3C, how do we read each point? What’s the meaning of the point $\eta_{10}$? Does this mean, 50% for $S_0$, 50% for $S_1$, and 50% for $S_R$? I think you need to draw three axes for each one to present the probability simplex. - In line 122, why did you use the residual stream before the layer norm? - In line 168, what do you mean by “degeneracies”?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

As they mentioned, there is no extension to LLMs. This work just suggests that the residual streams contain the information of HMMs linearly.

Summary

The paper studies the geometric structure of how belief states are represented in the residual stream of transformers using relevant theory from computational mechanics. The main object of investigation of the paper is a transformer model trained on data generated using an HMM. Following this, the theory of optimal prediction is used to make predictions about the possible geometry of belief states in the residual stream. The results report that belief states are linearly represented in the residual stream and contain information about the future beyond the next token prediction.

Strengths

- Fresh perspective and well-thought-out toy model: The paper introduces a way to look at learned representations from the perspective of computational mechanics by studying the geometric structure of the belief states for a toy model trained on the task of language modeling on data from an HMM. The control exercised by the toy model allows for having a notion of the expected optimal structure as given by the optimal predictor on the task. - The geometric structure of the activations in the residual stream implies that the model accurately recovers the parameters of the HMM and that the parameters of the HMM are linearly decodable from the activations. This is a neat demonstration that a linear projection of the activations of the residual stream aligns well with the ground-truth belief geometry of the Mess3 process. - The visualizations are well made.

Weaknesses

- There is little to no discussion or intuition provided on what it would mean for the belief state geometry to have a "fractal structure"? Is it a domain-specific expectation to have given the data distribution structure imposed by the HMM? - It is unclear how the finding holds beyond data distributions generated using an HMM. How does the representation of the belief structure change for real-world / theory-of-mind tasks? How would one measure it? - It is unclear how the implicature "the residual stream contains information about the future beyond just the next token" is made. - There is no discussion of related work in the introduction and the paper, making it a bit difficult to situate this work and compare its findings in the context of other related papers.

Questions

- What aspects of the findings from the paper will carry over to real tasks? - How does the recovered geometry change over different points in the residual stream? Are the belief states consistently linearly represented?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The discussion of the limitations of the work is well done. I especially like the point the authors make about "Computational Mechanics is primarily concerned with optimal prediction, but LLMs in practice are not perfectly optimal. Further work is needed to understand how near-optimality and non-optimality influence belief state representations." Even in the context of the toy model trained on data from the HMM the ground-truth fractal structure is not "perfectly" recovered. This perhaps hints at the model learning a near-optimal representation, rather than the optimal representation. Nonetheless, very interesting!

Summary

This paper proposes that the mixed-state presentation that is grounded in optimal prediction can explain the representations and computational steps in transformers over a sequence generated by some simple hidden Markov chains. This is done by first training the transformer until convergence, and then fitting a linear regression from residual stream activations to the optimal belief state. The authors observe that this procedure mimics the theoretical prediction from mixed-state presentations and reveals similar fractal patterns. Here, what is remarkable is that there is only a single linear transformation between hidden representations of the transformer and the optimal belief state geometry. The authors also reveal that for some tasks with degeneracy, residual streams across layers need to be concatenated to have the same observation.

Strengths

1. The paper presents an extremely interesting empirical observation on the hidden activations of transformer architecture. The results seem to be well-verified. The fact that the optimal belief-state geometry appears gradually is also very interesting. 2. This empirical observation points to a mixed-state presentation as the theoretical framework for how to think about transformer computations. Crucially, here, the computation is with respect to the input sequence rather than layer to layer. This is very important for the interpretability of transformers. The authors also note the possibility of generalizing the observations in the paper to other architectures, which would be interesting to see. 3. Overall, paper is really well-written and the discussion is very insightful.

Weaknesses

1. The paper has no discussion of the related work. I like the directness of the presentation, but there must be some prior or related literature? I am adding two papers that I thought were relevant to the discussion in the paper: Bachmann, Gregor, and Vaishnavh Nagarajan. "The pitfalls of next-token prediction." arXiv preprint arXiv:2403.06963 (2024). Hu, Jiachen, Qinghua Liu, and Chi Jin. "On Limitation of Transformer for Learning HMMs." arXiv preprint arXiv:2406.04089 (2024). 2. The claims on next-token prediction L216-224 are not necessarily realized for many tasks and settings, which forces many authors to devise other schemes to train (e.g., chain-of-thought). This might be a problem that comes with the fact that one cannot minimize harder problems than the ones considered in this paper due to difficulties in optimization or the sample complexities needed, which are not at all discussed in the paper. Related to the previous point above, these should be discussed more thoroughly. 3. The findings are limited to simple problems; transformers are trained until convergence, which has very good generalization capabilities, but I think the paper has a lot of interesting insights already.

Questions

1. How does the context length interact with optimal prediction? In Section 2.2., you talk about optimal prediction w.r.t. the whole trajectory observed, or does the transformer have access to the last subsequence that is of length equal to the context size in the paper? 2. I cannot see why the MSP for Random-Random-XOR has a finite number of states. Do you have a quick explanation? It could be useful to explain somewhere in the paper for readability and accessibility.

Rating

8

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors seem to acknowledge the limitations very well in Section 4.4.

Response

Thank you for the clarification. I’ve adjusted my score.