Generalized belief states in transformers

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 mjx-co