Transformers represent belief state geometry in their residual stream

📅 2024-05-24
🏛️ arXiv.org
📈 Citations: 4
Influential: 0
📄 PDF

career value

189K/year
🤖 AI Summary
This work investigates the computational structure of data-generating processes implicitly learned by large language models (LLMs) during next-token prediction training. Method: The authors formalize this structure as a meta-dynamics of belief updating over hidden states, and systematically characterize its representation within the linear residual stream via optimal prediction theory, linear probing of residual flows, and multi-layer activation space decomposition. Contribution/Results: The study provides the first empirical confirmation that Transformer residual streams explicitly encode global belief states—whose geometric structure can be high-dimensional or even fractal—and that these states contain long-range future token information beyond local prediction. Experiments across non-Markovian data sources verify the linear separability of belief states, revealing both layer-concentrated and distributed cross-layer encoding. This work establishes a theoretical bridge between the structural properties of training data and the geometric organization of internal model activations.

Technology Category

Application Category

📝 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.
Problem

Research questions and friction points this paper is trying to address.

Transformers represent belief states
Belief state geometry in residual streams
Connection between training data and activation geometry
Innovation

Methods, ideas, or system contributions that make the work stand out.

Transformers represent belief states
Linear belief state representation
Framework connects data structure
🔎 Similar Papers
2023-12-17Bulletin of the American Mathematical SocietyCitations: 59