Large language models internally track hidden state probabilities
Large Language Models Develop Belief State Geometry In-Context
Machine LearningComputation and Language
Summary
Large language models (LLMs), which predict text word-by-word, surprisingly learn to understand hidden patterns in the data they see during prompting. The authors tested LLMs by feeding them sequences generated from hidden Markov models, which have unseen hidden states influencing the output. They found that LLMs internally represent these hidden state probabilities in a way that can be linearly extracted and that this internal belief helps the models make better predictions. This shows that the LLM’s text prediction involves an implicit form of reasoning similar to a statistical model.
What this means in practice
- •For language model engineers: Identify and manipulate internal belief state representations to improve model interpretability and control prediction behavior.
- •For natural language processing tool builders: Design systems that exploit internal probabilistic structures in LLMs for better context-aware text generation and understanding.
Authors
Daniel Balcells, Andrew Jun Lee, Chirag Rastogi, Paul M. Riechers, Adam Shai, Xavier Poncini
Abstract
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.