Transformers represent belief state geometry in their residual stream (2405.15943v2)

Abstract: What computational structure are we building into LLMs 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.

• The paper reveals that transformers form linear, fractal-like representations of belief states during next-token prediction.
• It validates these structures using linear regression on residual activations from HMM-generated processes like Mess3 and RRXOR.
• The findings imply that belief state geometry can enhance model interpretability and guide improvements in transformer design.

Analysis of "Transformers Represent Belief State Geometry in their Residual Stream"

This paper presents a theoretical and empirical exploration of the geometric representations within the transformer model’s residual streams, particularly focusing on how these models internalize the belief state dynamics from data-generating processes. The authors argue that when transformers are trained on next-token prediction, they tend to develop internal structures that linearly represent belief states, often possessing fractal characteristics, over the course of training.

Core Findings and Methodology

The research begins with an investigation into the geometric encoding of belief states within transformers. Using the theory of optimal prediction, the authors conceptualize that transformers assimilate belief geometries related to the hidden states of data-generating processes. This hypothesis is particularly explored in the context of edge-emitting Hidden Markov Models (HMMs), which produce sequences through transitions between hidden states governed by probability matrices. The Mixed-State Presentation (MSP) formalism is used to encapsulate the belief state dynamics in a probability simplex.

Experiments involved training transformers on sequences generated by known ground-truth HMMs, such as the Mess3 and Random-Random-XOR (RRXOR) processes. The authors utilized linear regression to project the internal activations of transformers' residual streams into a low-dimensional belief simplex, thereby verifying that transformers represent these belief state geometries internally. The results demonstrated that this belief state geometry could be both explicitly represented in a single layer or distributed across multiple layers.

Significant Experimental Results

Noteworthy results include the empirical confirmation of the fractal-like structure of belief states in residual streams. For instance, in the Mess3 process, a 2D projection of a 64-dimensional residual space closely mirrored the theoretically predicted fractal structure of belief states. The authors meticulously controlled for artifacts and confirmed the non-triviality of these structures through various stages of training and cross-validation exercises.

The paper also revealed that transformers maintained distinctions in belief states that were degenerate concerning next-token predictions—particularly evident in the RRXOR process. These belief state geometries were not captured by surface-level predictions but were distributed throughout the deeper layers of the model.

Theoretical Implications

The work advances a compelling argument for a fundamental characteristic of transformer models: their capacity to encode complex belief state geometries beyond mere next-token information. This capability highlights the potential for transformers to engage in sophisticated inference processes over hidden states, reflecting an understanding of the entire future distribution of sequences rather than local token predictions alone.

One of the theoretical implications is that the understanding of how transformers synchronize with hidden states could inform improvements in model interpretability and efficiency. The representation of belief states within the residual streams might necessitate architectures with sufficient dimensional capacity for these geometries, potentially serving as a metric for evaluating model complexity and training efficacy.

Future Directions

The paper sets the stage for further examination into how these belief state geometries influence model behaviors in real-world applications involving more complex and non-ergodic data-generating processes. There is an identified need to explore larger HMMs, potentially expanding the vocabulary spectrum significantly beyond the toy models studied. Further investigation might involve diversified neural network architectures to verify the generality of these findings and to explore how these geometries interact with other facets of model learning, such as feature extraction and multi-token prediction tasks.

Conclusion

This paper contributes to a deeper understanding of the internal workings of transformer models, illustrating how they inherently encode the probabilistic structures of the data processes they are trained on. By linking data generation and computational geometry, the researchers provide a foundation for more advanced insights into model interpretability, potentially impacting how future models are designed and optimized for complex prediction tasks. Such explorations underscore a transformative approach to understanding the latent dynamics and belief models internalized by neural architectures during the training process.