Many thanks to Xavier Poncini, Paul Riechers, Fernando Rosas, and Matt Smith for useful conversations.

1. Introduction

Despite the recent successes of transformers and other sequence models, there are still many open questions about their basic principles of operation. For example, what internal representations do these models use for prediction and generation? Can we decode, interpret, and manipulate such representations? Is it even possible to consider a notion of representation without reference to an external world model?

To make progress on these questions, I propose a general way to define, decode and manipulate representations in sequence models. The approach is centered on the concept of "generalized belief states", statistics of past observations that satisfy a few criteria: (1) they are sufficient for predicting future data, (2) they can be recursively updated as new observations come in, and (3) they are low-dimensional, allowing for visualization and interpretation. I show that such generalized belief states can be mapped to the internal activations and dynamics of a sequence model, allowing for decoding and manipulation of the model's representations. I illustrate the ideas using a tractable toy model, a transformer trained on a simple HMM stochastic process, where we will see that the KV cache plays a key role.

This post is directly inspired by the 2024 paper "Transformers Represent Belief State Geometry in their Residual Stream" by Adam Shai, Paul Riechers, and collaborators. That paper introduced the idea of studying belief states in transformers, and proposed a method for doing so using tractable toy models. I recommend reading that paper if you haven't already, though it is not a prerequisite for understanding my post. Section 5 compares this approach with the one developed in that paper.

2. Belief states

We start at a general level, by considering the problem of sequence prediction. Imagine that we observe a history of previous tokens and then wish to predict the next token . We assume that sequences are drawn from a stochastic process . For now, it does not matter whether this process comes from real-world data or a generative model.

Before proceeding, we introduce a bit of notation. indicates the token alphabet, indicates the set of finite-length sequences (including the empty history ), and indicates the set of probability distributions over the alphabet . We write for the set of histories with positive probability under the stochastic process.

2.1 Hidden Markov models

To optimally predict the next token after observing history , we must compute the conditional distribution . We may now ask: what information should we track in order to compute this distribution? How should this information be updated as new tokens arrive? Does this information represent some unobserved latent aspect of the stochastic process?

These questions have a straightforward answer when the stochastic process is described as the output of a hidden Markov model (HMM). Given an HMM description, we can write the probability of any history as ^[1] where is the hidden state of the HMM at step , is the probability of emitting token and transitioning to hidden state from hidden state , and is the initial distribution over hidden states. Using this expression and Bayes' rule, the next-token distribution can be written in a factorized form as Here, is the Bayesian posterior distribution over the hidden state, conditioned on the observed history . This posterior distribution corresponds to the beliefs of a Bayesian observer about the underlying states of the generating process, and in the literature, it is called a belief state. ^[2]

Belief states satisfy two important properties: sufficiency and recurrence. Sufficiency means that the belief state allows us to predict the next token just as well as if we knew the whole history , as seen from the prediction map in Eq. (1a). Recurrence means that, after observing a new token , the updated belief state can be computed from and alone. In particular, using Bayes' rule, we can write the update map as

In addition to sufficiency and recurrence, belief states have several other attractive properties. For instance, the prediction map (1a) is time-independent and linear in . Time-independence lets us interpret belief states at different steps in the same way; linearity lets us interpret predictions as convex mixtures of those made by individual hidden states. The update map (1b) is also time-independent and "nearly" linear, being a ratio of two linear functions of the belief state. (In fact, without loss of generality, we may instead consider "unnormalized" belief states , which update in a purely linear way . ^[3] )

In practice, belief states are often encoded as nonnegative vectors: where is the number of hidden states of the HMM. In this way, each belief state is a nonnegative -dimensional vector that sums to 1, so it belongs to a -dimensional affine subspace of . For HMMs with few hidden states, belief states belong to a low-dimensional space, allowing for easy visualization and interpretation. Also, in this space, belief states often exhibit a remarkable fractal structure, which arises from the recursive nature of the update map. ^[4]

2.2 Example

To make things concrete, we now introduce a running example: a stochastic process called MRRXOR that is generated by an HMM with five hidden states. ^[5] The HMM emits two independent random bits, then their XOR, and then resets. In addition, at each time point, the hidden state can jump to another random hidden state with probability ; if this random jump occurs, it does not emit a token and is immediately followed by an ordinary transition. The HMM is shown in the following figure, with each transition labeled by its emitted bit and probability. (For visual clarity, the random jumps between hidden states are not shown.)

Next, we illustrate the geometry of belief states for this HMM. We compute the belief states after each of the possible histories of length . Each belief state is a normalized nonnegative vector that belongs to a four-dimensional affine hyperplane of .

We visualize these vectors using principal component analysis (PCA), showing all six pairs of the four non-trivial principal components. Each belief state is colored by the last two tokens of the history, with its point area proportional to the history probability. Black dots correspond to each of the five hidden states A–E.

PC1 largely tracks the probability of the "reset" state A, i.e., the belief that an XOR triple has just been completed. PC2 separates histories by their last token. PC3 separates histories ending in 00 or 11 from those ending in 01 or 10, thus coding for the XOR of the last two tokens. In fact, the next-token probability depends only on PC3. PC4 contrasts belief in B,C, where one random bit has been emitted, with D,E, where the XOR bit is pending. The PC2–PC3 projection (top middle) has a clean quadrant structure, with fine-grained fractal structure appearing within each quadrant.

2.3 Next-token prediction versus long-term future prediction

Above, we motivated belief states by considering the problem of predicting the next token after a history . Alternatively, we may consider the problem of predicting a long future continuation that follows that history. What information should an optimal predictor maintain in that case?

In general, sufficiency for next-token prediction does not imply sufficiency for the entire future continuation. For example, for MRRXOR, it can be shown that the third principal component (PC3) of the belief state is sufficient for next-token prediction, but it is not sufficient for the whole future.

However, it can be shown that the belief state -- or, more generally, any state that is recurrent and sufficient for next-token prediction -- is sufficient for the entire future continuation. ^[6] In essence, such states have enough information to update over time, so they can be "unrolled" in order to optimally predict the entire future.

To state the result in a general way, we introduce the following definition.

Definition. A recurrent representation of a stochastic process is a tuple with state space , initial state , prediction map , and update map such that where the state reached after a history is defined by the recurrent update

The state space contains the states of interest; it can be as small as the set of states induced by the histories in , or it can be larger, e.g., including mixtures of these. We now show that a recurrent representation is sufficient for the whole future.

Proposition. Given a recurrent representation of a stochastic process , for any history , the state is sufficient for any future continuation :

Proof

We first show that If every prefix with has , this follows from the chain rule and Eq. (2a). Otherwise, for some , and both sides vanish.

Then, by recurrence (2b), every state in the product is obtained from and the intervening tokens, Hence, each term depends on only through . Since is arbitrary, the same holds for the entire future.

This result is quite general: it applies not only to belief states, but also to "generalized belief states" (discussed next) and "recurrent states" of autoregressive models (Section 4). It also holds for time-dependent update and prediction maps: although and do not explicitly depend on the position , the state can maintain a counter.

3. Beyond belief

As discussed above, belief states formalize representations used by an optimal Bayesian predictor. Despite their attractive features, however, they face two serious limitations.

First, belief states are defined in terms of an HMM -- that is, in terms of some "ground-truth" generative model. However, in many cases, we are only given the stochastic process itself, without access to any ground-truth model. It is possible to infer an HMM for the stochastic process, but generally there is no unique solution, because the same process can be generated by many non-equivalent HMMs. ^[7]

The second limitation is that, even if an HMM description is available, the resulting belief states may be unnecessarily high-dimensional. We may illustrate the problem using a simple example: suppose we first sample the hidden bias of a coin, , then flip a sequence of coins with that bias. After a history of flips, the probability that the next coin lands on heads can be computed as Only two numbers -- the numbers of observed heads and tails -- suffice for predicting the next coin, and these numbers have a very simple update (depending on the next flip, either add one to or to ). However, under the natural HMM representation where the hidden state is the fixed bias , the belief state is a probability density over -- an infinite-dimensional object! In fact, it can be shown that the coin-flipping process has no finite-dimensional HMM representation. This limitation reflects the rigid constraints of the HMM description: the belief state is a nonnegative vector, the prediction map is linear, and the update is a normalized linear map. Such a description cannot accommodate the nonlinear prediction map in Eq. (3).

3.1 Generalized belief states

To overcome these two limitations, we introduce the concept of generalized belief states (GBSs). Like belief states, GBSs provide a recurrent representation of the stochastic process (i.e., they are sufficient for next-token prediction and have a recurrent update). However, GBSs differ from belief states in two ways.

First, GBSs are not defined with reference to a ground-truth generative model, such as an HMM. Instead, GBSs are chosen to be "simple", in the sense that they are low-dimensional. Formally, we seek a recurrent representation with state space for some relatively small .

Second, GBSs do not have to satisfy the rigid constraints of HMM belief states. In particular, they do not need to be normalized nonnegative vectors, and they may have nonlinear prediction and update maps. For example, for the coin process discussed above, a natural GBS is provided by the tuple .

Definition. A -dimensional GBS representation of a stochastic process is a recurrent representation with state space .

Here we have chosen to define "simplicity" in a geometric way, as low dimensionality. There is some subjectivity to this choice, and there are other ways of imposing a simplicity bias. For example, in computational mechanics, simplicity is defined as information-theoretic minimality: one seeks so-called "causal states", recurrent representations that have the smallest amount of memory about the past. ^[8] While causal states are compressed, they do not have an interpretation in terms of geometric embeddings.

It should also be said that the requirement of low dimensionality is not meaningful without constraints on the prediction and update maps. Otherwise, we can always define a trivial one-dimensional GBS representation, e.g., by encoding the entire history into the digits of a single number and decoding it again. To prevent this trivial solution, we typically require that the prediction and update maps belong to some well-behaved class, e.g., maps that are Lipschitz continuous and have bounded image. These continuity conditions also make the geometry interpretable, in the sense that similar GBSs give rise to similar predictions and updates.

In general, a GBS is not a probability distribution, so it cannot be interpreted directly as a Bayesian posterior over latent states of some generating process. Nonetheless, if an HMM description is provided, we can find a map between GBSs and belief states that lead to the same future predictions. ^[9] For instance, in the coin example, the tuple is not itself a probability distribution, but it can always be mapped to a probability distribution in the two-dimensional manifold of belief states . In this sense, we may think of GBSs as low-dimensional parameterizations of belief states.

Several variants of the GBS concept have previously appeared in the literature. Linear GBSs are the states of generalized HMMs (GHMMs), ^[10] while GBSs with nonlinear maps are related to "nonlinear predictive state representations". ^[11] Also, when the stochastic process is a sequence of independent draws from a distribution with an unknown parameter, as in the coin example, a natural GBS is given by the hyperparameters of the "conjugate prior". ^[12]

3.2 Identifying GBSs

The problem of identifying a GBS representation can be formulated as an optimization problem. Specifically, to find a GBS with -dimensional state space , we minimize the following objective over initial states , update maps , and prediction maps : Here is a measure of error, such as KL divergence or mean squared error, and is defined via the recursion . Finally, is some nonnegative weight assigned to each history. In our examples, we weight histories using the actual stochastic process, for any history of length or less, but other options are also possible (e.g., uniform or discounted weights). As discussed, the update and prediction maps should be restricted to some well-behaved class.

If , then there is a -dimensional GBS for the allowed class of update maps, prediction maps, and history weights. Otherwise, the value of provides a bound on the error of an "approximate" -dimensional GBS.

Observe that Eq. (4) is the objective of a recurrent neural network (RNN) with a -dimensional recurrent state. Therefore, in practice, a GBS representation can be found by fitting an RNN. There is also an important case when the problem can be solved in closed form. When the prediction maps are linear, the update maps are linear plus normalization, and there are no nonnegativity constraints, the GBS representation is called a generalized HMM. A minimal GHMM (with the smallest dimension) can be identified by SVD of the Hankel matrix. ^[13]

3.3 Example

We now illustrate GBSs on the MRRXOR example.

First, we compute a minimal GHMM using the SVD algorithm. For the MRRXOR process, it is 5-dimensional. We visualize the resulting GBSs for each history of length 16 by directly plotting all six pairs of coordinates 2–5. We omit the 1st coordinate because it mainly captures the mean prediction across histories and is nearly constant.

This visualization is remarkably similar to the HMM plot in Section 2 -- even though this GHMM was derived without any knowledge of the 5-state HMM description, using only the observed stochastic process . In fact, in this particular example, the GHMM has the same dimensionality as the belief states of the 5-state HMM, so the two geometries are related by an invertible linear map. ^[14]

Next, we look for low-dimensional nonlinear GBSs. For simplicity, we restrict updates to token-specific affine maps (), while allowing nonlinear prediction (an MLP with one hidden layer of 64 tanh units and a sigmoid output). We jointly optimize the initial state, the affine maps, and the prediction network. The following plot shows the best fitted losses for . As expected, the loss decreases with dimension and approaches around .

We visualize the three-dimensional GBSs for all histories of length 16. Each point's area is proportional to the history probability , and its color identifies the final two tokens .

Details of RNN experiment

The initial state is learned. For each token , the update has a learned matrix and offset . The prediction map is , where and is the logistic sigmoid. The models have 198, 271, 348, and 429 parameters for , respectively.

We minimize next-token KL divergence using exact MRRXOR next-token probabilities. Each minibatch contains 256 length-32 sequences sampled from stationary MRRXOR, and the training loss is averaged over all 33 prefixes, including the empty and complete histories. For each dimension, we train four random initializations and one warm start for 100,000 Adam updates each, with constant learning rates 0.001 and 0.003, respectively, and default moments. The warm starts for widen previously fitted 32-unit prediction heads, while the warm start for embeds a three-dimensional fit into four dimensions. All parameters are refined jointly. We select checkpoints by the lowest mean prefix KL on 1,024 fixed independent length-32 sequences.

4. From prediction to mechanism

In the previous section, we focused on the general problem of predicting a stochastic process, without considering how the predictive model is implemented. We now use ideas developed above to study the representations used by a specific implementation, e.g., a trained transformer.

In the following, we consider a given model as an optimal predictor of the same stochastic process that the model generates. Given this process, we then find a mapping between (1) a low-dimensional GBS representation and (2) the model's high-dimensional internal activations and dynamics. This structure, which we call a GBS abstraction, allows us to decode and manipulate the model's beliefs in a low-dimensional space.

4.1 Autoregressive recurrent models

Our analysis applies generally to autoregressive (AR) models. Formally, an AR model is a parameterized statistical model that specifies next-token distributions , one for each history length , where denotes the model parameters. AR models are often implemented using architectures like transformers, RNNs, or state-space models.

Via the chain rule, an AR model specifies a stochastic process over sequences: This may be understood as the "intrinsic" stochastic process of the model, the one that the model predicts optimally (with zero KL error). Alternatively, this stochastic process can be understood as the one generated by the model when it is sampled token by token.

In the following, we always focus on representations used to predict the intrinsic process , rather than some exogenous "ground-truth" distribution . At the same time, we note that if an AR model is successfully trained on sequences sampled from some ground-truth distribution , then and will be close.

Autoregressive training

AR models are typically trained by minimizing next-token cross-entropy. However, for a fixed horizon , next-token cross-entropy equals the KL divergence between sequence distributions, up to a term independent of the model: Thus, minimizing next-token cross-entropy also minimizes sequence-level KL divergence. The intrinsic and training distributions coincide at an exact optimum if the model class can represent the training distribution exactly.

In addition, we assume that the AR model can be expressed in recurrent form. That is, it has an internal state after observing history . This state updates according to some parameterized update map , and predicts the next token according to some parameterized prediction map : Together with the state space , these maps define a recurrent representation of the intrinsic process (as defined in Section 2.3).

In principle, an AR model can always be trivially expressed in recurrent form, e.g., by taking the whole history as the recurrent state, . Instead, we seek nontrivial recurrent representations that refer to the actual variables used by the computation. By manipulating these variables at a given time point, we should be able to change future predictions and updates -- that is, to steer a model in the middle of a sequence, without rerunning it on a different history.

In an RNN or state-space model, the recurrent state is the hidden activation carried between steps. A transformer also has a recurrent representation, but it must be defined in terms of the key-value (KV) cache rather than the residual stream. ^[15] Recall that each new token interacts with previous history through the attention mechanisms, by querying the keys and values stored in the KV cache. Thus, the KV cache along with an embedding of the new token is sufficient both for predicting the next token and for recurrently updating itself with newly generated keys and values. Formally, we define the recurrent state as where indicates the KV cache generated by the first tokens of a nonempty history and is an embedding of the current token. For an empty history, the initial state consists of an empty cache and a beginning-of-sequence (BOS) token. ^[16] From this perspective, the main difference between RNNs and transformers is that the dimensionality of the recurrent state is fixed in RNNs and grows extensively in transformers.

4.2 GBS abstraction

The recurrent state of an AR model is sufficient for prediction, but it may be high-dimensional and difficult to interpret. For example, below we will study a transformer trained on the MRRXOR process. At position 16, its recurrent state is 963-dimensional: 960 for the KV cache and 3 for the one-hot embedding of the token (0, 1, or BOS). In the following, we visualize the recurrent states by projecting them onto the top two principal components (16% of the variance). As we observe, there is no obvious low-dimensional organization.

We now introduce the concept of a GBS abstraction -- a mapping between a low-dimensional GBS and a model's recurrent state that allows us to decode and manipulate the model's representations.

Let indicate the recurrent representation of an AR model, and let indicate its associated stochastic process. Suppose there is also a low-dimensional GBS representation that we wish to use to decode from and encode into the model's states.

Definition. A GBS abstraction is a decoder and an encoder such that , for every , and for every state that can be reached with positive probability by updating from or any .

Simply put, the decoder maps the model's recurrent states onto those of the GBS, while the encoder maps GBS states onto the model's recurrent states. The decoder is a left inverse of the encoder, meaning that the GBS is preserved after an encoding-decoding round trip; we do not require , because the decoder may be many-to-one. Eq. (5) states that the decoder preserves next-token predictions and commutes with updating; its conditions can be summarized via the following commutativity diagram.

This construction is closely related to causal abstraction, which compares high- and low-level models through the effects of corresponding interventions. ^[17] Here the GBS is the high-level state, the decoder maps the model's dynamics onto the GBS dynamics, and the encoder realizes an assignment of the GBS by replacing the model's entire recurrent state.

Moreover, by induction, it can be shown that the above local conditions imply three global properties:

1. For every history , the GBS can be decoded from the model's state: .
2. For every GBS , the model started from the encoding assigns the same probability to every future continuation as the GBS started from .
3. If the model starts from and the GBS from , and both read the same input of positive probability, then the decoding of the model's final state is equal to the GBS's final state.

To summarize, encoding the GBS of history at the current position sets the model's predictions to those it would make after observing . For a transformer, this means overwriting the KV cache and the current-token coordinates, and then letting generation continue.

As with GBSs, we typically seek decoders and encoders that are sufficiently well-behaved, such as affine maps or small networks. Without this restriction, a GBS abstraction may be somewhat trivial. For example, the first layer of a transformer's KV cache typically determines the history, so a decoder could in principle recompute from the cached tokens, and conversely an encoder can write the cache of a history with the given GBS.

In practice, there are different ways to identify a GBS abstraction. For example, one can first fit a GBS of the intrinsic process (using an RNN, as in Section 3.3), then fit the decoder and encoder to it. Alternatively, one can fix only the model and fit the decoder and encoder directly, then use them to define consistent GBS prediction and update maps:

4.3 Example

We illustrate our approach on a transformer trained to predict MRRXOR. After training, we identify a simple decoder and encoder that map this transformer into an approximate 4-dimensional GBS representation. For details of the setup, see the callout box below.

To fit the decoder and encoder , we first find the GHMM approximation of the transformer's intrinsic stochastic process (very similar to the one visualized in Section 3.3). We then initialize the decoder and encoder by regressing this GHMM approximation; afterwards, we optimize their parameters by gradient descent using the objective described below. After fitting, the update and prediction maps of the GBS are defined using Eq. (6).

To define a position-independent decoder , we exploit the fact that attention provides a fixed-dimensional interface to the KV cache. Specifically, the decoder queries attention with a fixed set of optimized queries, combines the result with the current-token embedding, and then passes the result through a single-layer MLP.

For the encoder , we fit an affine map from the GBS to recurrent transformer states for histories of length . Note that the encoder replaces the entire existing cache and current-token vector, writing a cache of 15 positions plus the current-token vector. (Preserving the original cache length would instead require a position-dependent encoder , which could be considered in future work.)

In the following figure, we compare the GHMM states used for initialization (left), the GBSs decoded from the transformer (middle), and those same GBSs after encoding into the recurrent model state and then decoding again (right). Each panel shows coordinates 1 and 2 of the four-dimensional GBS. We see that the overall organization of the GHMM remains recognizable, even after encoding into the transformer and decoding again.

Next, we illustrate how the GBS abstraction allows for steering. We now consider two histories: (1) history ends one bit into an XOR triple, so that its next bit is random, and (2) history ends two bits into a triple with bits 1, 0, so that its next bit must be 1. In the following figure, the background shows the decoded GBSs of all histories of length 16, as in the middle panel of the previous figure.

The left panel shows the GBSs corresponding to these two histories, and . In the middle panel, we take the transformer fed with and intervene by encoding into its recurrent state; we see how this moves its decoded state to . In the right panel, we let the Intervened transformer read the tokens 1, 1, 1. The decoded GBSs (black) follow the path of the Reference transformer fed and then the same tokens (blue rings). Without the intervention, the same tokens take the Un-intervened model along a completely different path (gray).

Details of transformer training and GBS abstraction experiment

Transformer. We fit a transformer to sequences sampled from the MRRXOR process. As in Shai et al. (2024), we use a transformer with four pre-LayerNorm layers of width 64, one attention head of width 8, and ReLU MLPs of width 256. However, we use RoPE rather than learned position embeddings to impose position-independent representations. The transformer is trained on 32-token contexts for 100,000 Adam updates (batches of 512, learning rate 0.001 with warmup and cosine decay). Code can be found at github.com/artemyk/belief-state, including the trained model checkpoint.

Decoder and encoder. We write for the combined decoder and encoder parameters.

For each history , the decoder reads the cache with four learned queries per layer and passes their responses, together with the current token, through an MLP: Here is the output of the layer- attention head for a query at the current position. The parameter vector includes 8,836 decoder parameters: four queries of width 8 per layer and the parameters of the MLP with 131 inputs, 64 tanh hidden units, and four linear outputs.

The encoder is an affine map (4,815 parameters) from the four GBS coordinates to a full 963-dimensional recurrent state: a cache of 15 positions and a current-token vector, whose written coordinates need not be one-hot.

Initialization. A five-dimensional GHMM is fitted to the transformer's conditional future probabilities by SVD and least squares, as in Section 3.3 (using 512 sampled sequences, four-token futures for the SVD, and five-token futures for the operators). Dropping its constant coordinate leaves four, which we whiten over training histories and denote . The decoder is first fitted to these coordinates by regression, and the encoder is then fitted with the decoder frozen, using the GHMM's predictions, the round-trip error, and the consistency of decoded successor states.

Joint refinement. We then jointly optimize the decoder and encoder by minimizing , with the transformer parameters fixed. We consider two modes. In single insertion, the GBS is written once and the transformer then runs on with its growing cache. In repeated compression, the GBS is written, advanced by one token, and read back at every step, so that it evolves through the induced maps and of Eq. (6); these are the dynamics of the final GBS representation. For refinement, write for the continuation following history , and let count continuation steps. The loss is Here, , and the expectation is over histories of length 16, sampled in proportion to their MRRXOR probabilities within the training split, and continuations sampled from the original transformer conditioned on . The tracking term requires that, after a single write, the decoder keeps reading the GBS of the actual history as the transformer runs on with its own cache. Without this term, predictions after a write are preserved but the decoded states drift away from those of the actual history after two tokens. The coordinate term keeps decoded history states close to their initial GHMM coordinates , which fixes the coordinate system. We set to balance the scales of the predictive and squared-error terms. Training uses two stages of 4,000 Adam updates each, the first with the read-back error along the repeated-compression rollout in place of the tracking term, with batches of 64 histories, cosine learning-rate decay by a factor of ten (from for the decoder and for the encoder), gradient clipping at norm 1, and a fixed bank of 65,536 training and 4,096 validation draws; checkpoints are selected by the validation loss.

5. Comparison with "Transformers Represent Belief State Geometry in their Residual Stream"

As mentioned in the introduction, this post was inspired by the 2024 paper "Transformers Represent Belief State Geometry in their Residual Stream". That paper proposed that, when a transformer is trained on data from a ground-truth HMM, the HMM's belief states will be linearly represented in the model's residual stream. To test this, one can fit a linear probe from the residual stream to the belief state and then measure its fit on held-out data.

Although the method developed in the paper seems quite promising, it has two limitations.

First, the probe target is defined as the belief state of an HMM description of the training process. However, as discussed above, we'd like to identify transformer representations without reference to a ground-truth HMM or ground-truth training data. For this reason, we propose to focus on a GBS representation of the model's intrinsic process, which can be computed purely from the model's own predictions (Sections 3 and 4.1).

Second, recall that the recurrent state of a transformer is the KV cache, not the residual stream. For this reason, it is unclear why or when the residual stream might encode a recurrent representation, such as a belief state. ^[18] In fact, in practice, belief states are not always recovered with high fidelity from the residual stream; in my experiments on the MRRXOR process, the held-out for a belief state probe ranges over –, depending on model architecture (e.g., model size and choice of position embedding) and probe design (e.g., held-in/held-out split). Furthermore, because it does not store the recurrent state, the residual stream does not provide a natural target for causal manipulation of belief states. In essence, an intervention on the residual stream leaves the cache untouched, allowing later layers and positions to access earlier history.

To sidestep these issues, we propose to focus on the KV cache, which provides a recurrent representation of the model's stochastic process. At the same time, our analysis suggests a possible explanation for why, in practice, belief states are often recoverable from the residual stream. First, the residual stream generates the new keys and values that are written to the KV cache for later use, so it should carry more information than is needed to predict the next token. Second, since the residual stream queries the KV cache at each attention layer, the cache might "leak" a significant amount of internal information into the residual stream, especially in transformers with a high number of layers, attention heads, and embedding dimensions.

6. Summary

In this post, I proposed a way to identify, decode, and manipulate representations in autoregressive models. The approach was organized around the concepts of generalized belief states (GBSs), i.e., low-dimensional states that are recurrent and sufficient for prediction, and GBS abstractions, which allow us to connect such states to the internal activity of a statistical model. Using a small transformer trained on MRRXOR, I showed that a low-dimensional GBS abstraction can be implemented with simple decoding and encoding maps, allowing the model's beliefs to be set in the middle of a sequence.

This approach suggests the geometry of KV cache states -- possibly after mapping through attention -- as an interesting direction for further study. For example, the structure of the learned decoder and encoder, which we have not examined here, may shed light on how GBSs are encoded across cache layers, heads, and positions.

Finally, recent work has found that transformers have an inductive bias toward low-dimensional recurrent representations ^[19] ; in other words, they appear to favor low-dimensional GBS representations on their own. It would be interesting to understand where this bias comes from, and whether the concepts of GBS abstractions can help explain it.

---

1. We use the edge-emitting convention: tokens are emitted on transitions between hidden states. State-emitting HMMs, as in Rabiner (1989), are the special case in which factorizes into an emission probability for given and a transition probability from to . ↩︎

2. For belief states in POMDPs, see Åström (1965), Smallwood and Sondik (1973), and Kaelbling et al. (1998). In computational mechanics, belief states are often called mixed states. ↩︎

3. The unnormalized belief state is the forward variable in HMM terminology; see Rabiner (1989). ↩︎

4. The belief-state measure goes back to Blackwell (1957), The entropy of functions of finite-state Markov chains. Belief updating is a place-dependent iterated function system: the probability of applying each map depends on the current belief (Jurgens and Crutchfield, 2021, Shannon entropy rate of hidden Markov processes). For fractal dimensions, see Marzen and Crutchfield (2017) and Jurgens and Crutchfield (2021), Divergent predictive states. ↩︎

5. MRRXOR stands for "Mixing Random-Random-XOR". It is based on the RRXOR process of Crutchfield and Feldman (2003), also studied by Shai et al. (2024). Initial hidden states are sampled from the stationary distribution. ↩︎

6. See also Shalizi and Crutchfield (2001), Corollary 2. ↩︎

7. Non-uniqueness persists even among HMMs with the fewest hidden states. For equivalence and identifiability, see Ito et al. (1992). ↩︎

8. See Crutchfield and Young (1989) and Shalizi and Crutchfield (2001). ↩︎

9. If a GBS representation and an HMM belief-state representation describe the same process, and distinct reachable states in each induce distinct distributions over the entire future, their reachable states are in one-to-one correspondence. This correspondence preserves future predictions and commutes with updating whenever the observed token has positive probability. ↩︎

10. Generalized HMMs keep the linear prediction and normalized linear update maps of HMMs but drop the nonnegativity constraints; see Upper (1997) and Shai et al. (2026), who call their states predictive vectors. They include the non-classical (quantum and post-quantum) generative models studied by Riechers, Elliott, and Shai (2025). Observable operator models (Jaeger, 2000) and linear predictive state representations (Littman et al., 2002; Singh et al., 2004) are equivalent linear representations. The dimension of a minimal GHMM is the rank of the matrix of future probabilities given histories, a characterization that goes back to Gilbert (1959) and Heller (1965); see also Ito et al. (1992). ↩︎

11. Rudary and Singh, A nonlinear predictive state representation (2003). ↩︎

12. Conjugate priors with finitely many hyperparameters exist essentially only for exponential families (Diaconis and Ylvisaker, 1979). This reflects the Pitman–Koopman–Darmois theorem: for independent draws from a smooth parametric family with fixed support, a sufficient statistic of bounded dimension exists only for exponential families. ↩︎

13. For spectral learning of linear representations by SVD, see Hsu, Kakade, and Zhang (2009) and Singh et al. (2004). For MRRXOR, we apply an SVD to the matrix of conditional future probabilities , with one column per length-16 history, scaled by , and one row per length-4 future. This matrix has rank five, and each history is represented by the projection of its future vector onto the five leading left singular vectors. Up to column scaling, this is the Hankel matrix of joint probabilities . ↩︎

14. The dimension of a minimal GHMM is at most the number of hidden states of any HMM that generates the process (Littman et al., 2002, Theorem 1). When the rank equals , as for MRRXOR, it is not hard to show that the GHMM state is an invertible linear function of the belief state. ↩︎

15. Oren et al., Transformers are multi-state RNNs (2024). ↩︎

16. A transformer needs an input token to produce an output, so a beginning-of-sequence (BOS) token is prepended to obtain the prediction for the empty history; it plays the role of the initial state . In the model of Section 4.3, to enforce position-independent representations, later positions cannot attend to the BOS entries, so the cache after holds exactly those tokens. The current token is represented by three one-hot coordinates, for 0, 1, and BOS. ↩︎

17. A. Geiger, H. Lu, T. Icard, and C. Potts, Causal abstractions of neural networks, NeurIPS 34, 9574–9586 (2021), arXiv:2106.02997. ↩︎

18. Some stochastic processes, such as the Mess3 process considered in Shai et al. (2024), have the special property that their belief state is an affine transformation of the next-token distribution. Since the residual stream encodes the next-token distribution, we do expect the belief state to be recoverable from the residual stream in these special cases. ↩︎

19. See Riechers, Elliott, and Shai, Neural networks leverage nominally quantum and post-quantum representations (2025) and Shai et al., Transformers learn factored representations (2026). ↩︎