The Geometry of Nonergodic Composition

Crossposted from Belief Updates, the Simplex blog, where several of the figures are interactive. By Kyle J. Ray, Paul M. Riechers, and Adam S. Shai (Simplex, Astera Institute). Telescoping cones recovered by linear regression from a transformer’s residual-stream activations. Each component grows and shrinks in accordance with in-context evidence. Introduction Perhaps the defining feature of LLM pretraining data is its heterogeneity. The training corpus spans not only the collected and varied textual works output by the whole of humanity, but also those generated by machines, data collection devices, and more. Such a large and varied corpus is often appealed to as an explanation for the abilities of modern LLMs [1] . But the statistical structure of data created by a diverse set of generators also implies a particular computational structure for the next token prediction task, and, as we will see, for the geometric arrangement of the internal activations in LLMs. In order to understand the structure of the next token prediction task over data generated from many different sources, and its implications for the geometric structure of activations in neural networks, we will: Start by introducing the concept of nonergodicity, which is an important property of LLM training data. Nonergodicity formalizes the notion of a generator of data made of many sources. Derive the belief geometry that the prediction task over such data implies: per-source belief geometries whose magnitude scales in accordance with how strongly the context supports each component. This gives rise to the telescoping geometric structures shown in this post, and to feature geometry that is eventually sparse and multi-dimensional. Treat this geometry as a falsifiable prediction. If a network represents beliefs linearly, as we have found before[5], the geometry of its activations should take this form. By using data with known mixture structure and known per-component geometry, we are able to predict and then observe a transformer building that geometry. We are excited about these results because nonergodicity is a fundamental statistical aspect of real pretraining data. It is, in some sense, the structure that makes in-context learning both necessary and powerful. [6] argued that when pretraining data is a mixture of latent generative concepts, in-context learning should be thought of as the model implicitly performing Bayesian inference over those concepts. In [7] and [5], we showed that next-token pretraining forces this kind of inference even within a single concept. Nonergodic data calls for both at once. Here, we investigate the specific computational and geometric implications of this type of hierarchical inference. LLM Training Data is Nonergodic Consider a natural language sequence beginning with Do not make…. This ambiguous opening reveals little about the source. These four completions share a prefix, but they do not share a future. Note how each continuation gives information about the generator of that data sample, in this case, a Redditor, an instruction manual, etc. These different sources create token sequences with different correlation structure, which is revealed through additional context. The language of a document, the genre of a story, and the identity of a speaker are all initial choices that jointly constrain the subsequent tokens. In the language of stochastic processes, these situations correspond to nonergodic compositions: mixtures of distinct generative processes, where the identity of the active process is fixed at the start of each sequence generation and never revisited. Much of our previous work has dealt with inference over a world model composed of a single generator; here we extend the discussion to include the meta process of inference about which of multiple generators in the world we should be modeling at all. In the rest of this post, we will explain how the theory predicts a telescoping geometric structure for beliefs over this kind of data, and show some initial results consistent with the fact that transformers represent that geometry in their activations when trained on nonergodic data. Two coins: the simplest example of a nonergodic process In order to understand inference over such nonergodic data sources, we will start with the simple example of data generated from one of two coins. Imagine you know that I have two coins, coin mjx-container[jax="CHTML"] { line-height: 0; } mjx-container [space="1"] { margin-left: .111em; } mjx-container [space="2"] { margin-left: .167em; } mjx-container [space="3"] { margin-left: .222em; } mjx-container [space="4"] { margin-left: .278em; } mjx-container [space="5"] { margin-left: .333em; } mjx-container [rspace="1"] { margin-right: .111em; } mjx-container [rspace="2"] { margin-right: .167em; } mjx-container [rspace="3"] { margin-right: .222em; } mjx-container [rspace="4"] { margin-right: .278em; } mjx-container [rspace="5