Transformers represent belief state geometry in their residual stream

TL;DR

This paper reveals that large language models encode belief states linearly in their residual streams, capturing complex geometries and future information, thus linking training data structure to internal activation geometry.

Contribution

It introduces a framework showing belief states are linearly represented in transformer residuals, even with fractal geometries, and connects data structure to activation geometry.

Findings

Belief states are linearly represented in residual streams.

Transformers encode complex, fractal belief state geometries.

Inferred belief states contain information about the entire future.

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.