The geometry of nonergodic composition

Implicit Bayesian inference, made explicit

Telescoping cones recovered by linear regression from a transformer’s residual-stream activations. Each component grows and shrinks in accordance with in-context evidence.

Kyle J. Ray1,2, Paul M. Riechers1,2, Adam S. Shai1,2

Simplex, Astera Institute

1 Simplex, 2 Astera Institute

September 9, 2026 · 33 min read7,782 words

Kyle J. Ray, Paul M. Riechers, Adam S. Shai September 9, 2026

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 LLMs1. 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. If a network represents beliefs linearly, as we have found before[5], the geometry of its activations should take this form.
  • Treat this geometry as a falsifiable prediction. 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.
generators Component Mess3AComponent Mess3BorACCBBABBBBBCCCBAAAAACAACbelief in whichcomponentbelief geometryper component10wA = 0.83wB = 0.17×ηAηB=predictionexperimentwA ηAwB ηBactivationstest
The geometric structure of prediction over nonergodic generators. In the setting studied here, each token sequence comes from one of multiple generators (here, Component Mess3A or Mess3B). As tokens are read, an ideal next-token predictor holds two beliefs at once: a belief about which component is generating the tokens (the weights wAw_A, wBw_B) and a belief about what internal latent state that component is in (a point ηA\boldsymbol{\eta}_A or ηB\boldsymbol{\eta}_B in that component’s belief simplex; drawn are all the belief states the component can reach). Their products, wAηAw_A \boldsymbol{\eta}_A and wBηBw_B \boldsymbol{\eta}_B, are the predicted geometry. The belief states over all weights fill a cone, and the slice at the current weight slides along it as evidence accumulates; one grows while the other shrinks. Together they are the composition’s belief, η=(wAηA, wBηB)\boldsymbol{\eta} = (w_A \boldsymbol{\eta}_A,\, w_B \boldsymbol{\eta}_B). Right: A linear probe for the belief geometry, applied to the residual-stream activations of a transformer trained on the nonergodic data, with each activation colored by the belief the ideal predictor would have; the highlighted points are the contexts whose posterior matches the one being read.

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. Xie et al. [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 Riechers et al. [7] and Shai et al. [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.

Do not make…
  • my mistake, buy the bigger one. edit: I know, please stop DMing me. a redditor
  • adjustments while the unit is connected to power. an appliance manual
  • too much of what M. said at dinner. You always do this. a private journal
  • the reader hunt through the appendix. Also, this appears in my 2019 paper. Weak reject. referee #2

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 AA and coin BB, each of different fixed biases. I secretly choose one at random and start flipping it. You see only the outcomes: H, T, H, H, T, H, H,…. Your task is to predict the next heads or tails. To do that, it would be useful if you could figure out if it was coin AA or coin BB that was responsible for the flips you’ve seen so far. At first, you have no idea which coin is being flipped, and all you have is your prior: “it could just as easily be either coin”. But as flips accumulate, the frequency of heads tilts toward one of the two biases, and you will become more confident about which of the coins is active. That process is the process of sequentially updating your posterior to a strong belief about the world: “I’m pretty sure I know which coin it is.”

This is the simplest nonergodic composition; the coins are stand-ins for more generic ergodic components2 that may themselves carry nontrivial latent structure. We will tackle that case momentarily, but here we have two memoryless3 components. The only memory is a hidden “which coin” latent that is set once for each sequence generation and never changed. This choice is hidden because you never see the initial selection directly, but Bayesian inference eventually resolves your uncertainty from observation statistics alone.

Together, the two coins can be thought of as one generator with two hidden states and no way to move between them. Once you know that one of the coins generated the sequence, there is nothing more to know, so your knowledge of the system is fully determined by your belief about which coin is generating the data, which is a point on a segment that slides toward one end as evidence accumulates.

You can try this yourself: flip as many times as you like to gather evidence, set your belief about which coin is responsible for generating the data, and then reveal the Bayes-optimal posterior and the coin.

Generators (coins)
or
Belief in which coin
generated the sequence

← now the coin
← start here
← again?
Guess the coin. Two coins of different bias, drawn as hidden Markov models. One of them, chosen uniformly at random and kept hidden, generates token sequences consisting of heads and tails. Press flip a few times to extend the sequence, and set your belief about which coin is active with the slider (or drag the bars), then reveal the Bayes-optimal posterior and the coin in play. People on the Simplex team (most, not all!) tend to undercommit when the Bayes optimal is very confident.

For these coins, the optimal Bayesian posterior can simply be written down given a sequence of observations; the counts of heads and tails are all it needs (you might remember this from your statistics class). In our case, we have the two coins, AA and BB, with biases pAp_A and pBp_B and a prior μA,μBμ_A, μ_B over which one is active. After nHn_H heads and nTn_T tails, the posterior on coin AA is:

wA(nH,nT)⏟belief it’s coin A  =  μA⏞prior on coin A  pAnH(1−pA)nT⏞likelihood of data given coin AμA pAnH(1−pA)nT⏟coin A alone  +  μB pBnH(1−pB)nT⏟coin B alone\underbrace{\htmlData{sym=w}{w_{\textcolor{#5b7b9c}{\htmlData{sym=A}{A}}}(\htmlData{sym=h}{n_H},\htmlData{sym=t}{n_T})}}_{\htmlClass{ann}{\htmlData{sym=w}{\text{belief it's }\textcolor{#5b7b9c}{\text{coin A}}}}} \;=\;\frac{\overbrace{\textcolor{#5b7b9c}{\htmlData{sym=muA}{\mu_A}}}^{\htmlClass{ann}{\htmlData{sym=muA}{\text{prior on }\textcolor{#5b7b9c}{\text{coin A}}}}}\;\overbrace{\textcolor{#5b7b9c}{\htmlData{sym=likA}{\htmlData{sym=pA}{p_A}^{\htmlData{sym=h}{n_H}}(1-\htmlData{sym=pA}{p_A})^{\htmlData{sym=t}{n_T}}}}}^{\htmlClass{ann}{\htmlData{sym=likA}{\text{likelihood of data given }\textcolor{#5b7b9c}{\text{coin A}}}}}}{\underbrace{\textcolor{#5b7b9c}{\htmlData{sym=aloneA}{\htmlData{sym=muA}{\mu_A}\,\htmlData{sym=pA}{p_A}^{\htmlData{sym=h}{n_H}}(1-\htmlData{sym=pA}{p_A})^{\htmlData{sym=t}{n_T}}}}}_{\htmlClass{ann}{\htmlData{sym=aloneA}{\textcolor{#5b7b9c}{\text{coin A}}\text{ alone}}}}\;\htmlData{sym=plus}{+}\;\underbrace{\textcolor{#a86a6a}{\htmlData{sym=aloneB}{\htmlData{sym=muB}{\mu_B}\,\htmlData{sym=pB}{p_B}^{\htmlData{sym=h}{n_H}}(1-\htmlData{sym=pB}{p_B})^{\htmlData{sym=t}{n_T}}}}}_{\htmlClass{ann}{\htmlData{sym=aloneB}{\textcolor{#a86a6a}{\text{coin B}}\text{ alone}}}}}

The form of this equation shows one of the fundamental lessons of this post. When formally answering the question “what is the probability that coin A generated the sequence?”4 the numerator only depends on information about coin A: coin A’s own likelihood times coin A’s prior. It notably does not depend on information about coin B! The denominator, in contrast, normalizes this numerator by a sum that depends on both coins, and thus couples the belief in coin A with information about both coins. We will see that the form of this belief update, containing a part that has to do with each component independently, and then normalized by a part that has to do with all components, is general.

For the case of the coins, the order of the flips doesn’t matter at all. This is not general, and is atypical of the real world. Most environments that we need to identify have sequential structure. I not you kid. Sorry, rather: I kid you not. Order matters.

In the more general case, the identity of a source lives in the detailed correlation structure of how its tokens follow one another, not just in counts of tokens. A Reddit thread, for example, has many hidden states — what account has replied and what was said are directly observable, but not whether the person behind the account is hungry, or tired. Once components have internal structure, just counting current symbols is no longer enough. We need the general answer to the question the coins raised: what, exactly, must you remember about the past in order to best predict the future? The answer to this is the belief state.

Nonergodic Generators of Data and the Task of Prediction over them

To concretize this into a falsifiable theory, we will need to formalize a general notion of a generator of data composed of many different sources. Each source should have its own internal latent structure, and should generate sequences of tokens. In addition, multiple sources need to be able to be composed in such a way that is consistent with the notion of one source being active, or another source, but not more than one simultaneously.

In the following section, we quickly review the mathematical structure of Hidden Markov Models (HMMs) as latent generators of token sequences, the task of prediction over those sequences, and the corresponding belief geometry associated with that prediction task5. This section is all a review of our earlier work [5], but is necessary to get to Section 4 where we use HMMs as building blocks for nonergodic composition of generators, and study the geometric structure of prediction over those.

HMMs as Latent Generators of Token Sequences

We are trying to capture the situation relevant to the task of prediction over sequences of data, especially when the data is generated by processes that are hidden to the predictor. As in our earlier work, we will use the framework of Hidden Markov Models (HMMs) as our fundamental generator component.

An HMM has a set of hidden states, S\mathcal{S}, and emits tokens xx from a vocabulary X\mathcal{X}. Its dynamics are given by one transition operator per token, T(x)T^{(x)}, whose entry Tij(x)T^{(x)}_{ij} is the probability that the generator moves from hidden state ii to hidden state jj and emits xx as it does so. These operators define both how the hidden states move and also how the state dynamics relate to token emissions. You may remember Mess3, the 3-state HMM shown below on the right, from our earlier work.

M  =  ( X,  S,  η(∅),  (T(x))x∈X)\displaystyle \mathcal{M} \;=\; \bigl(\, \mathcal{X},\; \mathcal{S},\; \boldsymbol{\eta}^{(\varnothing)},\; (T^{(x)})_{x \in \mathcal{X}} \bigr)
X\mathcal{X}tokens, {A, B, C}
S\mathcal{S}hidden states, {S1, S2, S3}
η(∅)\boldsymbol{\eta}^{(\varnothing)}initial distribution, (⅓, ⅓, ⅓)
Tij(x)T^{(x)}_{ij}probability of going from state ii to jj and emitting token xx
A:42%B:42%C:42%B:9%A:9%C:9%B:9%A:9%C:9%S1S2S3
samples
CCCBCACBBABA…
CCACCCCAABAA…
An HMM as a generator of token sequences, with the Mess3 generator as an example. Right: Mess3’s three hidden states, with each arrow labelled by its most likely token (not all arrows are shown for concision). Labels are token : probability. The full transition operators are in Section A.2.

Natural language is of course more complicated than these examples, but notably any stochastic process6 can be generated by some HMM.

The Task of Prediction and Belief State Geometry

Despite the name (GPT stands for generative pretrained transformer), transformers are actually (pre)trained to predict, not generate. A predictor observing sequences of tokens and trying to predict the next token cannot directly see the hidden state of the generator. What it can do is keep a belief η\boldsymbol{\eta}, a probability distribution over the hidden states, and update it with each token. As discussed in our previous work, an optimal predictor will update its belief, upon seeing a token xx, from η\boldsymbol{\eta} to η′\boldsymbol{\eta}', according to Bayes’ Rule.

η′  =  η T(x)η T(x)1\displaystyle \boldsymbol{\eta}' \;=\; \frac{\boldsymbol{\eta}\, T^{(x)}}{\boldsymbol{\eta}\, T^{(x)} \mathbf{1}}
η\boldsymbol{\eta}current belief (prior)
T(x)T^{(x)}operator for the token seen (likelihood)
η T(x)1\boldsymbol{\eta}\, T^{(x)} \mathbf{1}normalization ZZ, the token’s probability
η′\boldsymbol{\eta}'updated belief (posterior)
P(S1)
P(S2)
P(S3)
η∅
ηB
ηBC
ηBCC
Belief updating over Mess3. Starting from η(∅)=(13,13,13)\boldsymbol{\eta}^{(\varnothing)} = (\tfrac13,\tfrac13,\tfrac13), every context induces a belief in the simplex; the set of reachable beliefs form a fractal. Points are colored by the belief itself (red, green, blue for the probability on S1S_1, S2S_2, S3S_3). The marked path is an example belief trajectory induced by the token sequence BCC.

These belief states η\boldsymbol{\eta} are vectors that live in a probability simplex. The set of belief states that are reachable from the sequences a generator creates thus has a geometry, the belief state geometry. For instance, in the case of Mess3, there are an infinite number of distinct belief states, that arrange themselves in the probability simplex as a fractal.

Importantly, the information a belief state contains is everything the past tells you about the future; it is the general answer to the question the coins raised. For a coin the belief over its single state is trivially the number one, which is why counting heads and tails was all there was to do. The posterior over which coin was a belief of a different kind, a belief about which generator is active. As we will now see, in general a predictor has to carry both types of information.

Nonergodicity, Prediction, and Telescoping Geometry!

We now have all the pieces needed to create a generator composed of multiple sources/components. The high level approach will be to design a single HMM whose hidden states are the hidden states of all the components put together, and whose dynamics never move between components. The coin game from earlier is a simple example of this: pick a coin, then generate a sequence using only that coin. After we have an HMM that generates nonergodic data, we will figure out the geometric consequences for prediction.

Nonergodic Composition

The mathematical move to create generators of nonergodic data, called nonergodic composition, will be to compose component HMMs via the direct sum. We find that it is often helpful to see both the general theory and an example to keep intuition grounded, so we encourage you to flip between the tabs below.

General theory

Given NN component HMMs M1,…,MN\mathcal{M}_1, \ldots, \mathcal{M}_N, the nonergodic composition is a single HMM M\mathcal{M} whose token-labeled transition matrices are the direct sum of the components’ matrices T1(x),…,TN(x)T^{(x)}_1, \ldots, T^{(x)}_N:

T(x)  =  ⨁n=1NTn(x)  =  (T1(x)0⋱0TN(x)) .T^{(x)} \;=\; \bigoplus_{n=1}^{N} T_n^{(x)} \;=\; \begin{pmatrix} T_1^{(x)} & & 0 \\ & \ddots & \\ 0 & & T_N^{(x)} \end{pmatrix} ~.

The block-diagonal structure is the key property: since the off-diagonal blocks are zero, a state in block nn can never transition to block m≠nm \neq n. The process is permanently confined to whichever block it starts in 7.

We also need to compose the initial states, ηn(∅)\boldsymbol{\eta}_n^{(\varnothing)}. Each sums to 1 within its own component, but not across components. So to compose them we need to choose a weight for each component, μn\mu_n with ∑nμn=1\sum_n \mu_n = 1, which is the prior probability that component nn is the one generating the data. The initial state vector of the composition is then the concatenation of the components’ initial vectors, each scaled by the weight on its component,

η(∅)  =  (μ1 η1(∅),  μ2 η2(∅),  …,  μN ηN(∅)) .\boldsymbol{\eta}^{(\varnothing)} \;=\; \bigl( \mu_1 \, \boldsymbol{\eta}_1^{(\varnothing)},\; \mu_2 \, \boldsymbol{\eta}_2^{(\varnothing)},\; \ldots,\; \mu_N \, \boldsymbol{\eta}_N^{(\varnothing)} \bigr) ~.

Mess3 example

Let us consider the nonergodic composition of two Mess3 generators, each acting as a distinct source of token sequences. We will call them Component Mess3A\mathrm{Mess3}_A and Component Mess3B\mathrm{Mess3}_B. Each Mess3 will have different hyperparameter settings, as shown below. To make a single generator out of these components, in which every sequence is generated either by Mess3A\mathrm{Mess3}_A or by Mess3B\mathrm{Mess3}_B with 50/50 probability, we arrange the transition matrices of the two components in block-diagonal form8.

Component Mess3AA:42%B:42%C:42%B:9%A:9%C:9%B:9%A:9%C:9%S1S2S3
μA = ½
or
Component Mess3BB:33%A:33%C:33%B:33%A:33%C:33%S1S2S3
μB = ½
⇨compose
T(A) =
.42.03.03000.09.14.03000.09.03.140000000.08.08000.330.08000.33.080
AB
T(B) =
.14.09.03000.03.42.03000.03.09.140000000.33.08000.080.08000.08.330
AB
T(C) =
.14.03.09000.03.14.09000.03.03.420000000.08.33000.080.33000.08.080
AB
Two Mess3 components nonergodically composed into one HMM. Each component has distinct dynamics. Component Mess3A mostly stays in its current state; Component Mess3B mostly switches (note that not all arrows are shown here for concision). Right: the composition’s token labeled transition operators: in total six hidden states, with one 3×33 \times 3 block per component (Mess3A’s block shaded slate, Mess3B’s rose, rows and columns in the order S1,S2,S3S_1, S_2, S_3 within each block), and zeros everywhere else, so probability mass never moves between components. The prior μA=μB=12\mu_A = \mu_B = \tfrac12 determines which component the composition starts in.

This composite HMM is another HMM, a latent generator of sequences of tokens. Note that because the transition probabilities associated with one component always have zero probability of transition to any state in the other component (the off-diagonal terms are all zero by construction), it is impossible for the generator to move between components, once it has started in one.

Belief Geometry over Nonergodic Data

Next, we apply the belief update rule to such a composition of ergodic components. We will see that while the generator is permanently confined to whichever block it starts in, an observer’s guess about which component is active is not [8]. Like guessing the hidden coin from a sequence: the true coin is always the same, but as flips accumulate you change your belief about which one it is. In the belief geometry, this ends up coupling geometric structures associated with each component in a particular way.

The belief updating rule is the same as for a single component HMM,

η′  =  η T(x)η T(x)1 ,\boldsymbol{\eta}' \;=\; \frac{\boldsymbol{\eta}\, T^{(x)}}{\boldsymbol{\eta}\, T^{(x)} \mathbf{1}} ~,

but now both the initial state and the transition operators have block structure:

η(∅)  =  (μ1 η1(∅),  μ2 η2(∅),  …,  μN ηN(∅)) ,T(x)  =  ⨁n=1NTn(x) .\boldsymbol{\eta}^{(\varnothing)} \;=\; \bigl( \mu_1 \, \boldsymbol{\eta}_1^{(\varnothing)},\; \mu_2 \, \boldsymbol{\eta}_2^{(\varnothing)},\; \ldots,\; \mu_N \, \boldsymbol{\eta}_N^{(\varnothing)} \bigr) ~, \qquad T^{(x)} \;=\; \bigoplus_{n=1}^{N} T_n^{(x)} ~.

Let’s take a look at the belief state after a single token emission. Because the off-diagonal blocks of T(x)T^{(x)} are zero, the numerator of the belief update acts block by block:

η(∅) T(x)  =  (μ1 η1(∅)T1(x),  …,  μN ηN(∅)TN(x)) .\boldsymbol{\eta}^{(\varnothing)}\, T^{(x)} \;=\; \bigl( \mu_1\, \boldsymbol{\eta}^{(\varnothing)}_1 T_1^{(x)},\; \ldots,\; \mu_N\, \boldsymbol{\eta}^{(\varnothing)}_N T_N^{(x)} \bigr) ~.

Each component’s initial belief gets multiplied by its own operator, as if it were the only generator. The denominator of the belief update is a normalization, which sums over all entries of the numerator, and thus couples the belief updating across the components by a scalar. A small bit of algebra9 shows that we are again left with a concatenation of per component beliefs each scaled by a scalar wnw_n with ∑nwn=1\sum_n w_n = 1. The resulting belief state, and indeed all reachable belief states (due to the recursive nature of belief updating), can be expressed this way. We can always decompose a belief as

η  =  (w1 η1,  w2 η2,  …,  wN ηN) .\boldsymbol{\eta} \;=\; \bigl( w_1 \, \boldsymbol{\eta}_1,\; w_2 \, \boldsymbol{\eta}_2,\; \ldots,\; w_N \, \boldsymbol{\eta}_N \bigr) ~.

Because of this, our interpretation of the initial state carries over to all belief states, with the mixture prior μn\mu_n becoming a per-component mixture posterior wnw_n. In short: the belief is a distribution over all components’ hidden states that can always be expressed in terms of the probability wnw_n that the predictor puts on component nn, and the belief ηn\boldsymbol{\eta}_n over the states of component nn, conditioned on being in that component.

Component Mess3A
S1
S2
S3
Component Mess3B
S1
S2
S3

wA = 0.50 wB = 0.50

▸ belief states▸ full belief geometry▸ per-component insets
The possible belief states of a nonergodic composition of any two 3-state HMM components, which live in a 5-simplex. Shown are two 3D projections from the 5-D space, corresponding to a single component’s simplex (dashed, unscaled) scaled toward the origin by the weight on that component; the weights are coupled, wA+wB=1w_A + w_B = 1, so the slider moves both. Drag either panel to rotate both. The first toggle draws each component’s belief state geometry for our example Mess3s inside its scaled simplex, which is where the composition’s belief can sit at that posterior. The second toggle shows the full belief geometry of the composition: those belief states at every weight from 0 to 1. The third toggle adds an inset per panel with the component’s belief states on its own unscaled simplex.

From this we can see something important about the belief geometry. The beliefs of a nonergodic composition live in a simplex whose dimension is set by the total number of hidden states across all the components. For our two 3-state HMMs, that is 6 states, so the 5-simplex. From that 5 dimensional space, we can project the belief onto the coordinates of any single component, giving wnηnw_n \boldsymbol{\eta}_n. This is a point in a simplex, but shrunk toward the origin by the weight wnw_n. The wnw_n are not independent from each other: they sum to one. So, as the belief puts more weight on one component, its simplex grows in magnitude, and the others shrink towards the origin. Thus, the projection gives the belief geometry a telescoping effect. Above, we show where the belief vectors can live when looking at this projection for two arbitrary 3-state HMMs, for any given wAw_A. The toggles show and hide the specific fractals for a nonergodic composition of two Mess3s.

The result is that components that explain the observed data well accrue weight; components that don’t, lose it. Belief updating over such a composition has a characteristic signature: eventually sparse multi-dimensional features. Early in context, several components carry non-negligible weight wnw_n; as we see more tokens and evidence accumulates we expect wn→δn,n∗w_n\to \delta_{n,n^*} for the true component n∗n^*, and the geometry to collapse onto the active block only.

Does this geometry show up in trained models?

The framework above predicts a specific geometric structure for the belief geometry associated with prediction on nonergodic token sequences. When a transformer is trained on next-token prediction over such data, can we find that geometry in its activations?

Here we show our initial positive results. To test this in a transformer we use a nonergodic composition of two Mess3 generators. Mess3’s belief states form a fractal that fills the simplex, so the nonergodic composition of two Mess3 components should give two fractal-filled cones, each telescopically scaling with the weight on its component.

The interactive figure below shows the ground truth belief geometry, which serves as a nontrivial falsifiable prediction for what we should find in the transformer activations. The full beliefs live in 5 dimensions, and what is shown below are two 3D projections from the 5-simplex to the belief entries associated with each component. Drag either panel to rotate both, and use the ℓ slider to move through the context.

Component Mess3A
S1
S2
S3
Component Mess3B
S1
S2
S3
POSTERIOR OVER COMPONENTS, wnwA = 0.50wB = 0.5000.51081624ℓwAwB
x1:ℓ =BBCAAAACACAAACBAAACCCCCCℓ 0
▸ parameters x, α▸ full belief geometry▸ per-component belief insets
The belief geometry of the nonergodic composition of two Mess3 components. Shown are two 3D projections from the 5-simplex. Each panel is one component’s simplex scaled by the weight on that component, carrying that component’s reachable belief states; the filled marker is where the composition’s belief sits, projected onto that component, and the bottom graph shows the weights through the sequence. The sequence is drawn from Component Mess3A\mathrm{Mess3}_A. Click a token to cycle through realizations at that context position. The first toggle exposes each component’s xx and α\alpha, which reshape its fractal and, for the same tokens, the posterior; the second shows the full belief geometry: the component’s belief states (to depth 5) at every weight from 0 to 1, which is the cone the composition’s beliefs live in; the third adds an inset per panel with that component’s own belief on its unscaled simplex.

We trained transformers on a nonergodic composition of two Mess3 generators. Our theory predicts that the activations should track the belief states. Because we have ground-truth access to the generator, we know the exact belief vector associated with each context position. A linear map fit from the residual stream to these ground-truth belief states recovers them on held-out contexts with R² ≈ 0.985 (compare to an untrained network, which is at ≈ 0.45), and the predicted geometry appears in the residual stream over training:

The telescoping geometry emerging in early training, with the training loss shown beneath the cones. These are two 3-d projections of a 5-d geometry so each point appears in both cones. The points that are the face of one cone, appear as the low variance tip of the other cone.

We see our two telescoping cones, one for each component, scaling with the posterior weight on that component. This emerges as a direct consequence of pretraining on next-token cross-entropy alone. Nothing in the training objective tells the model directly about components, belief vectors, or simplices. The color here encodes the entropy that a Bayesian observer would have over which of the two components is active, given the context that led to that activation. The middle yellow region corresponds to contexts that are well explained by either component, and the states of maximum certainty are the darker tips and faces of the cones.

Below, you can explore the geometry of the converged model’s activations. The left panel is the cumulative variance explained by PCA of the final layer activations, drawn separately for contexts generated by each component (M1\mathcal{M}_1 and M2\mathcal{M}_2); the two right panels show those same activations passed through the learned linear map to the predicted geometry. You can filter activations by ground-truth posterior entropy or by context position. Try dragging the maximum posterior entropy slider down toward 0. This means keeping only contexts where the evidence supports committing largely to one component. We see that the CEV curves climb faster, meaning the activations effectively fill fewer dimensions as the model hones in on a single component. Meanwhile, in the scatter plot, the cone for the now-unlikely component collapses toward the origin. Filtering to late sequence positions tells a similar, but noisier story. Additional context is increasingly likely to support just one component or the other, but it is also possible to observe long sequences that have similar likelihood under either component, or are even flat out misleading (perhaps you ran into some of these in the coin game?).

Interactive — Activation geometry explorer. Filter contexts by generating component, posterior entropy and sequence position; the cumulative-variance curves and the two cones update together. The variance concentrate into fewer dimensions when only confident, single-component contexts are included. The cumulative variance shown is always calculated using the full residual stream, rather than just the subspaces associated with the belief geometry.

Did it have to be this way?

The way we derived the nonergodic belief geometry, it may seem almost as if there was no alternative for what the neural network should represent 10. In light of this, it is worth explicitly pointing out that the predicted geometry is not something that just has to be present for the model to output correct next-token probabilities. This representation manifestly carries more distinctions between contexts than are implied by their differences in next-token prediction. While the beliefs live in 5 dimensions (a distribution over 6 hidden states), the next-token distribution lives only in 2 dimensions (a distribution over 3 possible tokens). Of course, this geometry does perfectly contain the distribution over the next token– but it also represents distinctions in the token after, the 10th token, and the joint probability distribution of the 3rd and 11th tokens conditioned on the 9th. It fully contains all distinctions that can be made between distributions over the future, yet it emerged only by looking one token ahead.

See the plot below, which shows the first three principal components of the 5-dimensional predictive geometry (right), colored by the associated next-token distribution. A small region in the next-token simplex (left) can correspond to significantly different parts of the full-future predictive geometry.

You can scroll around the next-token simplex yourself. You will see that some regions in the next token simplex correspond to unambiguous distributions over the future and other next-token distributions permit many different distributions over the full future.

Interactive — The belief geometry is richer than next-token predictions. Left: the contexts placed at their next-token distribution in the 3-token simplex. Right: those same contexts in the 5-dimensional belief geometry, shown as its top three principal components and colored by their next-token distribution. It opens on one highlighted group: they all predict essentially the same distribution over next tokens, yet lie far apart in the belief geometry. Click through and hover over the simplex to pick any other group.

We note that these canonical low-dimensional representations emerge most cleanly when we initialize network weights to be small, perhaps placing the network in the “rich” feature learning training regime studied in deep learning theory as opposed to the “lazy” one [9, 10].

Parting thoughts

The next token prediction task over nonergodic data requires two levels of inference: figuring out which generator is currently active, while also tracking what state that generator is in. One geometric implication for the activations of neural networks is a per-component projective embedding of the belief geometry with each component’s scale being the posterior weight accorded that component. With this geometry as a falsifiable prediction, we trained transformers on nonergodic compositions, and found this geometry linearly embedded in the residual stream.

Real data is made of many more sources than two, and they will overlap in their structure to different, and quite complicated, degrees. Some components will share most of their structure and differ in a few probabilities; others will share almost nothing; many will sit somewhere in between11. Taken together, we should expect a rich, hierarchical inference process to emerge from that: weights over components, weights over groups of components that look alike, and within each, the component’s own belief updating. This is one way to see why pretraining on such data produces in-context learning [7].

In closing, let’s revisit the humble coin. Some data is closer to a bag holding infinitely many coins: you draw a bias pp from the continuum and start flipping (the problem Laplace solved in 1774); the sum over coins in the Bayesian updating equation becomes an integral, the finite set of weights (one for each component) become a continuum of weights, and the telescoping picture would need infinitely many cones. Yet, the formula for the weights would still only ever consults two numbers, the counts nHn_H and nTn_T, so the distinct beliefs an observer can hold about the future still form a finite-dimensional predictive geometry, described by two parameters: an estimate (the fraction of heads) and how certain it is (the total number of flips). Whether a model stores such beliefs or computes them from running tallies, and what that means for a continuum of memoryfull components with nontrivial internal structure, and for generalization, is the subject of a post to come.

Our story supports a refinement to the picture of transformer representations as sums of sparse one-dimensional features that motivates sparse autoencoders [12, 13, 14]. For data with nonergodic structure, the right ansatz seems instead to be sparse dense subspaces — multi-dimensional geometries that correspond to inference-time Bayesian updating over an underlying world model that includes mutually exclusive12 parts. The model uses many dimensions while a given component is in play, but eventually only a few components carry weight at any given time. Sparsity at the component level, density within each component. The same machinery extends naturally to compositions with internal factorization (each component itself a product of more elementary parts), and it predicts that models trained on factorizable data should discover those parts, represent them in correspondingly factored subspaces [15], and also simultaneously keep track of the meta dynamic over the components.

This picture, of transformer representations as a sparse sum of points within multidimensional subspaces of activation space, is consistent with recent work extending the “linear representation hypothesis” [12] to accommodate observations of multidimensional features in language models [16, 17]. These works suggest that neural network activations be modeled as sums of multidimensional features, whose value is represented as a point in subspaces of dimension greater than 1, but where most such features don’t have a defined value (or have value ~0) on most activations (they are sparse). We find this picture emerges naturally from theory as a consequence of performing prediction over a process consisting of nonergodic components.

Appendix

Acknowledgments

This post draws on joint work at Simplex on the geometry of belief states in nonergodic sequence tasks. Particularly, we thank Javan Tahir, Casper Christensen, Loren Amdahl-Culleton, and Andrew Jun Lee for helpful discussions; Eric Michaud, Jasmina Urdshals, and Selma Maizioud for helpful comments on this blog; and Eric Michaud for input on our discussion of sparse autoencoders and the multidimensional linear representation hypothesis.

A version of this problem has served as a take-home question for Simplex job and MATS applications; the theory and results presented here were developed beforehand and are independent of any applicant work. We thank the applicants for the care and creativity they brought to the problem.

We used LLMs (Opus 4.5+, Opus 5.0, and Fable) to design and run experiments, design this blog, and draft this post. Most prose in this version was written by the authors. We take all responsibility for the content.

The Mess3 process

The Mess3 process [5, 18] has three hidden states S={1,2,3}\boldsymbol{\mathcal{S}} = \{ 1, 2, 3\}, and three observable tokens X={a,b,c}\mathcal{X} = \{ a, b, c \}.

The process is defined by two parameters, α\alpha and xx, with dependent quantities β=(1−α)/2\beta = (1-\alpha)/2 and y=1−2xy = 1-2x. The two components used throughout this post are drawn from this family: the first uses x=0.15,α=0.6x=0.15, \alpha=0.6 and the second uses x=0.5,α=0.66x=0.5, \alpha=0.66.

The labeled transition matrices are:

T(a)=[αyβxβxαxβyβxαxβxβy]T(b)=[βyαxβxβxαyβxβxαxβy]T(c)=[βyβxαxβxβyαxβxβxαy] .\begin{aligned} T^{(a)} &= \begin{bmatrix} \alpha y & \beta x & \beta x \\ \alpha x & \beta y & \beta x \\ \alpha x & \beta x & \beta y \end{bmatrix} \\ T^{(b)} &= \begin{bmatrix} \beta y & \alpha x & \beta x \\ \beta x & \alpha y & \beta x \\ \beta x & \alpha x & \beta y \end{bmatrix} \\ T^{(c)} &= \begin{bmatrix} \beta y & \beta x & \alpha x \\ \beta x & \beta y & \alpha x \\ \beta x & \beta x & \alpha y \end{bmatrix} ~. \end{aligned}

Training details

Data. Sequences are drawn from the nonergodic composition of the two Mess3 components defined above, mixed with equal weight. Each training sequence begins with a BOS token and then stays inside a single component for all 127 subsequent tokens; the two components share the same three-token alphabet, so no individual token reveals which component is active.

Model. A four-layer decoder-only transformer (TransformerLens HookedTransformer): dmodel=128d_{\text{model}}=128, four attention heads of dimension 32, gated GELU MLPs of width 512, RMS normalization, rotary position embeddings, context length 128, and a vocabulary of four tokens (three emissions plus BOS). Weights are initialized from a Gaussian with standard deviation 0.02, about 3.5×3.5\times smaller than the TransformerLens default of 0.8/dmodel≈0.070.8/\sqrt{d_{\text{model}}} \approx 0.07.

Optimization. AdamW (β1=0.9,β2=0.999\beta_1=0.9, \beta_2=0.999, no weight decay) at a constant learning rate of 10−310^{-3}, with batches of 512 sequences.

Metrics over Training. By step 10,000 the model’s next-token distribution sits within a few 10−410^{-4} nats per token of the optimal loss, and at the step-45,000 checkpoint used for the figures the gap is about 10−410^{-4}. While the loss is falling, the regression error tracks it: 1−R21-R^2 drops roughly as the excess loss to the 0.60.6 power in every layer past the first (every layer in which we find a belief representation). Once the loss reaches its floor, each layer settles onto a floor of its own, later for deeper layers. The animations in this post and the activation explorer use the residual stream after the third of the four blocks, where a linear map recovers the weighted belief vectors with held-out R2≈0.985R^2 \approx 0.985; at the output of the final transformer block, it reaches ≈0.99\approx 0.99.

Held-out error of the linear readout of the weighted belief vectors, by layer over training across eleven checkpoints. While the loss falls, every layer past the first tracks it as roughly (excess loss)^0.6 (black dashed guide); Block 3 is the residual stream queried for the post’s activation figures. Step 1 is one optimizer step from random initialization; the untrained network reads out at R² ≈ 0.45 from recent-token features alone. The first block never rises much above that ceiling.

Generators of Data and the Geometry of Beliefs

The three subsections below restate (at two levels of formality which you can toggle between) the mathematical machinery of our work [5]: hidden Markov models as latent generators of token sequences, belief updating as the structure of prediction, and the geometry of those beliefs.

Fundamentally we are trying to capture the situation relevant to the task prediction over sequences of data, especially when the data is generated by processes that are hidden to the predictor. As such, whatever our formal notion of a generator is, it should have an internal latent space that is hidden from the predictor, a set of rules for how that latent space changes through time (or context position), and a set of rules for how changes in the latent space relate to the observations (or tokens) emitted.

Everything below also holds for generalized HMMs (GHMMs), in which the transition operators may carry negative entries and the predictive vector need not be a probability distribution. Finite GHMMs represent a strictly wider class of processes than finite HMMs — the non-classical geometries of our companion post — but every example here is an ordinary HMM.

In the following sections, we present this work at two levels of formality. If you are interested in the concepts without necessarily the formal mathematics, we suggest choosing the Worked example tab. If, instead, equations are what spark joy in you, choose the General theory tab.

HMMs and their transition operators

General theory

HMMs are an extremely flexible model class: with enough hidden states, essentially any distribution over token sequences can be represented by one. An HMM is defined by the tuple

M  =  (X,  S,  η(∅),  (T(x))x∈X) ,\mathcal{M} \;=\; \bigl( \mathcal{X},\; \mathcal{S},\; \boldsymbol{\eta}^{(\varnothing)},\; (T^{(x)})_{x \in \mathcal{X}} \bigr) ~,

where X\mathcal{X} is the token alphabet, S=Rd\mathcal{S} = \mathbb{R}^d is the latent space, η(∅)∈Rd\boldsymbol{\eta}^{(\varnothing)} \in \mathbb{R}^d is an initial state vector, and each T(x)T^{(x)} is the operator describing latent dynamics for emission xx. The net transition operator T=∑x∈XT(x)T=\sum_{x\in\mathcal{X}} T^{(x)} must have a right eigenvector 1\mathbf{1} with unit eigenvalue. We can then interpret the dynamical systems latent space as carrying a conserved probability mass for which 1\mathbf{1} is the integrator. This means we can interpret the probability of any token sequence x1:Lx_{1:L} as being expressed by

QM(x1:L)  =  η(∅) T(x1)⋯T(xL) 1 .Q_\mathcal{M}(x_{1:L}) \;=\; \boldsymbol{\eta}^{(\varnothing)} \, T^{(x_1)} \cdots T^{(x_L)} \, \mathbf{1} ~.

It is in this sense that we say that an HMM generates a stochastic process.

Worked example

Consider token sequences of the form 0, then 1, then a random bit, and repeats:...0 1 R 0 1 R …. Importantly, sequences can start at any of the three phases. We will call this the Z1R process (for “zero one random”). Here, we are interested in a latent generator of such data. As discussed above, it should have latent states, and dynamical rules telling us how those latent states change through time, and how those changes relate to token emissions.

One such generator for this particular data is a hidden Markov model (HMM). It has three latent states: S0S_0 , S1S_1, and SRS_R, which can be represented by circles in a graph as shown below.

Graphical representation

arrow from state iS0S1SR to state jS0S1SR: token x01, probability p01½

Algebraic representation

T(0) =to statefrom stateS0S1SRS0010S1000SR½00T(1) =to statefrom stateS0S1SRS0000S1001SR½00

T(x01)iS0S1SR jS0S1SR = P(st+1 = jS0S1SR, xt+1 = x01 | st = iS0S1SR)= 0= 1= ½

The Z1R process, drawn as a graph (left) and written as token-labeled transition matrices (right). Every nonzero matrix entry is one arrow. Hover an entry or an arrow to fill in the line beneath each and light up its counterpart.

Sequences of tokens are generated by starting in a particular state (or a distribution over states), then following the arrows according to the probabilities on them. Upon choosing an arrow, the system moves to another latent state (which could be the same one), and emits a token, x∈{0,1}x\in \{0,1\}.

One can represent this system algebraically as well, as a set of token-labeled transition matrices, with one matrix, T(x)T^{(x)}, per token. The entries of these matrices, Tij(x)T^{(x)}_{ij}, are the probability that the system, sitting in state ii, takes the arrow to state jj and emits the token xx. In the figure above on the right, you can see the transition matrices for an HMM that generates the Z1R process.

The only other part needed to define an HMM is the initial state, denoted η(∅)\boldsymbol{\eta}^{(\varnothing)}. In general this can be any probability distribution over the latent states of the system. When the HMM is generating a sequence, you can think of its starting state as being sampled from this initial state13.

Prediction Over Data Generated by HMMs

In the previous section we discussed generators of token sequence data. What is the computational structure of the prediction task, relative to the structure of the latent generator of the token sequence data?

Here we review the answer we established in our previous work: the information that a predictor must represent in order to take in sequences of token and predict future token sequences is given by beliefs, η\boldsymbol{\eta}, over the hidden states of the latent generator of that data.

General theory

We are interested here in the task of prediction of future token sequences given observations of past token sequences. Formally, the conditional probability of any future sequence given the observed context is

QM(xℓ+1:L∣x1:ℓ)=QM(x1:L)QM(x1:ℓ)=η(∅) T(x1)⋯T(xℓ)η(∅) T(x1)⋯T(xℓ) 1 T(xℓ+1)⋯ T(xL) 1Q_\mathcal{M}(x_{\ell+1:L}|x_{1:\ell})=\frac{Q_{\mathcal{M}}(x_{1:L})}{Q_{\mathcal{M}}(x_{1:\ell})} = \frac{\boldsymbol{\eta}^{(\varnothing)} \, T^{(x_1)} \cdots T^{(x_\ell)}}{\boldsymbol{\eta}^{(\varnothing)} \, T^{(x_1)} \cdots T^{(x_\ell)} \, \mathbf{1}}\, T^{(x_{\ell+1})}\cdots \, T^{(x_L)} \, \mathbf{1}

We call the vector encoding the past information the predictive vector (for an HMM, where it is a probability distribution over the hidden states, this is the belief state η\boldsymbol{\eta}):

η(x1:ℓ)  :=  η(∅)T(x1)⋯T(xℓ)η(∅)T(x1)⋯T(xℓ)1 .\boldsymbol{\eta}^{(x_{1:\ell})} \;:=\; \frac{\boldsymbol{\eta}^{(\varnothing)} T^{(x_1)} \cdots T^{(x_\ell)}}{\boldsymbol{\eta}^{(\varnothing)} T^{(x_1)} \cdots T^{(x_\ell)} \mathbf{1}} ~.

This vector is the general answer to the question the coins raised. For the memoryless coins it collapses to the head/tail counts; in general it is everything the past tells you about the future, and nothing more.

Token by token, the same vector updates by one matrix multiplication and a renormalization,

η′  =  η T(x)η T(x)1 ,\boldsymbol{\eta}' \;=\; \frac{\boldsymbol{\eta}\, T^{(x)}}{\boldsymbol{\eta}\, T^{(x)} \mathbf{1}} ~,

where the denominator Z=η T(x)1Z = \boldsymbol{\eta}\, T^{(x)} \mathbf{1} is the probability the observer assigned to the token that just arrived — its next-token prediction.

Because the conditional probability above can be written as

QM(xℓ+1:L∣x1:ℓ)=η(x1:ℓ)T(xℓ+1)⋯ T(xL) 1Q_\mathcal{M}(x_{\ell+1:L}|x_{1:\ell}) = \boldsymbol{\eta}^{(x_{1:\ell})}T^{(x_{\ell+1})}\cdots \, T^{(x_L)} \, \mathbf{1}

iterating this update rule from η(∅)\boldsymbol{\eta}^{(\varnothing)} recovers exactly this closed form.

Worked example

We are interested here in the task of prediction of future token sequences given observations of past token sequences. Intuitively, if we see a sequence of tokens from Z1R in context, like 0110, we would do well to figure out which of the latent states the generator is in. Once we have that, we can then make a prediction for what the next token will be.

In general, given a particular sequence of tokens you will not be able to figure out exactly which latent state the HMM that generated that sequence is in. For instance upon seeing a 0 in context, which can be generated by taking arrows from either S0S_0 to S1S_1 or SRS_R to S0S_0, we won’t know if the HMM is in state S1S_1 or S0S_0. But we can have an optimal belief about which state the HMM is in, in the form of a probability distribution over those states.

Toggle for an example of belief updating by hand

Let’s do this by hand on Z1R. Before any tokens arrive, our belief about what state the HMM is uniform, η(∅)=(13,13,13)\boldsymbol{\eta}^{(\varnothing)} = (\tfrac13,\tfrac13,\tfrac13). After seeing a 1, each entry of the belief is multiplied by the chance that its state emits a 1, the mass moves along that state’s arrow according to the transition matrix for the token 1, T(1)T^{(1)}, and the result is renormalized:

η(∅)T(1) = (⅙, 0, ⅓)
↓ normalize
η(1) = (⅓, 0, ⅔)

That 1 came either from S1S_1 (which emits 1 with certainty, moving the process to SRS_R) or from SRS_R (which emits 1 only half the time, moving to S0S_0); weighing the two likelihoods leaves belief 23\tfrac23 on SRS_R, 13\tfrac13 on S0S_0, and none on S1S_1. Observe a second 1:

η(1)T(1) = (⅓, 0, 0)
↓ normalize
η(11) = (1, 0, 0)

Had the process been in S0S_0, the next token would have been a 0 — the second 1 rules it out. The observer now knows the hidden state exactly: belief has synchronized, and it stays synchronized forever after, hopping deterministically around the corners of the simplex as the cycle turns.

Mathematically, belief updating is Bayes’ rule, with the HMM’s transition matrices as the likelihood. To update your belief upon seeing a new token, multiply your current belief by that token’s transition matrix and renormalize:

η′ = ηT(x) ηT(x)𝟙

η current belief (prior) T(x)  operator for the token seen (likelihood) Z normalization (prob. of that token) η′ updated belief (posterior)

This update rule gives us a belief updating dynamic. The predictor has some current belief about the latent state of the generator, it sees a new token, and it dynamically updates its belief in the service of future token prediction.

The geometry of beliefs

Beliefs are vectors, so they have a geometry.

General theory

Two contexts with identical predictive vectors make identical predictions about all future tokens; contexts with similar predictive vectors η,η′\eta,\eta' make similar predictions because the probability for any future word differs by an amount proportional to η−η′\boldsymbol{\eta}-\boldsymbol{\eta'}. The collection of predictive vectors over all possible contexts forms a geometric arrangement in the latent space, determined entirely by the data-generating process. For a dd-dimensional latent space this arrangement lives in d−1d-1 dimensions (since predictive vectors are normalized).

Worked example

In fact, from the stationary start only seven belief states are ever reachable: the center; three partially-resolved points — (13,23,0)(\tfrac13,\tfrac23,0) after 0, (13,0,23)(\tfrac13,0,\tfrac23) after 1, (12,12,0)(\tfrac12,\tfrac12,0) after 10 — and the three corners. Every context, of any length, lands on one of these seven. This finite constellation in the 2-simplex is the belief state geometry of Z1R. Simple processes give finite constellations; richer processes (like the Mess3 process) fill their simplex with fractal ones; the machinery is identical either way.

The seven reachable belief states of Z1R, as a machine (left — dashed edges are belief updates, labeled token : probability) and placed in the 2-simplex (right). Green states are transient, pre-synchronization beliefs; blue states are the synchronized corners the belief walks forever after. Subscripts name the context that first reaches each state.