Transformers Represent Belief State Geometry in their Residual Stream
Authors: Adam Shai
Lucas Teixeira
Alexander Gietelink Oldenziel
Sarah Marzen
Paul M. Riechers
Venue: NeurIPS 2024 poster (2024)
Keywords: Interpretability
Computational Mechanics
Belief State
Features
Representation
Primary area: interpretability_and_explainability
TL;DR: Transformers trained on next-token prediction learn to represent the geometry of belief updating over hidden states of the data-generating process in their residual stream.

Abstract: What computational structure are we building into large language models when we train them on next-token prediction? Here, we present evidence that this structure is given by the meta-dynamics of belief updating over hidden states of the data- generating process. Leveraging the theory of optimal prediction, we anticipate and then find that belief states are linearly represented in the residual stream of transformers, even in cases where the predicted belief state geometry has highly nontrivial fractal structure. We investigate cases where the belief state geometry is represented in the final residual stream or distributed across the residual streams of multiple layers, providing a framework to explain these observations. Furthermore we demonstrate that the inferred belief states contain information about the entire future, beyond the local next-token prediction that the transformers are explicitly trained on. Our work provides a general framework connecting the structure of training data to the geometric structure of activations inside transformers.

Transformers Represent Belief State Geometry in their Residual Stream
Adam S. Shai Simplex PIBBSS
Sarah E. Marzen Department of Natural Sciences Pitzer and Scripps College
Lucas Teixeira PIBBSS
Alexander Gietelink Oldenziel University College London
Paul M. Riechers
Simplex
Timaeus
BITS
Abstract
What computational structure are we building into large language models when we train them next-token prediction? Here, present evidence that this structure is given by the meta-dynamics of belief updating over hidden states of the data- generating process. Leveraging the theory of optimal prediction, we anticipate and then find that belief states are linearly represented in the residual stream of transformers, even cases where the predicted belief state geometry has highly nontrivial fractal structure. We investigate cases where the belief state geometry is represented the final residual stream or distributed across the residual streams multiple layers, providing framework to explain these observations. Furthermore we demonstrate that the inferred belief states contain information about the entire future, beyond the local next-token prediction that the transformers are explicitly trained on. Our work provides a general framework connecting the structure training data to the geometric structure of activations inside transformers.
Introduction
In this work, we present a rigorous and concrete theoretical framework that connects the structure of training data to the geometry of activations in trained transformer neural networks. Our framework is grounded the theory of optimal prediction. suggests that transformers pretrained on next-token prediction will develop internal structures characterized by the meta-dynamics of belief updating hidden states of the data-gen genttatnneratng process. In other words, pretrained models should learn more than the hidden structure of the data generating process they must also learn how to update their beliefs about the hidden state of the world as they synchronize to it in context.
To test this framework, we conduct well-controleed experiments where we train transformers on data generated from processes with hidden ground truth structure, and then use our theory to make predictions about the geometry of internal activations. Even in cases where the framework predicts highly nontrivial fractal structure, our empirical results confirm these predictions (Figure []. This provides comprehensive explanation of how transformers encode information beyond the local next-token predictions they are explicitly trained on.
In Section 2 we review the relevant theory from computational mechanics which motivates our prediction the geometry of activations in the residual stream. Core to our work is the mixed-state presentation, which describes the metadynamic of beliefs in the probability simplex over states of the
Principles of Intelligent Behaviour in Biological and Social Systems
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Data Generating Structure
Theoretical Prediction
SCCACCMMB4444998
Residual Stream Geometry
Training
Figure I: (Top) Given a hidden data-generating structure, our framework predicts unique belief state geometry in probability simplex. Often these have highly nontrivial fractal structure as shown this example, (Bottom) Our main experimental result is that we find that the fractal geometry of optimal beliefs is linearly embedded in the residual stream, and emerges over the course of training.
data generating process. This formalism leads to a geometric structure that forms a natural hypothesis for the internal states of neural networks trained on prediction tasks.
In Section we verify that this geometry is linearly represented the residual stream of transformers. We implement two main experiments that control for different aspects of the hidden structure of the data generating process and make concrete predictions about the internal states of networks trained on these datasets. In some cases the geometry of belief state updating is found in the final residual stream, while others it is spread out across multiple layers. We detail how our framework explains this phenomenon. After demonstrating the initial success of this framework, Section 4 discusses the theory and implications in more depth.
2 Theory and methods
2.1 Data generating processes
In this study, we assume that our training data generated by an edge-emitiing hidden Markov model (HMM Paths through the HMM produce sequences of tokens from predefined vocabulary. Tokens (x E X) are emitted as we transition between hidden states (8 E S). These transitions are governed by token-labeled transition matrices (T(e)), where each T! Pr(a, i sy(ş;) represents the joint probability of emitting token z and transitioning to state sy. given that the HMM was in state s,,
As simple example, consider the HMM shown in Figure 2] that we call the Zero-One-Rande (ZIR) process. The ZIR process generates strings of tokens of the form ... www.01111.... where is
a randomly generated or 1. This HMM has states: So. S,, and Sg. Arrows of the form m-p8 s' denote the probability Pr(z, s) p% of moving to state s' and emitting the token x, given that the process was state
2.2 Mixed-state presentations
There are competing intuitions for what transformers should represent. Are they stochastic parrots 2 Do they build a world model [13]? One natural intuition would be that transformers represent the hidden structure of the natarggnn process-a a world model-in order to predict the next token well. For instance, Ilya Sutskever has said: "Predicting the next token well means that you understand
For arbitrarily large, possibly non non-ergodic HMMs with arbitrary initial distribution over the latent states, this does not limit the sophistication of training data.

---

Hidden Markov Model
Transition Matrices
/0
T
0
0:308
4
1ON
11300
T
0
0
Example Training Data
0.0100001100001101100...........
Figure 2: An illustration of a hidden Markov model (HMM) and its components. The left side shows the HMM with states So. St, and S3, and their respective transition probabilities. The right side displays the transition matrices T and T(1) corresponding token emissions 0 and 1. Example training data is provided at the bottom, demonstrating a sequence generated by the HMM.
the underlying reality that led to the creation of that token 32 " This type of intuition is natural, but not formal.
Computational mechanics is a formalism that was developed in order to study the limits prediction in chaotic systems, and has since expanded to deep and rigorous theory of computational structure for any process [4] One of its contributions is in providing a rigorous answer to what structures are necessary to perform optimal prediction B1I.
Interestingly, computational mechanics shows that prediction substantially more complicated than generation [19 30 Informally, the reason for this is that even if an agent knows the underlying hidden generative structure that creates the data they are trying to predict, the agent must still perform computational work in order to infer which state the generative process is in, given finite observations of data. Thus, there are conceptually two distinct inference processes related to prediction: learning the hidden structure of the data generating process, and subsequently inferring which state the data generating process is in given finite observations of data.
Computational mechanics formally captures the computational structure of the latter inference process with the mixed- state presentation (MSP). Whereas the HMM of data generating process describes a generative structure BigureBA), the MSP is the structure of prediction (Figure The MSP answers the question of how an optimal observer updates their beliefs over the states of the data-generating process given finite observations of tokens [27 15. If the observer is in belief state given by probability distribution r (a row vector) over the hidden states of the data generating process, then the update rule for the new belief state 1/' given that the observer sees new token is:
nT
(I)
T)
where 1 is column vector of ones of appropriate dimension, with the denominator ensuring proper normalization of the updated belief state. In general, starting from the initial belief state 7g. we can find the belief state after observing sequence of tokens Tp, II,...,,N:
No ((e)) ... T(2N)
7
7 (sol)(e) ..ne).......... T(#N)
This process belief updating itself another HMM, where hidden states are associated with belief states, and paths leading to particular belief states are the observed token sequence. The probability of transitioning from belief state 7 to n' = n7(((=) (r][(==) 1 of Eq. [I is simply given by the probability r)T()) of observing an x from For stationary processes, the optimal initial belief state is given by the stationary distribution n π over latent states (the left-eigenvectorr of the transition matrix T==E. T(x) associated with the eigenvalue of I).
Each belief state of the MSP is probability distribution over the states of the data generating process. Belief states thus inherit a natural geometry; we can plot these belief states in probability simplex, as indicated in Figure We refer to the geometric arrangement of the points on the simplex as the belief state geometry. Note that each belief state-as a distribution over generator states induses a probability density over all possible futures. Distributions over the entire future clearly carry more than just next-token information. We can see this in our example in 3..ureBD. The states π, 720 and Noi all have the same next-token prediction despite being distinct belief states.
3

---

A
Process Ganerativa Structure
B
Mixed-State Presentation Predictive Structure
Distributions Probability Simplex
D Mixed-State Singlex
Predictive Geonetzy
1008 Sn
(())) P(S)
((S)
100% Sp
100%
Figure 3: (A) An example generative structure called the zero-one-random process (Z.IR), since generates data of the form- 01R01R01R. where R is random bit. (B) The generative structure implies unique metadynamic over belief states as a predictor synchronizes the hidden state of the world as it observes more context. This predictive structure is called the mixed-state presentation (MSP). We label belief states with 7u where w is the shortest string of emissions which leads to that belief state. (C) Belief states are distributions over generator states and can be embedded in a probability simplex. To read off this distribution for a given belief state, 77, one measures the perpendicular distance from the point to each edge of the simplex, shown as blue, green, and red lines this example. These distances directly give the probabilities for each state. Thus, the vertices represent states of certainty only one of the edges of the simplex. (D) Plotting the belief state distributions in the probability simplex gives the belief state geometry.
generator state, since the perpendicular distance is nonzero to
2.3 Finding the belief' simplex in the residual stream
As described in the previous section, the MSP and belief state geometry are formalizations of the computational structure of the belief updating process an observer must perform the service of optimal prediction. Consequently, a natural hypothesis would be that transformers, trained on data generated from given ground-truie truth HMM, will represent the MSP structure internally. Our procedure for finding the belief states in the residual stream of a pretrained transformer is depicted schematically in Figure [
We begin by training a transformer on data generated by ground-truth HMM. We consider activations of the residual stream (in either single layer, or a concatenation of layers) induced by all possible input sequences (Figure 4AB), and at all context window positions. Thus our dataset consists of of activations, IR Bmed when we consider a single layer given position, where dresid the residual stream dimension. For each input sequence, our framework tells us which belief state the input corresponds to, and the probability distribution over hidden states of the generative process that corresponds to this belief state. These probability distributions can be thought of as points b R [R|)) where (S] denotes the number of hidden states in the generative process.
Since every input has belief state associated with if we can label (or color, as in Figure 4C) each activation by the associated ground- belief state. We then use standard linear regression to find an affine map from a to of the form & Av Wa + e, where W E III xdes is a weight matrix and c E R|S) is bias vector. Parameters W and e are found by minimizing the mean squared error between the predicted belief states and the true belief states. The matrix W projects from the residual stream to (as best as possible) the probability simplex, shown diagrammatically in Figure 4CD.
3 Results
3.1 Belief state geometry is linearly represented in the residual stream
To investigate the computational structure learned by transformers we conducted an experiment using a simple 3-state hidden Markov model (HMM) called the Mess3 Process (00 (Figure 5). We used the data generated by this process (Figure 5 train a transformer model.
In general, multiple inputs will lead to the same belief state.

---

B Rosidual activations at all context window positions for all imputs
rind linear projection that best presorvos simplex geonetry
Oztgut
Piot the projection
Recidual
MLP
Actantion
Altoesion
Ixput
Figure 4: To verify if transformers represent belief state geometries in their residual streams, we record (A) residual stream activations at all context window positions over all inputs. (B) These activations live in a high dimensional space. (C) Each input has ground-truth optimal belief state, which is probability distribution over states of the data-generating process. In this way we can label, or color, each activation by the ground-truth belief associated with the input. (D) Using linear regression we then find a linear subspace of the activation space that best preserves the belief state geometry of the simplex.
A Mess3 Process
Belief State Geometry
C Residual Stream
BBCCACCCBBBBMABBBBBB....
Figure 5: The residual stream of trained transformers linearly represents the belief state geometry of the mixed-state presentation. (A) The Mess3 Process has 3 hidden states and generates sequences in a token vocabulary of (A. B, C}. (B) The ground truth belief state geometry of the Mess3 Process has intricate fractal structure. Each point in this plot is belief state a probability distribution over the hidden states of the Mess3 Process. Points are colored by taking the belief probability distribution and using them as RGB values. (C) We find linear projection of the final residual stream activations contains a representation of the ground-trutt belief geometry. Points are colored according to the ground-truth belief states.
The Mess3 MSP has a fractal structure, shown in Figure providing a highly non-trivial test of our theory's prediction. Each point in this geometry corresponds to probability distribution over the hidden states of the data generating process, and thus lies in a 2-simplex.
To test these predictions, we analyzed the final layer of the transformer's residual stream, before the layer norm and unembedding. Using linear regression, we identified 2D subspace of the 64-dimensional residual activations that best matched the ground-truth belief distributions from the MSP. Remarkably, the geometry of this 2D subspace closely resembled the predicted fractal structure (Figure 5C), providing strong evidence that the transformer had indeed learned to represent the geometry of the belief states in residual stream
We ran a number of controls, shown in Figure 6 to make sure these results were not artifacts. First, we performed our analysis at multiple points through training and found that the simplex structure emerged gradually (Figure (AA), suggesting that the detailed fractal structure found at the end of training was not trivial consequence the model architecture or initialization. Second, quantifying the mean squared error of the regression revealed a decrease in the errors of belief-state geometry representation throughout the course of training (Figure 6D, first bars). In addition, more faithful
'These results hold when analyzing the activations after the final layer well, Figure S1
5

---

A
Training
D
Traning
Creas Vol
Stuflo
Shuffle
s.00s eaiss
a.eets baioa anios o3i3t
Mean Emor
Figure 6: The representation belief state geometry is nontrivial. (A) Projected activations at different stages of training shows the emergence of belief state geometry. (B) cross-validation of our main result. (C) We shuffle the belief states in our linear regression procedure, preserving the overall ground truth fractal shape while getting rid of the associated context. The projection collapses the data, showing that the fractal's 's appearance in the residual stream is not an artifact of projecting high- nigh-dimmnsiooal data to desired shape. (D) Mean squared error (averaged across input sequences) hetween (i) the position of projected activations and (ii) the ground-truth position of the corresponding belief state.
representation the belief state geometry in the residual stream corresponded to lower cross entropy loss (Figure S2). Third, performed cross-vale validation, where only 20% of all inputaactivation pairs were used to train the regression, and then the held out 80% data was used to visualize and analyze the result, and found that the geometry of the belief state was still well represented (Figure (6B) and that the mean squared error was similar to that of the full regression (Figure|6), red vs. purple bars). Finally,
control that preserved the fractal geometry in the simplex but shuffled the input-point correspondences, resulting in the regression mapping all points to the simplex center due to the lack of discoverable structure (Figure 6C).
These results provide compelling evidence for our central claim: transformers trained on data with hidden generative structure will learn to represent the geometry of belief states in their residual stream. The close match between the predicted fractal structure and the empirical geometry of the residual activations, even for the highly complex Mess3 MSP, suggests that this geometry is fundamental aspect of how transformers build predictive representations of their input data.
3.2 Belief state geometry represents information beyond next-token prediction
Often, distinct belief states will have the same next-token prediction associated with them. Our framework suggests that transformers will keep internal distinctions in the representations these belief states, despite the fact that transformers trained explicitly next-token prediction.
The Random-Random-
(RRXOR) process, shown Figure 7A has these types of degenera- cies (28). The MSP of the RRXOR process has 36 distinct belief states, and can be geometrically represented a 4-simplex. In Figure B we visualize the simplex geometry by projecting down to 2 dimensions. After training, represented in the residual stream. Unlike our previous results for the Mess3 process, we find that the belief state geometry is not represented in the final residual stream before the unembedding, However, when we take the residual stream activations of all layers and concatenate them, we find a veridical representation (Figure (C).
regression analysis to see if the belief state geometry is
The geometry we find not explainable by next-token predictions. quantify this, we ask the pairwise distances between belief state representations in the transformer can be explained by pairwise distances the next-token predictions, or if they are better explained by ground truth belief state distances. First, we compute the Euclidean distance between pairs of ground truth belief states.

---

Randdont-aandon XOR
Geametry
Residall Strcam Reseosonttton
6:33
ELINCE
3:531
Dentor of Mass
6.33
D
ARXOR
RACXR
Mess3
Muan
Mosi
Activation
Figure 7: The belief state geometry can be represented across multiple layers when the belief state structure has next-token degeneracies. (A) The correlation, but has interesting higher-order structure. The MSP of this process has 36 distinct states, with many of them degenerate in terms of their optimal next-token predictions. (B) The belief state geometry lies in a 4-simplex. To visualize we project down 2-dimensions, with each belief state colored uniquely. (C) The transformer linearly represents this belief state geometry. Small dots correspond to individual activations and are colored according to the ground truth belief state. Large dots correspond to the center of mass of all activations associated with particular ground-truth belief state. (D)We compare Fuclidean distances between pairs of ground truth belief states and see if those distances preserved in the belief state representation in the transformer (left scatter plots), showing good correspondence for both RRXOR and Mess3 processes. The right scatter plots compare distances in ground truth next-token probabilities to distances in the transformer's belief' state representations, showing low correlation for RRXOR (R² = 0.31) and high correlation for Mess3, indicating that RRXOR belief state geometry is not captured by next-token predictions alone. (E) Mean squared error of the linear regression procedure, applied to individual layers and, at the far right, applied to the concatenation of activations across layers.
-OR process has zero pairwise
For each pair, we also compute distances between the corresponding belief state representations in the transformer. In (Figure 7 Left) we compare the ground truth pairwise distances to the represented pairwise distances, and find that for both the RRXOR and Mess3 process, there is good correspondence.
Next, for every pairwise distances of belief state representations, we compute the corresponding distance in the ground truth next token probabilities. The scatter plots on the right of Figure 7 compares these distances. For the RRXOR process, the correlation with next- predictions is relatively low (R² 0.31) compared to the correlation with the belief state geometry (R² = 0.95). indicating that belief state geometry captures structure in the residual stream not captured by next- token predictions. For the Mess3 process, the pairwise structure of next-token predictions preserved in the belief state representation. This difference between the RRXOR and Mess3 results is expected from the theory. Since the RRXOR process contains distinct belief states with the same optimal next-token distribution, the discovered belief state geometry is not well-explained by the next token predictions.
3.3 Belief state geometry can be spread across layers of the residual stream
The framework presented here suggests that belief states should be represented in transformers, but does not say that they must be represented in the final layer. In cases of belief states that are degenerate in their next-token predictions, distinctions can be lost before the unembedding. Thus, in the Mess3 process we expect belief state geometry to persist to the final layer of the residual stream, whereas in the RRXOR trained transformer belief state information can be collapsed before the unembedding.
7

---

To quantify this, we report the mean squared error of the belief geometry regression over the different layers of the transformer as well as the concatenation of residual activations from all layers in Figure E We find that the RRXOR belief state geometry is not well represented in the final layer of the transformer. In contrast, the Mess3 belief state geometry is well represented in indidividual layers of the transformer, and persists through to the final layer before the unembedding. Interestingly, the RRXOR belief state geometry is not represented in any of the individual layers of the transformer, but is represented in the concatenation of the layers. This suggests that the belief state geometry is spread across multiple layers of the residual stream.
4 Discussion
4.1 Belief state geometry: Why?
During training, models do not directly see the hidden structure generating the data they only see data. So how is it that the pretrained model infers belief states in the simplex of the generator states? In fact, there are infinitely many distinct HMMs that generate the same stochastic process [5]. Despite this, any stochastic process has canonical vector-space representation in the space of probability densities over all possible future token sequences 3 This corresponds to a minimal HMM representation, and we should expect that transformers learn belief state geometry in the simplex over these minimal generator states. This canonical geometry for process provides an explanation for why we were able to find the belief state geometry represented in our transformers.
In the jargon of computational mechanies, it is notable that the recurrent belief states map onto the causal states of the process' 6-machine [4]. The recurrent states must be distinguished for prediction, even if there exists a generative HMM that is smaller than the e-machine. Typically, data can be generated by a much smaller mechanism than is required to predict it [19, 33). Accordingly, transformers and other pretrained models learn much more than just an accurate generative model
of the wheld-they also learn how to synchronize to the hidden state of the world through observed
context.
We expect our main results do not depend strongly on the particular neural network architecture of transformers, except that they (i) utilize residual stream and (ii) were pretrained future- token prediction. Indeed, we expect to find the same belief state geometry in other neural network architectures like Mamba [1), and recent work building on discovery has shown that recurrent neural networks do represent the belief state geometry for the processes tested in our experiments (26). The generality of our results is thus a powerful architecture- inneeeeere-ddeeedeeent insight for interpretability. Pretraining will generally induce mixed-state geometry. We should then be able to work backwards from this knowledge and the knowledge of the neural network architecture to determine both its world model and how it performs Bayesian updating over its hidden states.
Corollary 2 of (31) implies that, for 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. The implication for feedforward networks like transformer architectures is much less obvious, but it motivates the hypothesis that to perform as well as possible on next-token prediction requires that all information about the past necessary to predict the entire future as well as possible must be present across the layers of each context window position's 's residual stream. Our results here support this interpretation.
Intuitively, why would a model learn about the entire future if it is only trained on next-token prediction? The simple answer is that, 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. If the initial nuance were discarded, then the ability to distinguish next-token probabilities in the future would be compromised. To minimize loss, pretrained models need to not only nail the next-token probability distribution, but also retain all information necessary to get the correct next-token distributions in the distant future. It turns out that this requires distinguishing all pasts that induce distinct probability distributions over the entire future BI].
It is tempting to think that all of this information will persist at the final layer of the residual stream, but this is not strictly necessary nor borne out in practice as we showed in Figure 7 In general, the belief state is spread across the layers at each position. This is consistent with 225. which found that
8

---

the predictive accuracy of the residual stream with respect to far-future tokens peaks somewhere in intermediate layers. Performing as well as possible on multi-token prediction, as in [T]. would imply the same belief-state structure as next-token prediction, but it should change how this information is spread across the layers and, in practice, performance may be different when loss is not minimized.
It is also interesting to consider the linearity of the belief' state representations we found, and if the linearlity is strictly necessary for network optimize next-token cross-entropy loss. In fact, recent work has constructed a transformer by hand (i.e. human being selected the weights) that perfectly solves the RRXOR task, but where the belief state geometry is not linearly represented [T]] is thus a non-trivial empirical finding that transformer architectures trained via stochastic gradient descent find solutions in which the belief state geometry is linearly represented.
4.2 Further implications
Our finding of the belief-state simplex in the residual stream strongly suggests that any model capable of minimal loss requires as many residual-stream dimensions as the number of states in a minimal generative model of the stochastic process sampled by the training data. Moreover, we should be able to use this new understanding to determine the minimal loss with fewer residual-stream dimensions. Likewise, using an adaptation of the rate distortion theory applied to the belief-state simplex (20), we anticipate that our framework should be able to make non-trivial predictions about the evolving geometry of activations during training.
To the extent that different users are represented in the training data, they can be thought of as disconnected ergodic components of non-ergodic stochastic process. Even after the model has learned, needs to synchronize (in context) to both (i) the ergodic component representing the user and (ii) the current latent state of that component. Each will show up in different ways in the MSP Using this framework, we should be able to predict rates of adaptations to different users-i.e. rates of convergence in context-as identified by the reduction in next-token entropy as context position increases.
4.3 What are features?
growing literature in LLM research studies features these systems represent 3 Importantly, there is currently agreed upon definition of feature. Consequently, it has been difficult to study this issue in a principled manner in which ground truth features are known. While the MSP provides ground truth understanding of the computation next-token predictors are asked to accomplish, the relationship between belief states and features is not necessarily one-to-one. A single belief state may encompass multiple features, and the same feature may be represented across multiple belief states. The exact mapping between belief states and features is likely to be complex and dependent on the specific architecture and training data of the transformer model, making it an exciting open problem for future research.
4.4 Limitations and applicability to more realistic settings
Our experimental validation focused on small-scale systems, using HMMs with only 3-5 states and vocabularies of 2-3 tokens. While these systems exhibited meaningful complexity through infinite Markov order in both the RRXOR and Mess3 processes, they represent a significant simplification compared the sophisticated architectures and massive vocabularies (>50,000 tokens) of modern language models trained on natural language data. This scale limitation raises important questions about how our framework extends larger and more complex settings.
For realistic systems like board games or natural language, the dimensionality of the belief' state simplex would exceed the dimension of the residual stream. This suggests that if transformers repre- sent belief state geometry in these domains, they must do so in a compressed form. Understanding the nature of this compression-wve what information is preserved and what is discarded represeens crucial direction for future work.
As the size of the minimal generating structure and the context window length grows, explicitly computing ground-truth belief states becomes intractable. We will need new methodologies for validating our framework in settings where cannot directly compute the MSP structure.

---

We limited ourselves to the simpler stationary ergodic setting in this paper to empirically verify important basic features of transformer representations. However, the framework presented in this work will naturally extend to non-stationary 29 and non-ergodic processes [6] which is important for the extension to real-world tasks. For example, in in-context learning the underlying process is inherently nonergodic. In the more general setting, we also expect transformers to represent belief- state geometry, although that geometry will then reflect the more general types of non- non-stationary and non-ergodic processes representative of real-world data like natural language 7 g
In the experiments we ran, we generated training data using HMMs. A natural question is what to expect with training data generated from non-HMM sources. It is important to note that any dataset made of sequences of tokens used to train transformer can be represented as being generated from an HMM, if one allows for non- ergoodici,,, non non- non-staiionarity. and infinite states. For example, in the case of training data that takes the form parenthesis matched strings, one would usually conceive of this as being generated by a push-down automaton. However, this can equívalentaly be represented by an HMM consisting of an infinite chain of hidden states. Another interesting point that we did not focus on in our presentation for the sake of simplicity, is that even if one does not assume an HMM for the generating machine, the question of the computational structure of an optimal predictor is naturally HMM, with hidden states given by distinct distributions over the future (whether the generator is an HMM or a Turing Machine, or something else) [3].
In this work we verified that belief states are linearly represented in the residual stream, even in cases where the belief states form intricate fractal geometries. Although this is quite suggestive that these representations are used to implement some form of Bayesian updating, we did not directly test this. In other words, we not yet know if these LLMs develop something like a circuit for implementing Bayesian updates, as would be natural for MSP dynamics. So far, we have only established the representation of belief states and their geometry.
Computational mechanics is primarily (but not completely, see Ref. (21)) concerned with optimal prediction, but LLMs in practice are not perfectly optimal. Further work is needed understand how near-optimen oott-optiaaltyy and non- non-optimality influence belief state representations. Progress in this direction could offer insights into evolving structures during training.
Moving beyond stationary ergodic HMM training data, we still generically expect fractal-like structure in the residual-stream activations. Why? The fractal structure is result of Bayesian updating applied to the "non-deterministic" computational structure of HMMs [16]. With stationary processes, this same Bayesian map is folded into itself time and again, producing very clean fractal. For natural language, even in the absence of stationarity, there will nevertheless be resonances of non- deterministic computational structure, which would imply a type of folding beliefs in a self-similar manner. However, a more rigorous version of these statements and the empirical validation is left for future work.
5 Conclusion
In this work, we introduced a theoretical framework that establishes a clear connection between the structure of training data and the geometric properties of activations in trained transformer neural networks. Through experiments, we validated that the geometry of belief states is linearly represented within the residual stream of the transformer architecture. This finding suggests that transformers construct predictive representations that surpass simple next-token prediction. Instead, these representations encode the complex geometry associated with belief updating over hidden states, capturing an inference process over the underlying structure of the data-generating process.
Our work takes significant step towards coneretizing the understanding of the computational structures that drive the behavior of large language models (LLMs) and how these structures relate to the data on which they are trained. This concretization is crucial, as the rapid advancements in LLMs have led to models with increasingly sophisticated behavioral capabilities. However, without a clear understanding of the underlying computational structures and their relationship to the training data, becomes challenging to interpret, trust, and further improve these models.
10

---

References
[1] Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls next-token prediction. arXiv preprint arXiv:2403.06963, 2024.
[2] Emily M Bender, Timnit Gebru, Angelina McMillan-Major and Shmargaret Shmitchell. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM conference on fairness, accountability, and transparency, pages 610-623, 2021.
[3] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- erly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfill-Doddss, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke. Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https//ttrassfommer- http://220/mmooosmaa--n-feftttie
[4] J. P. Crutchfield. Between order and chaos. Nature Physics, 80Jauary):17-24, 2012. doi: 10. 1...0//
[5] P. Crutchfield, C. Ellison, Mahoney, and R. G. James. Synchronization and control in intrinsic and designed computation: An iiformatiion-hhooetii analysis of competing models of stochastic computation. CHAOS, 20(3):037105 2010. Santa Institute Working Paper 10-08-015; 11ti.orrg:007.5533
cond-mat.stat-medd
[6] James P Crutchfield and Sarah Marzen. Signatures of infinity: Nonergodicity and resource scaling in prediction, complexity, and learning. Physical Review E, 91(5):050106, 2015.
[7] Lukasz Debowski. Information theory meets power laws: Stochastic processes and language models. John Wiley & Sons, 2020.
[8] Eukasz Debowski. simplistic model of neural scaling laws: Multiperiodic santa fe processes. arXiv preprint 20X0v22302.00049, 2023.
[9] Joshua Engels, Isaac Liao, Eric J Michaud, Wes Gurnee, and Max Tegmark. Not all language model features are linear. arXiv preprint arXiv:2405. 14860, 2024.
[10] Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. arXiv preprint arXiv:2404.19733 2024.
(11) Rick Goldstein. Handcrafting a network to predict next token probabilities for the random- random-xor process.
June 2024. Research submis- sion to the Computational Mechanics Hackathon! research sprint hosted by Apart, PIBBSS. and Simplex.
[12] Albert Gu and Tri Dao. Mamba: Sinear-time sequence modeling with selective state spaces. arXiv preprint arXi::2312.007552
2023. 2023..
[13] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 2018.
[14] Jiachen Hu, Qinghua Liu, and Chi Jin. On limitation of transformer for learning hmms. arXiv preprint arXiv:2406. 240..0006.0000, 2024.
[15] A. M. Jurgens and J. P. Crutchfield. Shannon entropy rate of hidden Markov processes. Journal of Statistical Physics, 183(2):32, 2021.
[16] Alexandra M Jurgens and James P Crutchfield. Divergent predictive states: The statistical com- plexity dimension of stationary, ergodic hidden markov processes. Chaos: An Interdisciplinary Journal Nonlinear Science, (1)8), 2021.
[17] Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring sequence model trained on synthetic 2022.
task. arXiv preprint

---

[18] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata, arXiv preprint arXiv:2210.10749) 2022.
[19] Wolfgang Löhr. Predictive models and generative complexity. Journal of Systems Science and Complexity, 25:30-45, 45, 2012.
[20] S. E. Marzen and J. P. Crutchfield. Nearly maximally predictive features and their dimensions. Phys. Rev. E, (5 2017. doi: 10.1 1ew.1100sseeee..5..... 95.013011. SFI Working Paper 17-02-007; arxiv.org:1002.008666 702.08866 mond-mmatstat-meehh)
[21] Sarah E Marzen and James P Crutchfield. Nearly maximally predictive features and their dimensions. Physical Review E, 95(5):051301, 2017.
[22] Neel Nanda and Joseph Bloom. Transformerlens.
https://github. com/
2022.
[23] Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models. arXiv preprint arXIv:2309.00941, 2023.
[24] Pedro A Ortega, Jane X Wang, Mark Rowland, Tim Genewein, 7eb Kurth-Nelson, Razvan Pascanu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, et al. Meta-learning of sequential strategies. arXiv preprint 1905 030000 2019.
[25] Koyena Pal, Jiuding Sun, Andrew Yuan, Byron C Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state, arXiv preprint arXiv:2311.08897, 2023.
(26] Keenan Pepper.
RNNs represent belief state geometry in their hidden states. https//aparrtescaarh..oo//// wow//vve/////ommmccc.... June 2024. Research submission to the Com- putational Mechanics Hackathon! research sprint hosted by Apart, PIBBSS, and Simplex.
[27] M. Riechers and J. P. Crutchfield. Spectral simplicity of apparent complexity, Part I: The nondiagonalizable metadynamics of prediction. Chaos, 28:033115, 2018. doi: 10.1063/1. 4985199.
(28) M. Riechers and Crutchfield. Spectral simplicity of apparent complexity, Part HI: Exact complexities and complexity spectra. Chaos, 28:033116 2018. doi: 10.1063/1.4986648.
[29] Paul M Riechers and James P Crutchfield. Fraudulent white noise: Flat power spectra belie arbitrarily complex processes. Physical Review Research, 3(1 3(1:0131770, 2021.
[30] Joshua B Ruebeck, Ryan G James, John R Mahoney, and James Crutchfield. Prediction and generation of binary markov processes: Can finite-state fox catch a markov mouse? Chaos: An Interdisciplinary Journal of Nonlinear Science, 28(1), 2018.
[31] C. R. Shalizi and J. P. Crutchfield. Computational mechanies: Pattern and prediction, structure and simplicity. J. Stat. Phys., 104:817-879, 2001.
[32] Ilya Sutskever. Building AGI, alignment, future models, spies, Microsoft, Taiwan, & enlighten- ment. Podcast Dwarkesh Patel Podcast, March 2023. URL https://www youtube com/ vatch?v=VEUucIaaj.cm Sc Accessed: 2024-05-22.
[33] D. R. Upper. Theory and Algorithms for Hidden Markov Models and Generalized Hidden Markov Models. PhD thesis, University of Califoria, Berkeley, 1997. Published by University Microfilms Intl, Ann Arbor, Michigan.
12

---

A Appendix supplemental material
A.1 Supplemental Figures
Deforn LaperNorn Oae poger)
Albee Layeehoom
Regresntation Etrme
0.00300
0.00255
0,9000
8.0015
Youn
039000
Figure S1: To test the effect of using residual stream activations from before or after the final LayerNorm in our analysis, we compared the belief state geometry representations in both cases. (A) Projection of the final residual stream activations before LayerNorm onto the belief state simplex, as presented in the main paper. (B) The same projection for activations after LayerNorm, showing a qualitatively similar structure. (C) Mean squared error of the linear fit capturing the belief state geometry for both cases. The representation after LayerNorm shows slightly lower error (0.0003) compared to before LayerNorm (0.0004). indicating that the belief state geometry is preserved and marginally better represented after the LayerNorm operation. These results demonstrate that our findings are robust to the choice of using pre- or post-LayerNorm activations.
Relationship Between Loss Representing Belief State Geometry RRXOR
Mess3
0.550
MSE of Projection
0,010
8.008
0.008
6.05
0.004
0.000
0.000
0.079
1.00
1.22
3,04
1.08
1.00
1.10
1.00
1.03
1.30
3.35
1.20
Normalized Validation Loss
Normalized Volidation Loss
Figure S2: To quantify the relationship between transformer performance on next-token prediction and representing the belief state geometry in the residual stream, we plotted normalized validation loss vs. mean squared error (MSE) of projecting residual stream activations onto the belief state geometry for Mess3 (left) and RRXOR (right) processes. Validation loss is normalized to the theoretical optimal, such that 1.0 characterizes transformer with optimal loss. As the validation loss decreases (moving left on the x-axis), the MSE of the regression also decreases, indicating that better performance on the next-token prediction task correlates with a more accurate representation of the belief state geometry in the residual stream.
A.2 Relations to Previous Work
number of connections to previous work in computational mechanics have been mentioned in the Discussion section. Here we will continue the discussion of related works, focusing connections to neural network research.
Transformer Internal Representations Recent work has significantly advanced our understanding of how and what information encoded within transformer models. Ref. (25) introduced the "Future Lens" framework, demonstrating that individual hidden states in transformers trained on natural language contain information about multiple future tokens. This finding consistent with transformers representing belief states, since belief states contain information well beyond the next token. Ref. 9 found that internal representations of days of the week LLMs lie in particular geometric arrangements (in a circular pattern). Since our work draws concrete relation between the structure of data and geometric arrangements inside of transformers, promising avenue of research would be to explain the days of the week result using framework. Other work has studied internal
13

---

representations in the board game Othello, finding representation of the board state in the residual stream [3]. This finding is interpreted by the authors as the transformers containing a "world model". Conceptually extending our work to such setting is quite natural. In sequential game with full knowledge, such as Tic-tac-toe, Othello, and Chess, each board state uniquely determines the future distribution of moves. Thus, the belief states correspond to board states, providing an explanation for why transformers trained on gameplay should contain explicit representations of board states in their residual stream. However, this also makes clear limitation of our approach. The simplex associated with these belief states would, for any realistic board game, be of larger dimension than the residual stream. If transformers they must be doing so in a compressed fashion. This is an important theoretical problem for future work. Importantly, our work provides unique model-agnostic point of view to study interpretability.
representing the belief state geometry in those cases,
Behavioral Limitations of Transformers and Chain of Thought A number of papers have studied the limitations of transformers learn certain algorithmic tasks, such as graph path-finding tasks and HMMs [ [1]. These papers provide important empirical results that are interesting to think through the perspective of the MSP and belief state geometry. In particular, Ref. [14] provides an upper-bound for the depth of a transformer needed to capture HMM generated data up to a particular context- context-window length. This makes sense given that transformers are feedforward machines, and the MSP is, in general, an automata containing recurrent components. Thus, what the transformer should be capturing is really an unrolled version of the MSP, which naturally gives an upper bound for how many layers one would need to capture the MSP to a certain depth. This paper also explicitly performs experiments on belief-state inference. In some of their experiments the explicitly train transformers to map sequences of emissions from an HMM to sequences of the associated belief states. other experiments they train in the normal autoregressive manner solely on the observations, like in our study. Our work provides a nice theoretical connection between the two tasks.
In Ref. M. the authors test and explain a lack of ability for transformers to perform seemingly simple graph planning tasks. In this work, and others like it that study particular algorithmic tasks with particular formal structures like Emite-automate [8]. addition, and the like, we think its interesting to think of how our framework would apply. In general, our approach treats the sequential training data as first class citizen. This makes clear that to relate studies of formal structures (like graph planning) to our framework, one must explicitly consider the particular tokenization used to convert from the formal structure to sequences of tokens. Then the main consideration becomes the computational structure of the sequential nature of the tokens generated by the chosen tokenization scheme (as opposed to the formal structure directly). From our perspective what one learns, and the limitations of being able to perform well on next-token generation given that tokenization, is directly related to the tokenization and not directly to the formal structure that was tokenized. Thus, these results provide important empirical results for our framework to explain, or good tests to find where our framework fails and needs to be amenned/eexannded.
Meta-learning In Ref. 24 the authors provide a theoretical framework for predictors and agents based on a Bayesian framework which builds state machines of sufficient statistics. This is analogous to the mixed-state presentation we have presented here. Indeed, work in computational mechanics has shown that the states of the MSP correspond not only to beliefs over the generator states, but also to the minimal sufficient statisties for predicting the future based on the past. One exciting area of further research this exposes is the potential for computational mechanics to inform our understanding of artificial RL agents.
A.3 Details of data generating processes used in experiments
The transition matrices for the Mess3 process ((A), TTB), and (C))) are defined as follows:
0.765
0.00375
0.00375
0.0675
0.0425
0.00375
((4)_
0.0425
0.0675
0.00375
: (()) =
0.00375
0.765
0.00375
and
0.0425
0.00375
0.0675
0.00375
0.0425
0.0675
0.0675
0.00375
0.0425
T(
0.00375
0.0675
0.0425
0.00375
0.00375
0.765
14

---

The RRXOR process has transition matrices defined as
0.5
0
0
0
0.5
0
0
0
0
0.5
0
0
0
0
0
0
0.5
and T:(1)
0
0
0
0
1
0
0
0
0
where (()) and T are the transition matrices for emissions and 1, respectively.
A.4
Belief State Simplex
Belief states lie within the probability simplex A || defined as:
A 51| = {0 RSS
b = 1 and b; 2 0 for all i}
A.5 Regression
For the regression procedure used in this paper we used standard linear regression, assuming an affine model for the data and then minimizing:
min W.e
1b.. aWa; + 2
where the summation is over all input sequences in the dataset.
A.6 Details of transformer architecture and training
In our experiments, we trained transformer model using the following hyperparameters and training parameters: The model had a context window size of 10, used ReLU as the activation function, I attention head in each
and had head dimension of 8 and model dimension of 64. There of 4 layers. The model had MLPs of dimension 256 and used causal attention masking. Layer normalization was applied. For training, we used the Stochastic Gradient Descent (SGD) optimizer with batch size of 64, running for 1,000,000 epochs, and a learning rate of 0.01 with no weight decay. For each batch we generated 64 sequences from the Mess3 RRXOR HMM, choosing an initial hidden state from the stationary distribution. We instantiated the networks and performed our analysis using the Transformerl library 22
A.7 Details for analysis
For the shuffle and cross validation analyses used in Figure 6 the cross validation procedure 1000 times over, randomly shuffling performing the train test split independently 1000 times.
we performed both the shuffle and
A.8 Reproducibility Instructions
To install the necessary dependencies and the code, follow these steps:
A.8.1
Installation
Navigate to the repository folder and run:
pip install
A.8.2 Reproducing Figures
To reproduce the figures presented in the paper, step through the provided Jupyter notebooks:
eegure6.ipyn
15

---

Figure7 ipyyb
These notebooks utilize the saved models located in the /examples/mon/ee modees directory.
A.8.3 Training the Models
The models trained using legacy code found in the following scripts:
Lagerr//ensllll--rrcmmmmmmmmmmmr py
htgeey//psllaa..tttttmccrrrrrsss _ew.eP py
These instructions should ensure that the environment is correctly set up and that the figures can be reproduced as described the paper.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims have do with the representation of belief state geometry as defined by the mixed-state presentation. Our experiments show that this framework is relevant transformer neural networks.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We include limitations section the discussion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
17

---

Answer: [No]
Justification: Although the paper does contain theory. much of it relies on established results about the structure and limits of prediction. We provide explanations of the theoretical work that allow readers to follow the logic of our experiments and interpret the results. We reference the previous theoretical work that our work is built on.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include details for the transformer architecture and training, well as the HMM structures we used. Additionally we include code that recreates the figures in this work.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
18

---

reproducible, they some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have included code in the submission that reproduces all results. It should be noted that we will continue to work on cleaning up this code for final submission (though the code very much works as is and recreates the figures in the submission).
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
Swwp///ooeSeisss.eoooaa wwwpsubmmmssossnniii ssiiiioiccciiiiiiicc for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We discuss these paramaters in the paper, and also provide code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
We include explanations of error bars and the statistical methods used.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
19

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: These experiments small toy models and thus can be on any modern hardware.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp/////.earreeeccc/ppbbcbheaa
Answer: [Yes]
Justification: I've reviewed the code of ethics and conform to them.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is quite theoretical and the experiments are in toy models in highly controlled settings. Thus we don't think our work has implications for socictal impact the kind mentioned in the guidlines.
Guidelines:
20

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work poses no such risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators of this work including the code on this paper. When external libraries we cite them.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
21

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.comm/aaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We include explanations and documentation for the code we submitted.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
22

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "vZJFZGypTV", "reviewer": "Reviewer_wEQw", "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.", "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?\n\n2. 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.", "limitations": "The authors seem to acknowledge the limitations very well in Section 4.4.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. 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.\n\n3. 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:\n\nBachmann, Gregor, and Vaishnavh Nagarajan. \"The pitfalls of next-token prediction.\" arXiv preprint arXiv:2403.06963 (2024).\nHu, Jiachen, Qinghua Liu, and Chi Jin. \"On Limitation of Transformer for Learning HMMs.\" arXiv preprint arXiv:2406.04089 (2024).\n\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review of our paper. We appreciate your recognition of our work as \"extremely interesting\" and \"well-verified,\" as well as your acknowledgment of the paper being \"really well-written\" with \"very insightful\" discussion. Your positive remarks about the importance of our observations for transformer interpretability are particularly encouraging.\n\nWe're glad you found value in our approach of using mixed-state presentation as a theoretical framework for understanding transformer computations. Your point about the computation being with respect to the input sequence rather than layer-to-layer is indeed a crucial aspect of our work, and we're pleased it resonated with you.\n\nWe have carefully considered your feedback, including both the strengths you highlighted and the areas for improvement. We address your specific points below:\n\n## Weaknesses\n\n**Lack of discussion of related work**\n\nWe strongly agree that this was lacking in our original submission. Please see the global response for how we plan to address this issue.\n\n**Issues related to next-token prediction and Chain of Thought (CoT)**\n\nWe thank the reviewer for bringing up this interesting point, and plan to discuss this issue explicitly in the revised paper. As the reviewer points out, many tasks do not have a structure that, at least intuitively, seems to be naturally captured by next-token prediction dynamics. For instance, tasks may have structures in which many intermediate steps must be considered before arriving at the correct answer, like graph traversing tasks or arithmetic tasks. The structures of these tasks have been both theoretically studied using methods from circuit computational complexity theory (Feng et al., 2024), as well as empirically studied using controlled training data structures (Bachmann et al., 2024). In many of these cases, using seemingly alternative approaches to the traditional token-by-token autoregressive structure, like chain-of-thought, allows transformers to solve these tasks. \n\nFrom the perspective of our work, if CoT examples are prevalent during pre-training, then our framework still applies. That is to say, CoT examples are sequential data just like any other, and have a natural interpretation as step-by-step reasoning that gives rise to particular correlation structures between tokens that will be capturable in a particular generating HMM. We therefore expect our framework to apply; a transformer trained on such data will learn the belief state geometry associated with the mixed-state-presentation of that HMM.\n\nThe difference between CoT and more traditional approaches exists in the difference between the training data distribution, which will be captured by differently structured generating HMMs, and then reflected in their respective MSPs. We expect that a transformer trained on data containing a mix of CoT and non-CoT examples for the same problems, will implement a mixed-state-presentation that contains structures related to each. An exciting research project that we intend to pursue in the future is to study the types of MSPs and belief state geometry structures induced by CoT vs non-CoT training data, and then use this to explain why and how CoT helps transformers perform more complicated tasks.\n\nHere we have focused on simpler training data structure in order to present and experimentally verify the overall framework and its applicability to pretrained transformers. We are very excited to study the geometric structures induced by particular types of data (like chain-of-thought) or training paradigms (like RL methods after pretraining) from the perspective of our framework in future work.\n\n## Questions\n1. **Context length interaction with optimal prediction**: This is an important point to clarify. In our experiments transformers are trained in the normal way, such that a particular context window position only has access to positions within the context position earlier to it (and itself). So for example, in our experiments with a context window of length 10, position 5 in the context window only gets input from 5 tokens. The MSP maps onto this setup quite naturally, as it describes the process of updating beliefs as one observes more and more tokens, starting from no observations. In the MSP, there are transient belief states that arise from the fewer number of observations, which are distinct from the belief states accessible after observation of a full context window.\n\n2. **Why does MSP for RRXOR have finite states?**:  Yes, there is a quick explanation for the finite states in the RRXOR MSP which has an interesting general lesson, having to do with a property of HMMs called unifilarity. An HMM is unifilar if given a current hidden state and a particular emission, there is at most one next-state to go to (i.e. the conditional distribution of the next state given the current state and an emission is deterministic). Interestingly, the recurrent component of the MSP will always be the minimal unifilar HMM that produces its data. The RRXOR HMM is minimal and unifilar, and so we know that the recurrent component of its MSP will have a finite set of states, in this case 5 states. Thus, once enough data is seen such that one is synchronized to the generating HMM, we are assured to be in a recurrent component with 5 states (in this case 5). Showing that there are a finite number of transient belief states is less obvious, but can be worked out for each particular case. In contrast to the RRXOR process, the Mess3 HMM is not unifilar and has infinitely many belief states. Data-generating HMMs that are not unifilar generically have infinite-state MSPs.\n\nBachmann, G., and V, Nagarajan. \"The pitfalls of next-token prediction.\" arXiv  (2024).\n\nFeng, Guhao, et al. \"Towards revealing the mystery behind chain of thought: a theoretical perspective.\" Advances in Neural Information Processing Systems (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nPwbw7qDBq", "reviewer": "Reviewer_exaE", "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.", "questions": "- What aspects of the findings from the paper will carry over to real tasks?\n- How does the recovered geometry change over different points in the residual stream? Are the belief states consistently linearly represented?", "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.\"\n\nEven 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!", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n- 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.\n- 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? \n- 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? \n- It is unclear how the implicature \"the residual stream contains information about the future beyond just the next token\" is made.  \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your thoughtful review of our paper. We are particularly grateful for your recognition of our work as offering a \"fresh perspective\" and presenting a \"well-thought-out toy model.\" Your positive comments on our visualizations are also much appreciated, as we care deeply about the presentation of our work. Additionally, your observation about our work potentially revealing near-optimal rather than perfectly optimal representations has given us food for thought.\n\nWe want to clarify an important point regarding your second bullet in Strengths. Our results don't show that the Mess3 HMM or its parameters are directly represented in the residual stream. In particular, we do not show that the parameters of the generative HMM are linearly decodable. Instead, we demonstrate that the belief state geometry associated with the mixed state presentation of the Mess3 HMM is represented. This is the meta-structure an observer of Mess3 emissions would construct to infer the HMM's state, not the HMM structure itself. Despite this nuance, we agree with your assessment that \"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.\"\n\nWe address your specific points and questions below.\n\n## Weaknesses\n\n**Why Fractals?**\n\nThanks for pointing out this opportunity for clarification. Moving beyond stationary ergodic HMM training data, we still generically expect fractal-like structure in the residual-stream activations. Why? The fractal structure is a result of Bayesian updating applied to the “non-deterministic” computational structure of HMMs (Jurgens and Cruchfield, 2021). With stationary processes, this same Bayesian map is folded into itself time and again, producing a very clean fractal. For natural language, even in the absence of stationarity, there will nevertheless be resonances of non-deterministic computational structure, which would imply a type of folding beliefs in a self-similar manner. The more rigorous version of these statements and the empirical validation will be a subject of our future work. We will include a discussion of this in the text upon revision.\n\n**Does this Extend Beyond HMMs?**\n\nWe agree that further discussion of this issue is needed. Please see the global response for details on how we plan to address this issue.\n\nWe are not exactly sure what a theory-of-mind task is, but are very interested\\! Could you provide some more detail or example? This sounds like an interesting thing to think about from the perspective of our framework.\n\n**Residual Stream Contains More Info Than Next Token**\n\nThis is an important point to make more clear, as it is central to the conceptual advance of this work. Distinct belief states can have the same next-token distribution. However, distinct belief states are always distinct in their predictions for the entire future, beyond just the next token. Why is this? Intuitively, since belief states are distributions over hidden states of the generative process, and the hidden states of the generative process differ in their distributions over the entire future (but not always in their next-token emissions), the belief states inherent this distinction in future distributions.\n\nAnother framing of the MSP states that was not discussed in this paper, but naturally arises in computational mechanics, is that the MSP states are equivalence classes of any-length histories that constrain future sequences in the same way (Upper, 1997). Building up these equivalence classes will give the same MSP structure as the one that arises from the Bayesian updating approach presented in this work. When defined according to these equivalence classes, it is clearly the case that distinct belief states contain information beyond the next token.\n\nAs a concrete demonstration of this, we show that transformers represent distinct belief states of the RRXOR process, even though many of these distinct belief states have identical distributions over next tokens.\n\n**Lack of discussion of related work**\n\nWe strongly agree that this was lacking in our original submission. Please see the global response for how we plan to address this issue.\n\n## Questions\n\n**What will carry over to other tasks?**\n\nPlease see global response for discussion on this issue\n\n**How does recovered geometry change over different points in resid. stream? Is it consistently represented?**\n\nThis is a great question. In our analysis of the transformer trained on the RRXOR task we presented evidence that the belief state geometry is not, in general, consistently represented across layers, and is instead represented in a distributed manner across the layers (Fig. 7E, black points). In contrast, for the transformer trained on Mess3 data, we found that the belief state geometry was already well represented by the first layer, and persisted throughout the layers (Fig. 7E, red points). However, as the question implies, there is more work to do in uncovering exactly which situations give rise to this phenomenon, and a more thorough description of the manner in which the belief states are represented. For instance, one question that can be asked in the case of Mess3 is if the representation of the belief state geometry found in Layer 1 is in the same subspace of the residual stream found in the other layers, or if the belief state geometry is similarly represented across context window positions.\n\nA more detailed investigation is beyond the current scope, though we agree these are natural questions to ask, and potentially connect quite directly with open problems in mechanistic interpretability.\n\n## Citations \n\nJurgens, AM., and JP. Crutchfield. \"Divergent predictive states: The statistical complexity dimension of stationary, ergodic hidden Markov processes.\" Chaos (2021).\n\nUpper, DR. Theory and algorithms for hidden Markov models and generalized hidden Markov models. UC Berkeley, 1997.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QMkgaIrmCG", "reviewer": "Reviewer_UE1W", "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.", "questions": "- In Figure 2, $T^{(1)}$ seems wrong.\n- 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.\n- 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’$?\n- In line 90,  $\\eta_{10}$ and $\\eta_{01}$ are not defined. Maybe you should define $\\eta_{x_0, x_1, \\dots, x_N}$.\n- I think you defined the belief state as a probability distribution $\\eta$. Then, what is the definition of belief state geometry?\n- What is Z1R in Figure 3?\n- 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.\n- In line 122, why did you use the residual stream before the layer norm?\n- In line 168, what do you mean by “degeneracies”?", "limitations": "As they mentioned, there is no extension to LLMs. This work just suggests that the residual streams contain the information of HMMs linearly.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n- 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.\n- More experiments are needed for a larger number of hidden states and vocab sizes.\n- It is only for the HMM model, not language models.\n- 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.\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Weaknesses\n\nWe have revised Sec. 2.2 for clarity (see responses to questions below). For issues dealing with simplicity of our experiments and applicability to LLMs, please see global response.\n\nRegarding belief states affecting next token prediction making the results unsurprising: It is true that every belief state has a next-token distribution associated with it. However, in general, multiple belief states have the same next-token distribution. Since transformers are pretrained on next-token prediction it is therefore perhaps not obvious that a transformer should represent distinctions in those belief states, since one might think that if you only care about next-token distributions, there’s no point in distinguishing belief states that have the same next-token distribution. In fact, it is possible to construct the weights of a transformer by hand that perfectly solves the RRXOR task where belief states are NOT linearly represented (Goldstein, 2024).  Our empirical discovery that this representation emerges over pretraining is thus nontrivial.\n\nR2 values in Fig. 7: We will amend the text to be more precise: The belief state geometry (BSG) explains the residual stream geometry beyond the amount that next-token distributions do. A positive correlation between next-token distributions and activations would be a weak test for BSG representation since representations that are not the BSG can also have high amounts of correlation to the next-token distributions.\n\nBecause of these concerns, we wanted to directly ask if the BSG explains the geometric structure of activations *beyond* the next token distributions. We compare pairwise distances between residual stream activations to both next-token distributions and ground-truth belief states. For RRXOR, while pairwise distances between next-token distributions weakly correlate with residual stream activations (R2 = 0.31), the correlation with ground-truth belief state distances is much stronger (R2 \\= 0.95).\n\nWe'll revise line 164: \"For the RRXOR process, the correlation with next-token predictions is relatively low (R2 = 0.31), indicating that BSG captures structure in the residual stream not captured by next-token predictions.\" We'll also update Fig. 7D caption and provide a more detailed explanation of our approach in the text.\n\n## Questions\n\n1. We will fix this  \n2. We will add further discussion of prediction vs generation and include references. The task of a generator is to emit sequences of tokens by transitioning between hidden states (Löhr, 2012; Ruebeck et al., 2018). In contrast, the task of a predictor is to predict the future while observing sequences of emissions, without direct access to the hidden states of the generator. Predictors don’t have access to the internal state of the generator, since they are conceived of as being observers of generators. Importantly, a small generative mechanism can be hard to predict. This is a major theme in computational mechanics, which we will give more attention to in our revision.  \n3. We will add to line 82 to further clarify the notation and meaning, and update the equation. Upon seeing a belief state η and observing an emission x, the updated belief state η′ is uniquely defined. Thus, the probability associated with that transition is uniquely determined by the previous belief state and the emission, without need to reference the particular belief state one transitions to. We will make this more precise by multiplying the right hand side of the equation by a Kronecker delta function that ensures Equation (1).  \n4. We will add the definition: “We label belief states with η_w where w is the shortest string of emissions which leads to that belief state. (In general, an ordering on tokens is required to make this labeling unique.)”  \n5. We will add this on line 87: “We refer to the geometric arrangement of the points on the simplex as the belief state geometry.”  \n6. We will edit line 49 and Fig. 3 to clarify what the Z1R process is. Please see attached PDF for details.  \n7. To clarify we will edit the text and Fig. 3 See attached PDF for details.  \n8. We used before layer norm (LN) as we wanted to study the \"internal state\" of the transformer. We include new analysis in Fig. R1 in the attached pdf. When we perform the same analyses after the LN our results stay qualitatively the same, and actually become a marginal amount better, in terms of the mean squared error of linearly capturing the belief state geometry.  \n9. Degeneracy refers to the fact that the RRXOR MSP contains multiple belief states that have the same distribution over the next token. For clarity, we will amend line 168 to read: “Since the RRXOR process contains distinct belief states with the same optimal next-token distribution, the discovered belief state geometry is not well-explained by the next token predictions.”\n\n## Limitations\n\nRegarding the statement that our work just suggests that the residual streams contain information of HMMs linearly: It is not just that information about the HMM is contained linearly. We provide a much more specific theory and empirical evidence for exactly what information is contained in the residual stream activations, and with what geometric structure. In particular the information contained in the residual stream is the information of the belief states in the mixed-state-presentation, and the geometry is that of the belief state geometry. It is important to note that the representations we find in the transformers are of beliefs over the data-generating HMM states, rather than information about hidden states or the transition probabilities directly. Figure 3 gives an example distinguishing between these different structures.\n\n## Citations\n\nLöhr, W. \"Predictive models and generative complexity\" (2012)\n\nRuebeck, JB., et al. \"Prediction and generation of binary Markov processes\" (2018).\n\nGoldstein, R. \"Handcrafting a Network ... for the RRXOR Process\" Apart Hackathon (2024)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p1wkkOQizh", "reviewer": "Reviewer_89Bt", "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”.", "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?\n\nHow can one measure similarity to the ground truth more quantitatively?\n\nHow hard is this analysis to perform in higher data-structure dimensions?\n\nHow much does one need to worry about false positives?\n\nHow might the analysis go wrong, or be misleading?\n\nHow does matching of the ground truth belief state correlate with performance on the task?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "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. \n\nAlso 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.“). \n\nWhile 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.\n\nThe 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.\n\nI 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.\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are deeply grateful for your insightful review of our paper. Your recognition of our work as \"well executed\" with \"clear writing and ideas\" is greatly appreciated. We're particularly pleased that you found our examples and figures helpful in explaining complex concepts like mixed state representation and belief state geometry, as we care deeply about the presentation of our work. We are similarly encouraged by your recognition of our controls.\n\nYour observation that our results have \"the flavor of being obvious in hindsight\" while acknowledging the importance of actually proving these intuitions is particularly insightful, and we wholeheartedly agree.\n\nWe have carefully considered your feedback, including both the strengths you highlighted and the areas for improvement. We address your specific points and questions below, and have additionally performed an extra analysis based on your comments, which you can find in the PDF attached to the global response (Fig. R2):\n\n## Weaknesses\n\n**Will findings generalize to more complex tasks**: We will include a larger discussion about this issue in our revision. Please see the global response for how we plan to address this issue.\n\n**Lack of discussion of related work**: We strongly agree that this was lacking in our original submission. Please see the global response for how we plan to address this issue.\n\n## Questions\n\nThese questions are very interesting, and many hit on active areas of research which we plan to address in future work. We don’t have final answers for many of these questions, but for now can give some ideas and interesting perspectives from our framework.\n\n**Multiple Belief Simplices**: The issue of multiple simplices is related to non-ergodic processes, and your intuition for meta-belief states is on the right track. In non-ergodic processes, the data generating HMM consists of multiple disjoint ergodic components. Intuitively, we can think of non-ergodic processes as sets of multiple ergodic HMMs. Data is generated by first choosing one of the ergodic components, and then generating a training example from that HMM. The training data in real LLMs has such non-ergodic structure. For instance, wikipedia articles and personal websites are naturally thought of as being generated by distinct processes. In the case of such non-ergodic processes the MSP describes the joint inference of both determining which ergodic component one is in, and also figuring out which hidden state of the ergodic component one is in. Thus the synchronization process consists of multiple scales of inference. From the perspective of computational mechanics, formal work has shown that this leads to power-law decay in the timescale of synchronization, compared to exponential decay in purely ergodic processes (Crutchfield and Marzen, 2015). Interestingly, this gives a natural explanation for some aspects of in-context learning, such as the empirical finding of power-law decay for loss in context in the largest LLMs (Reid et al., 2024). We are very excited to flesh out this story more formally and connect it to in context learning in future work\\!\n\n**Measuring Similarity to Ground Truth, Analysis in High Dimensions, and False Positives**: On the issue of measuring similarity to ground truth and difficulty of high dimensions, we tried to account for these issues with the shuffle, cross-validation controls and by looking at the activations over training. However in the future we plan to devise more methods of validation, and even unsupervised methods to study the represented belief states without knowing ground truth, like PCA. There has been some initial work towards this end, using genetic algorithms over a parametrized family of HMMs (Bowman and Lu, 2024). We are again excited to study and present this in future work\\!\n\n**Belief state geometry correlation to task performance**: Regarding how matching the belief state correlates with performance on the task: We've added an analysis to quantify this. Please see Fig. R2 in the attached PDF. Also, In Fig. 6D we show that over training the mean-squared error of the belief state representation decreases. Thus, at least in that example, performance on the task directly correlates with the degree to which the belief states are represented. There has also been research in computational mechanics which describes formally how coarse-graining of the belief state geometry relates to prediction (Marzen and Crutchfield, 2017), which is likely related to the fuzziness of the fractal changing over training, and other work which describes the tradeoff between non-optimal predictive features (analogous to belief states in the MSP) and predictive accuracy (Marzen and Crutchfield, 2016, Marzen et al., 2024).\n\n## Citations\n\nMarzen, Sarah E., and James P. Crutchfield. \"Nearly maximally predictive features and their dimensions.\" Physical Review E 95.5 (2017): 051301\\.\n\nMarzen, Sarah E., and James P. Crutchfield. \"Predictive rate-distortion for infinite-order Markov processes.\" Journal of Statistical Physics 163 (2016): 1312-1338.\n\nCrutchfield, James P., and Sarah Marzen. \"Signatures of infinity: Nonergodicity and resource scaling in prediction, complexity, and learning.\" Physical Review E 91.5 (2015): 050106.\n\nMarzen, Sarah E., Paul M. Riechers, and James P. Crutchfield. \"Complexity-calibrated benchmarks for machine learning reveal when prediction algorithms succeed and mislead.\" Scientific Reports 14.1 (2024): 8727\\.\n\nReid, Machel, et al. \"Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.\" *arXiv preprint arXiv:2403.05530* (2024).\n\nBowman, Dylan and Lu, Colin. “Unsupervised Recovery of Hidden Markov Models from Transformers with Evolutionary Algorithms.” Apart Research Hackathon (2024)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers,

We sincerely thank you all for your reviews on our paper; we believe the paper will be much stronger because of them. We were quite happy by the positive recognition of our work, described as “very important for the interpretability of transformers,” a “fresh perspective,” and “a well executed piece of research.”  While the reviewers were overwhelmingly positive, there was also room for improvements that we take seriously. In particular, we appreciate that all reviewers pointed out a lack of discussion of related work and the applicability of our framework to more complex settings, and reviewer UE1W pointed out important areas to improve on clarity. We take these issues to heart and have responded to them. 

We have added two new analyses and a new subpanel in Fig. 3, see attached PDF.

Due to character limits we were not able to respond as fully as we wanted to. Please ask follow up questions if something is not clear or more detail is needed!

We found two main concerns came up across multiple reviewers, and so will address them here:

**Lack of Discussion of Related Work**

We agree with the reviewers that our submission would be greatly strengthened by connecting it to the existing literature. We will discuss and reference much of the work we brought up in the individual responses with some additions:

* Pal et al (2023) demonstrated that hidden states in LLMs have information about multiple future tokens
* Multiple papers have studied internal representations in transformers, both theoretically (Elhage et al., 2022, Park et al, 2023), and empirically ( Li et al., 2022; Nanda, et al., 2023, Engels et al., 2024).
* In Ortega et. al., (2019), the authors provide a theoretical framework for predictors and agents with persistent memory and trained on nonergodic data, based on a Bayesian framework which builds state machines of sufficient statistics. This is analogous to the MSP which we show transformers learn implicitly despite their apparent lack of persistent memory in the traditional sense. States of the MSP correspond not only to beliefs over the generator states, but also to the minimal sufficient statistics for predicting the future based on the past. One exciting area of further research this exposes is the potential for computational mechanics to inform our understanding of artificial RL agents
* We will also discuss the issue in our response to Reviewer wEQw, section *Issues related to next-token prediction and Chain of Thought (CoT)*, which also touches on the issue of more complex data structures
* Refs to papers about in-context-learning, and discussion similar to *Multiple Belief Simplices* section in response to Reviewer 89Bt.
* There are a number of Computational Mechanics papers which we will discuss that are cited in the individual responses to the reviewers, and in the section immediately following.

**Applicability to More Complex Settings and LLMs**

The framework in this work extends to non-stationary (Riechers and Crutchfield, 2021) and non-ergodic processes (Cruchfield and Marzen, 2015), which is important for the extension to real-world tasks. We limited ourselves to the simpler stationary ergodic setting in this paper to empirically verify important basic features of transformer representations. In the more general setting, we also expect transformers to represent belief-state geometry, although that geometry will then reflect the more general types of non-stationary and non-ergodic processes representative of real-world data like natural language.

With regard to non-HMM data, it is important to note that any dataset made of sequences of tokens used to train a transformer can be represented as being generated from an HMM, if one allows for non-ergodicity, non-stationarity, and infinite states. For example, in the case of training data that takes the form of parenthesis matched strings, one would usually conceive of this as being generated by a push-down automaton. However, this can equivalentaly be represented by an HMM consisting of an infinite chain of hidden states.
Another interesting point that we did not focus on in our presentation for the sake of simplicity, is that even if one does not assume an HMM for the generating machine, the question of the computational structure of an optimal predictor is naturally an HMM, with hidden states given by distinct distributions over the future (whether the generator is an HMM or a Turing Machine, or something else) (Upper, 1997).

With respect to LLMs, we will provide discussion of natural language that are relevant to our framework (Debowski, 2012; Debowski, 2020), which motivates reasoning for why we believe that our framework will apply to that setting, as well as some difficulties. We will also include discussion similar to the Why Fractals section of our response to reviewer exaE, and discuss the Pal et al (2023) results.

Again, we appreciate the work the reviewers have put in here, and are excited to share these results with the world!

Sincerely, 

The Authors
## Citations

Crutchfield JP, Marzen S. Signatures of infinity: Nonergodicity and resource scaling in prediction, complexity, and learning. Phys Rev E Stat Nonlin Soft Matter Phys. 2015

Riechers PM, and Crutchfield JP. "Fraudulent white noise: Flat power spectra belie arbitrarily complex processes." Phys Rev Research

Upper, DR. Theory and algorithms for hidden Markov models and generalized hidden Markov models. UC, Berkeley, 1997.

Li, K., et al. "Emergent world representations: Exploring a sequence model trained on a synthetic task." arXiv (2022).

Nanda, N., A. Lee, and M. Wattenberg. "Emergent linear representations in world models of self-supervised sequence models." arXiv (2023).

Pal, K., et al. "Future lens: Anticipating subsequent tokens from a single hidden state." arXiv (2023).

Ortega, PA., et al. "Meta-learning of sequential strategies." arXiv (2019).

Elhage, N., et al. "Toy models of superposition." arXiv (2022).

Meta-review: This paper presents an empirical study showing that "belief states" (essentially, the hidden states of an HMM) are represented linearly in the residual stream of a large language model trained on next token prediction. This study is partially motivated by computational mechanics.

Generally, the reviewers felt that the contribution is solid and the findings are interesting. Some reviewers also expressed enthusiasm for computational mechanics as a potentially useful lens for studying LLMs.

The main reservations were about the simplicity of the setup, the treatment of related work, and the clarity of presentation. On the whole, I'd view the simplicity of the experimental setup as a feature rather than a bug. However, I'm inclined to agree strongly with the later two criticisms. In particular, I think this paper needs to engage substantially with othellogpt and the successor papers. In particular, it's not clear to me what additional information we learn from these experiments beyond what was already known.

Nevertheless, given the overall reviewer enthusiasm, I think it is appropriate to accept the paper and trust the authors to substanstively clarify the relationship and implications for the existing work in this area.
Decision: Accept (poster)