PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Language Models Develop Belief State Geometry In-Context

Daniel Balcells, Andrew Jun Lee, Chirag Rastogi, Paul M. Riechers, Adam Shai, Xavier Poncini

arXiv:2609.17376v1cs.LGcs.CL

TL;DR

The paper addresses the limited understanding of representations supporting in-context learning by probing LLMs prompted with HMM data for belief states. Across six open-source models and 40 HMMs, the states are linearly decodable, and interventions on their subspace preserve downstream prediction quality on the order of the untampered model. The findings provide representation-level evidence that HMM in-context learning approximates optimal Bayesian prediction.

  • Problem

    The representations underlying in-context learning remain poorly understood, with few theoretically derived, falsifiable predictions to probe.

    Source: §1 · p. 1

  • Method

    The authors prompt open-source LLMs with HMM-generated sequences, probe residual-stream activations for belief states, and intervene on the probe-identified subspace.

    Source: §1 · pp. 2–3

  • Results

    Across six open-source LLMs and 40 HMMs, belief states are linearly decodable with peak probe R2 values of 0.83–0.99, while targeted interventions preserve downstream prediction quality on the order of the untampered model.

    Source: §1 · pp. 2–3, §5.2 · p. 6

  • Takeaways & Limitations

    The findings extend computational mechanics’ account of belief-state geometry from toy transformers trained on HMM data to open-source LLMs encountering HMM data in context.

    Source: §9 · p. 14

    The evidence is scoped to HMM-generated data; the analyses separately test whether decoding reflects HMM-specific belief information or short-suffix structure.

    Source: §5.2 · pp. 6–7

Abstract

from arXiv · show

Large language models (LLMs) trained on next-token prediction exhibit remarkable in-context learning (ICL) abilities, yet the representations that support ICL remain poorly understood. We consider such representations in a controlled setting: prompting LLMs with data emitted from hidden Markov models (HMMs) and probing for the corresponding belief state -- the posterior distribution over the HMM's hidden states given the observed token history. Across six open-source LLMs prompted with data from 40 HMMs selected for non-trivial belief structure, we find that belief states are linearly decodable from residual stream activations, with peak probe $R^2$-values from 0.83-0.99 across HMM and LLM combinations, ranging from early to late layers. To establish functional relevance, we intervene directly on the probe-identified subspace via patching and steering, resulting in downstream prediction quality on the order of the untampered model, while controls degrade performance substantially. Together, these results provide representation-level evidence that ICL in open-source LLMs approximates optimal Bayesian prediction over a context-inferred generative model. More broadly, our findings extend prior results linking input-distribution structure to activation geometry: from toy networks trained explicitly on HMM data to production-scale LLMs.

1 Introduction

The paper asks what representations support in-context learning, using computational mechanics to predict that LLMs track HMM belief states. Across open-source LLMs, these states are linearly decodable and interventions on their subspace preserve prediction quality near the untampered model.

Motivation

  • The study predicts that LLMs learning to predict HMM-generated data in context represent belief states: posteriors over hidden states given the observed context.The prediction draws on computational mechanics, which treats belief state as the canonical representation of Bayesian inference over an HMM’s hidden state.Source: §1 · pp. 1–2
  • The work addresses a gap in theoretically derived, falsifiable predictions about representations underlying in-context learning, beyond extensive behavioral characterization.It proposes belief-state geometry as a representational target that can be probed and intervened upon.Source: §1 · p. 1

Contributions

  • Across six open-source LLMs and 40 HMMs, belief-state probes reach peak R2 values of 0.83–0.99, and targeted patching and steering preserve prediction quality on the order of the untampered model.Substituting an unrelated point on the belief simplex substantially degrades performance.Source: §1 · pp. 2–3

2 Background

The background defines HMMs and belief states, explains how beliefs support optimal next-token prediction, and introduces finite-order approximations as baselines for LLM predictions.

2.1 Hidden Markov models

  • An HMM has hidden states and observable tokens, with per-token transition matrices encoding joint emission and state-transition probabilities and an initial distribution over hidden states.The paper assumes a stationary initial distribution for the remainder of its analysis.Source: §2.1 · p. 3

2.2 Optimal prediction

  • The belief state is the posterior over hidden states given observed tokens; applying the emission matrix maps it to next-token probabilities, though that mapping need not be invertible.The paper studies reachable belief states as the representational target, restricted to subsets given LLMs’ finite context length.Source: §2.2 · pp. 3–4

2.3 Order-k approximation

  • Because HMM output sequences may depend on distant history, an order-k approximation predicts using only the most recent k tokens and serves as a baseline against LLM performance.The k = 0 case corresponds to i.i.d. token emissions under the initial hidden-state distribution.Source: §2.3 · p. 4

2.4 Transformer notation

  • Transformer residual activations at token position t depend on the input only through its prefix ending at t.The notation defines hℓ,t as the residual stream after the MLP and before the layer’s LayerNorm.Source: §2.4 · p. 4

3 Experimental setup

The experiments use 40 HMMs and six open-source LLMs, with sequences formatted as single-letter tokens and analyses averaged over ten sequences. Figure 2 describes how Qwen’s next-token KL is compared with k-HMM predictors as context grows.

HMMs

  • The study tests 40 HMMs with non-trivial belief geometries across Mess3, Arch, Wing, and Strata, using ten parametrizations per family.Arch, Wing, and Strata were designed for the study; one illustrative parametrization per family appears in the main text.Source: §3 · p. 4

LLMs

  • Six open-source LLMs are evaluated, with Qwen3.5-9B featured in main-text figures; some findings are weaker or absent for the smallest model, Gemma-4-E2B.The other models are Qwen3.5-4B, Llama-3.1-8B, Llama-3.2-3B, and Gemma-4-E4B.Source: §3 · p. 4

Evaluation

  • Figure 2 compares Qwen3.5-9B’s next-token KL against the ground-truth HMM and k-HMM predictors, showing KL over token position and converged KL versus k.The converged KL uses the final 5,000 tokens, with 95% confidence intervals over ten sequences.Source: Figure 2 · p. 5

Formatting and tokenization

  • Sequences contain 20,000 single-letter tokens, and analyses are run and averaged over ten sequences.The formatting is checked so each HMM token aligns with a single LLM token.Source: §3 · p. 5

4 Evaluating in-context prediction accuracy

This section evaluates whether LLMs predict HMM-emitted data accurately, measuring token-level KL divergence against ground truth and comparing post-convergence performance with optimal k-HMM predictors. It finds convergence and performance beyond low-order baselines, while noting that accuracy alone does not identify the underlying mechanism.

4.1 Method

  • Prediction accuracy is measured as KL divergence between the HMM’s ground-truth next-token distribution and the LLM’s output after observing each token history.A KL of zero means the LLM’s next-token distribution perfectly matches the ground truth; performance is compared with the k for which a k-HMM achieves the same accuracy.Source: §4.1 · p. 5, §4 · p. 5

4.2 Results

  • KL falls rapidly over the first 5,000–10,000 tokens and then plateaus, marking the point at which the model’s predictions are considered converged.The authors hypothesize that inference after convergence is primarily over a fixed generative model; the KL curve alone cannot distinguish the two forms of inference.Source: §4.2 · p. 5
  • For Qwen3.5-9B, converged KL matched k-HMM prediction at k = 4.2–11.0 across four main-text parametrizations, with median k = 3.4–4.7 across 10 parametrizations per HMM family.Converged KL was below the 0-HMM and 1-HMM; it was 0.009–0.018 nats for the four shown parametrizations. The other five LLMs were comparable, except Gemma-4-E2B, which was no more accurate than the 1-HMM on 6/10 Wing and 4/10 Strata parametrizations.Source: §4.2 · pp. 5–6
  • The authors distinguish matching prediction accuracy from establishing that the LLM uses a belief-state mechanism, since a short-suffix heuristic could also achieve that accuracy.Subsequent analyses are presented as evidence that LLMs compute and use the belief state for prediction.Source: §4.2 · p. 6

5 Linearly probing the belief state geometry

This section tests whether LLM residual streams encode HMM belief states beyond alternative representations such as next-token probabilities. Probes recover belief states with high held-out fit, while further tests examine HMM specificity and information beyond short-range structure.

5.1 Method

  • Probes fit separately at each layer map residual activations to beliefs using ordinary least squares, with a 20% training split and held-out R2 averaged across 10 sequences.Activations are sampled at positions t ≥ 15,000; each sequence contributes 1,000 training and 4,000 test tokens.Source: §5.1 · p. 6

5.2 Results

  • Peak per-HMM R2 ranges from 0.83 to 0.99 across HMM families and LLMs; shuffled-belief and random-belief controls have extremely low R2.Fits are high in early-to-middle layers and remain high thereafter.Source: §5.2 · p. 6, Figure 3 · p. 7
  • Across most layers, transfer-probe R2 matrices correlate with ground-truth cross-parametrization decodability, supporting HMM-specific information alongside some HMM-generic information.Transfer R2 values are substantially higher in magnitude than ground-truth transfer R2 values.Source: §5.2 · p. 7
  • For Wing and Strata, activation-based belief R2 exceeds ground-truth NTP and log-NTP R2; Arch results vary by parametrization, while Mess3 beliefs and NTP are linearly indistinguishable.For Mess3, the belief-to-NTP map is invertible, so linear probes cannot distinguish the two representations.Source: §5.2 · pp. 8–9
  • Across HMMs and LLMs, activation-minus-NTP or log-NTP R2 tracks the theoretically predicted surplus, with Pearson r = 0.85–1.00 except 0.43–0.88 for Arch under log-NTP.The comparison relates activation-based surplus belief information to the theoretical amount not explained by ground-truth NTP or log-NTP.Source: §5.2 · pp. 8–9
  • After sufficient k, full-HMM activations outperform 0-HMM and 1-HMM activations for most parametrizations, indicating belief information beyond the Markov order-one approximation.The comparison probes full-HMM beliefs using activations induced by full-HMM, 0-HMM, or 1-HMM sequences.Source: §5.2 · pp. 7–8, Figure 5 · p. 9
  • NTP and log-NTP are also linearly decodable from activations, but high log-NTP decodability does not follow from belief-state decodability.Beliefs map linearly to NTP, but not to log-NTP.Source: §5.2 · p. 9
  • Belief and log-NTP representations are jointly present, but this section leaves their causal relationship unresolved and frames beliefs as a hypothesized intermediate representation.The following section uses causal interventions to test that hypothesis.Source: §5.2 · p. 9

6 Controlling representations via the belief state

The section tests whether belief-state representations influence later log-NTP representations, using belief-space interventions and matched controls. Steering changes log-NTP decodability and affects representations beyond the NTP subspace.

6 Controlling representations via the belief state

  • To separate belief representation from NTP representation, the δ direction lies in the emission matrix’s kernel; random-direction controls match its steering magnitude.The δ intervention is designed to leave NTP and log-NTP unchanged.Source: §6 · pp. 9–10, §6.1 · pp. 10–11

6.1 Method

  • A layer- and sequence-specific affine embedding maps belief states into residual-stream activations, providing the subspace used for interventions.The embedding is identified by linear regression from beliefs to activations.Source: §6.1 · p. 10
  • The tests intervene on the final 5,000 positions and compare past-consistent targets with past-inconsistent targets that conflict with the input prefix.For log-NTP steering, the authors use past-inconsistent targets as the stricter propagation test.Source: §6.1 · p. 10
  • Steering adds the target-minus-current belief activation difference, whereas patching replaces the activation with the embedded target belief.Section 6 uses steering to preserve activation components outside the embedding’s image; patching is also used in Section 7.Source: §6.1 · p. 10

6.2 Results

  • Belief steering from later layers makes the target belief’s implied log-NTP highly decodable downstream, while earlier steering can trigger self-repair.Similar results hold across HMMs and LLMs except the Gemma class.Source: §6.2 · p. 11, Figure 6 · p. 11
  • Steering along an NTP-preserving belief direction installs its target component downstream, unlike matched random-direction steering, with distinct self-repair patterns.The δ-component effect is especially evident when steering in the final half of the network; Mess3 is excluded because its emission matrix is invertible.Source: §6.1 · pp. 10–11, §6.2 · p. 11, Figure 6 · p. 11

7 Controlling predictions via the belief state

The authors intervene on belief-state activations and evaluate whether predictions shift toward the injected belief. Predictions move toward target beliefs, including when they conflict with the unintervened prefix, while arbitrary-belief controls do not produce that shift.

7.1 Method

  • The study varies the intervened suffix over k = 1, 5, and 10 positions to test how much unintervened context can compete through attention.Longer suffixes reduce the conflicting context available from unintervened positions.Source: §7.1 · p. 12
  • Prediction shifts are measured by KL divergence to the NTPs implied by the target and original final belief states.The metric is evaluated at the final token position.Source: §7.1 · p. 12
  • Both past-consistent and past-inconsistent targets are tested to quantify the effect of conflict with the unintervened prefix.The two target conditions differ in consistency with that prefix.Source: §7.1 · p. 12

7.2 Results

  • Across intervention suffix lengths, patched predictions are closer to the target belief’s NTP than the original belief’s, with target KL on the order of unmodified-model KL.A random-belief control does not show the target-directed shift; steering produces similar results.Source: §7.2 · pp. 12–13
  • The target-directed shift grows with suffix length, and past-inconsistent targets still succeed despite conflicting with every unintervened position.The advantage of past-consistent targets at k = 1 shrinks as k grows and disappears by k = 10.Source: §7.2 · p. 13

8 Decoding predictions from intermediate layers

The authors test whether intermediate activations contain output-relevant structure by comparing tuned-lens early-exit predictions with belief-state decodability. Ground-truth-trained tuned lenses perform well, and their accuracy covaries with belief-state decodability, though this does not establish causation.

8.1 Method

  • A tuned lens learns a per-layer affine map that transforms residual-stream activations into vocabulary distributions, trained here toward ground-truth NTPs and shuffled or random controls.The authors also report a tuned lens trained toward the model’s own output distribution in Appendix Q.1.Source: §8.1 · p. 13

8.2 Results

  • Ground-truth-trained tuned lenses attain low KL across most layers, below the LLM’s own prediction KL, while shuffled- and random-NTP lenses perform substantially worse.This pattern holds across all HMMs and LLMs.Source: §8.2 · p. 13
  • Tuned-lens accuracy is highest where belief states are most decodable, and its per-layer KL covaries with probe R2 beyond early-exit controls.The covariation holds for most parametrizations and also appears with the canonical tuned lens.Source: §8.2 · pp. 13–14, Figure 8 · p. 14
  • The observed covariation indicates that belief-representing layers also contain output-relevant NTP structure, but does not show that belief representations cause early-decoding accuracy.Log-NTP decodability is another correlate, though the authors treat belief states as more fundamental because log-NTP is downstream.Source: §8.2 · p. 14

9 Discussion

The discussion argues that computational mechanics’ predictions extend from toy transformers to open-source LLMs encountering HMM data in context: belief states are linearly decodable and the corresponding subspace is causally relied upon for next-token prediction. It also emphasizes limits in the tested HMM regimes and the need for further evidence about the nature of the decoded representation.

9 Discussion

  • The tested LLMs’ belief-state subspace is linearly decodable and causally relied upon for next-token prediction, extending computational mechanics’ account to open-source LLMs in-context.The authors frame this account as a source of advance, falsifiable predictions about activation geometry in production-scale LLMs.Source: §9 · p. 14
  • The results concern low-entropy HMMs with non-trivial belief geometry; arbitrary HMMs, especially those with high-dimensional or indistinctive geometry, remain untested.A broader parameter sweep also found regimes with desired KL trends but no gap in short-suffix belief decodability.Source: §9 · p. 14

A LLM Usage

  • The authors acknowledge using LLM-powered tools for literature review, experimental design, meeting notes, code, and writing, with outputs reviewed and approved by human authors.Source: A LLM Usage · p. 19

B Related work

  • Prior work found linearly recoverable ICL representations and belief manifolds; this paper tests a computational-mechanics geometric prediction for richer HMM processes in pretrained LLMs.The cited work also includes toy transformers trained on HMM data and analyses of belief updates for Mess3.Source: B Related work · p. 20

C.1 Process definitions

  • The process definitions specify valid HMMs through per-token transition matrices whose rows are stochastic; Figure 9 depicts the representative belief-state geometries.Mess3, Wing, and Strata have three states, while Arch has four.Source: §C.1 · p. 20, Figure 9 · p. 21

C.2 Parameter selection

  • The 40-process sweep selects ten parameterizations per family across entropy and mixing-rate properties, targeting a region associated with small prediction convergence gaps.Figure 10 includes combined, transition, and emission entropy against mixing rate; Mess3 varies entropy at near-fixed mixing rate as a control.Source: §C.2 · p. 22, Figure 10 · p. 22

D All HMMs’ per-layer highest R2 decoded geometries for Qwen3.5-9B

  • For Qwen 3.5 9B, linear probes recover geometry at each HMM’s best-fitting layer, with held-out R2 averaged over 10 sequences (4,000 positions each); belief-state geometry is visibly weaker in the first 250 tokens.Control probes trained on shuffled or random belief states recover no structure.Source: Figure 11 · p. 23, Figure 13 · p. 25, Figure 14 · p. 26

E Controlling predictions via steering for Qwen3.5-9B

G In-context prediction accuracy for all models and all HMMs

The appendix compares each LLM’s next-token probabilities with the full HMM’s ground-truth predictions across all 40 HMMs. It also plots order-0 and order-1 HMM baselines against the full HMM.

G In-context prediction accuracy for all models and all HMMs

H k-HMM KL crossover for all models and all HMMs

The appendix defines a crossover order by matching an order-k HMM approximation’s KL to the full HMM with the LLM’s converged KL. It presents the distribution of that order for each model and HMM family.

H k-HMM KL crossover for all models and all HMMs

  • Figures 23–24 plot crossover-order distributions for Qwen 3.5 9B and Qwen 3.5 4B across the four HMM families.The crossover is the smallest k for which the k-HMM’s KL to the full HMM falls to the LLM’s converged KL.Source: Figure 23 · p. 32, Figure 24 · p. 32, H k-HMM KL crossover for all models… · p. 32
  • Figures 25–26 show crossover-order distributions for Llama 3.1 8B and Llama 3.2 3B across 10 parametrizations of each HMM family; k is the smallest Markov order whose predictor matches the LLM’s converged KL.Source: Figure 25 · p. 32, Figure 26 · p. 33
  • Figures 27–28 show crossover-order distributions across 10 parametrizations of each HMM family for Gemma 4 E4B and E2B; crossover k is the smallest Markov order whose predictor matches the LLM’s converged KL.Source: Figure 27 · p. 33, Figure 28 · p. 33

I Linear probes for all models and all HMMs

The appendix reports layerwise linear-probe R2 fits for belief states and shuffled- and random-belief controls. For each model, the plots cover all 40 HMMs.

I Linear probes for all models and all HMMs

J Cross-parametrization for all models and all HMMs

The cross-parametrization analyses test whether activations carry belief information specific to the HMM that generated the sequence and whether that information exceeds an order-1 approximation. The figures compare transfer between parametrizations and measure belief information beyond order one.

J Cross-parametrization for all models and all HMMs

  • The analysis correlates activation-based transfer-probe R2 matrices with ground-truth transfer matrices and compares decoding of generating versus alternative parametrizations’ beliefs.The ground-truth matrix uses belief states rather than activations.Source: J Cross-parametrization for all models… · p. 37
  • Probes predict the generating HMM’s beliefs better than the selected alternative HMM’s beliefs across the six model-specific figure sets.The alternative is the parametrization with the lowest ground-truth transfer R2 from activations on the generating HMM’s sequence.Source: Figure 35 · p. 37, Figure 36 · p. 38, Figure 37 · p. 38

K 1-HMM belief state comparisons for all models and all HMMs

  • The order-1 comparison measures the difference in belief-state probe fit between full-HMM and 1-HMM sequences at suffix length k = 20.Its theoretical counterpart is 1 − R2 for predicting the true belief state from the 1-HMM belief state on the same sequence.Source: K 1-HMM belief state comparisons for… · p. 41
  • Figures 41–46 show layer-wise ∆R2 and compare each sequence’s best-layer ∆R2 with its theoretical counterpart across six models.Each plotted point represents a sequence, with 10 parametrizations and 10 seeds per family.Source: Figure 41 · p. 41, Figure 42 · p. 42, Figure 43 · p. 43

L Next-token (log) probability comparisons for all models and all HMMs

The analyses test whether activations encode belief information beyond what next-token probabilities or an order-preserving steering direction can explain. They also examine whether steering belief-related directions changes later readouts and downstream predictions.

L Next-token (log) probability comparisons for all models and all HMMs

  • ∆R2 measures belief information in activations beyond what ground-truth NTP or log-NTP can explain, with the theoretical counterpart given by 1 − R2 from those probabilities.For Mess3, the belief-to-NTP map is invertible, making the theoretical counterpart zero and the correlation undefined.Source: L Next-token (log) probability… · p. 47, Figure 47 · p. 48

M Changing log-NTP by steering beliefs for all models and all HMMs

N Steering beliefs without changing NTP decodability for all models and all HMMs

O Random direction steering for all models and all HMMs

P Controlling predictions for all models and all HMMs

  • Patching and steering inject target beliefs at the last k ∈ {1, 5, 10} positions, and evaluate prediction KL to the target belief’s ground-truth NTP.The conditions are past-consistent, past-inconsistent, and random beliefs; lower KL means the prediction moved closer to the injected belief.Source: P Controlling predictions for all… · p. 72, Figure 71 · p. 73

Q.1 Canonical tuned lens

The canonical tuned lens provides a complementary readout of the model’s latent predictions, while ground-truth-supervised lenses and controls are evaluated across layers. The analyses also measure how lens KL tracks belief-probe decodability.

Q.1 Canonical tuned lens

  • At intermediate layers, the canonical tuned lens’s predictions are closer to HMM ground-truth NTP than the model’s final output.Because this lens is trained only to match the model’s own final-layer distribution, it cannot encode ground-truth information; the authors interpret the gap as latent predictions readable at early-to-mid layers.Source: §Q.1 · p. 79

Q.2 Layer-wise tuned lens for all models and all HMMs

  • The layer-wise lens analysis compares per-layer KL to HMM ground-truth NTP for ground-truth-supervised and shuffled- and random-NTP control lenses.The control lenses’ KL remains high across layers.Source: §Q.2 · p. 83, Figure 83 · p. 83

Q.3 Correlation of tuned lens KL and belief probe R2 for all models and all HMMs

  • The correlation analysis computes, for each sequence, the Pearson correlation between the lens’s KL-by-layer curve and the belief probe’s 1 − R2-by-layer curve.It includes ground-truth-supervised lenses and shuffled- and random-NTP controls.Source: §Q.3 · p. 86

Related papers

Loading 2609.17376v1…