Large Language Models Develop Belief State Geometry In-Context

TL;DR: Take an off-the-shelf LLM, prompt it with 20,000 tokens sampled from a hidden Markov model it has never seen, and look inside. This paper finds that the residual stream comes to contain a linearly-embedded copy of the Bayesian posterior over the HMM's hidden states — the exact quantity an optimal predictor would track — with probe R^2 of 0.83–0.99 across six open models and 40 processes. Crucially, the authors go beyond decoding: they steer along a direction that changes the belief but provably leaves the next-token distribution unchanged, and show the model treats that direction as meaningful. It's the strongest evidence yet that in-context learning in production LLMs is implemented as something like online Bayesian filtering, not a bag of suffix heuristics.

The theory that made a prediction before anyone looked

Most mechanistic interpretability works backwards: find a structure in activations, then invent a story for it. This paper works forwards, from a framework called computational mechanics (Crutchfield, Shalizi), which answers a precise question: if you want to optimally predict the next symbol from a stochastic process, what is the minimal thing you must keep in memory?

For a hidden Markov model the answer is the belief state: the posterior distribution \eta(x_{1:t}) over the HMM's hidden states given everything observed so far. Formally, with per-token transition matrices T^{(x)} (entry ij = probability of emitting token x while moving from hidden state i to j) and prior \eta^{(\emptyset)},

i.e., Bayes' rule applied token by token. The set of all reachable posteriors traces a characteristic shape inside the probability simplex — often a striking fractal — called the belief-state geometry. In 2024, Shai et al. showed that tiny transformers trained from scratch on HMM data linearly embed this geometry in their residual stream. Separately, Dai et al. (2025) showed that production LLMs predict HMM sequences near-optimally purely in-context, despite never training on them.

This paper closes the loop: does the theoretically-mandated representation appear in Llama, Qwen, and Gemma when the HMM shows up only in the prompt? That's a genuinely different claim. A model trained on one process can bake its belief update into weights. A pretrained LLM must, within a single forward context, (1) infer which process is generating the tokens and (2) run a Bayesian filter over that inferred process's hidden states — all with frozen weights.

Why "belief state" is a stronger claim than "predicts well"

Here's the subtlety that makes the paper interesting. The belief state determines the next-token probabilities (NTP) via a fixed linear map, the emission matrix M: p(x_{1:t}) = \eta(x_{1:t}) M. But M is generally not invertible — several distinct beliefs can imply the same next-token distribution. The belief state carries more information than the output: it encodes what the model should predict many steps ahead, not just one.

That asymmetry is the paper's experimental lever. If the model only represented output statistics, everything decodable about beliefs would live in the row space of M. But if it represents beliefs proper, there's a direction \delta \in \ker(M^\top) — a way to move the belief that changes nothing about the next token — and the model should still track it.

The setup

The authors design 40 HMMs across four families (Mess3, plus three new processes — Arch, Wing, Strata — deliberately unlikely to be in pretraining data), each with a non-trivial belief geometry, and choose parametrizations in the low-entropy, slow-mixing regime where Dai et al. showed LLMs predict well. Tokens are emitted as space-separated letters ("F Q V F Q ..."), verified to tokenize one HMM symbol per LLM token. Sequences are 20,000 tokens; six models are tested (Qwen3.5-9B/4B, Llama-3.1-8B, Llama-3.2-3B, Gemma-4-E4B/E2B).

Behavior first. The KL divergence between the LLM's output and the ground-truth next-token distribution falls rapidly over the first 5–10k tokens and plateaus at 0.009–0.018 nats for Qwen3.5-9B — matching an optimal Markov predictor of order k \approx 4–11 depending on the process. So the model isn't just memorizing bigrams; it's exploiting long-range structure.

Probing. Ordinary least squares from residual activations h_{\ell,t} (late in context, t \geq 15{,}000) to the ground-truth belief \eta(x_{1:t}), trained on 20% of positions, tested on the rest. Held-out R^2 peaks between 0.83 and 0.99 across every model × HMM combination, at layers ranging from early to late. Shuffled-belief and Dirichlet-random-belief controls score negative R^2.

High R^2 alone proves little — a 4096-dimensional space can linearly "explain" many low-dimensional targets. The paper's value is in the battery of alternative hypotheses it eliminates:

• "It's just a suffix heuristic." They generate control sequences from order-0 and order-1 approximations of each HMM (which preserve short-range token statistics but destroy long-range structure) and try to decode k-suffix beliefs from those activations. For most parametrizations, true-HMM activations decode beliefs beyond what order-1 activations can, and the size of the surplus correlates (r = 0.59–0.86) with the theoretical information gap between order-1 and full beliefs.
• "It's just the output distribution." They regress ground-truth NTP (and log-NTP) onto beliefs and compare. For Wing and Strata, activations decode beliefs better than the ground-truth NTP can — and the gap tracks the theoretical residual 1 - R^2(\text{NTP}\to\text{belief}) with Pearson r up to 1.00. Mess3 is the honest negative control: its M is invertible, so beliefs and NTP are linearly identical, and indeed no surplus appears.
• "It's generic to the HMM family, not process-specific." Transfer probes trained on one parametrization decode a similar parametrization's beliefs well and a dissimilar one's poorly, in a pattern that correlates (r \approx 0.64–0.78) with the ground-truth cross-predictability of the beliefs themselves — though transfer R^2 is higher than theory predicts, so some family-generic information is also present.

The causal experiments: the \delta-steering trick

The centerpiece is Section 6–7. The authors fit an affine embedding \text{emb}_\ell from belief space into activation space, then intervene:

• Steering: add \Delta = \text{emb}_\ell(\eta_{\text{target}}) - \text{emb}_\ell(\eta_{\text{current}}), preserving everything outside the belief subspace.
• Patching: overwrite the activation with \text{emb}_\ell(\eta_{\text{target}}) entirely.

Three results stand out.

1. Steered beliefs propagate to output representations. Steer beliefs at layer \ell toward those of a different, conflicting token history; probes at later layers then decode the log-NTP implied by the injected belief, not the true history's. This works when steering beyond roughly the first third of the network. Steer earlier and a "self-repair" mechanism erases the edit and reinstalls the original — consistent with the belief being actively recomputed from context in early layers.

2. The NTP-invisible component is functionally privileged. Steering along \delta \in \ker(M^\top) — changing the belief without changing what the model should output next — installs the target \delta-component at subsequent layers, in a pattern clearly distinct from steering along a magnitude-matched random direction v. The model transports this output-irrelevant coordinate through its layers as if it matters. It's hard to explain that if the "belief subspace" were merely an incidental linear shadow of output statistics.

3. Injected beliefs control predictions. Patch target beliefs into the last k \in \{1,5,10\} positions and measure the final-token KL. The KL to the injected belief's implied distribution falls well below the KL to the distribution implied by the actual token history, approaching the unmodified model's baseline KL — even when the injected belief contradicts every unintervened position in a 20k-token context. Injecting a random belief and scoring against the same target shows no such drop: you can't get there by perturbing arbitrarily; you must write the specific belief. At k=1, conflicting context still exerts some pull; by k=10 the injected belief fully dominates.

A complementary tuned-lens analysis (Section 8) shows early-exit prediction accuracy covaries layer-by-layer with belief decodability, and — a nice aside in Appendix Q — the canonical tuned lens reveals the model's intermediate-layer latent prediction is closer to the HMM ground truth than its own final output, as if late layers reimpose natural-language priors that slightly miscalibrate it on synthetic data.

What this means, and what to doubt

If it holds up, this is the cleanest bridge yet between a normative theory of prediction and the actual geometry of production LLM activations. Computational mechanics didn't just fit a post-hoc story — it named the representation, its shape (a specific subset of the simplex, sometimes fractal), and its algebraic relationship to the output (p = \eta M, with an exploitable kernel), all in advance. The result also gives concrete substance to the "ICL as implicit Bayesian inference" account (Xie et al., 2022): not merely behavior consistent with Bayes, but a linearly-addressed posterior you can read and write.

Reasons for skepticism, most of which the authors state themselves:

• The setting is maximally clean. Single-letter tokens, stationary low-entropy HMMs, 20k tokens of pure process data. Naturalistic ICL involves compositional, non-stationary, high-entropy latent structure. Nothing here shows belief-state geometry survives that transition — the authors flag factored processes as future work.
• Parametrization selection matters. Processes were pre-filtered to the regime where LLMs predict well, and Table 3 lists a nontrivial subset of parametrizations with good KL curves but no decodability gap over the order-1 approximation. The geometry story is regime-specific, not universal.
• Linear probes over-find structure in high-dimensional activations (the "dead salmon" concern they cite). The controls and, especially, the \delta-steering substantially mitigate this — a spurious regression fit shouldn't yield a direction that the network causally transports — but unsupervised discovery of the geometry is still missing.
• Log-NTP is jointly decodable, sometimes better than beliefs. The claim that beliefs are upstream of log-NTP rests on the steering results, which are suggestive but not a full circuit-level account of the update mechanism.
• Gemma-4-E2B, the smallest model, partially fails — behaviorally (no better than order-1 on many Wing/Strata parametrizations) and in the steering experiments (both Gemma models are excluded from the causal claims). This hints belief geometry sharpens with scale, but it's also a caveat on generality.

Where to spend your time: Section 6 and Figure 6 — the \delta-direction steering is the paper's sharpest and most original piece of evidence, and the design pattern (find a theory-defined direction that is invisible to the output, then show the network respects it anyway) is one interpretability researchers should steal for other domains. Figure 5's dual controls (suffix approximations and NTP confounds) are the second-best stop, a model of how to falsify the cheap alternative explanations before claiming the interesting one.