Technical Summary: Extracting Algorithms in Pre-trained LLMs: A Case on Hidden Markov Models

Problem Statement

Large Language Models (LLMs) demonstrate a striking ability to perform in-context learning (ICL) on sequences generated by Hidden Markov Models (HMMs), often approaching Bayes-optimal prediction performance. However, the specific algorithmic mechanism underlying this capability remains undetermined. Prior work has proposed several candidates—including spectral methods, belief-state tracking, and gradient-based feature learning—but there is no consensus on which algorithm is actually implemented. Crucially, previous proposals have not been validated against the model's internal activations, leaving a gap between observed external behavior and internal mechanistic understanding.

Methodology

The authors propose a three-stage pipeline to bridge this gap, connecting ICL behavior to internal mechanisms:

1. Empirical Benchmarking: The study systematically benchmarks 12 pre-trained LLMs (from Qwen, Llama, Gemma, and OLMo families) against a suite of classical and learning-based HMM predictors across 75 distinct HMM configurations. These configurations vary in transition entropy (mixing rate), emission entropy (noise), and alphabet size. The goal is to narrow the space of plausible algorithms by identifying which external baselines best match LLM behavior across different regimes.
2. Theoretical Construction: The authors derive theoretical connections showing that Transformers are capable of implementing the three leading candidate algorithm classes (Linear n-gram, Non-linear n-gram, and Spectral) via in-context learning. They provide explicit constructions expressing next-token probabilities as linear functions of nonlinear features built from past observations. This is validated by training a small Transformer on a single HMM and analyzing its convergence.
3. Mechanistic Probing (PAP): To determine which algorithm is actually used by pre-trained LLMs, the authors introduce the Principal Activations Probe (PAP). This is a layer-wise probing and intervention method that:
   • Projects residual stream activations onto their top principal components to isolate low-dimensional algorithmic signals.
   • Fits linear probes to decode candidate algorithmic representations (e.g., belief states, operators) from these activations.
   • Performs causal interventions (activation patching) to test whether these decoded representations causally drive the model's predictions.

Key Contributions

1. Empirical Narrowing of Mechanisms: The study demonstrates that no single classical algorithm explains LLM behavior across all HMM regimes. However, the behavior is collectively spanned by three algorithm classes: Linear n-gram, Non-linear n-gram, and Spectral. The Baum-Welch algorithm is ruled out as a candidate because its performance profile (strong at short contexts via explicit parameter estimation) differs fundamentally from the gradual in-context generalization of LLMs.
2. Theoretical Unification: The authors prove that Transformers can implement all three candidate ansatzes. They show that next-token probabilities can be expressed as linear functions of nonlinear features (Kronecker products of one-hot vectors) and that Transformers can learn these features and the associated linear maps via gradient descent-like processes. A small trained Transformer is shown to naturally converge to a finite-window Linear n-gram predictor.
3. Causal Identification via PAP: The Principal Activations Probe reveals that pre-trained LLMs do not implement a single static algorithm but recruit different representations depending on the HMM regime:
   • Soft n-gram (Linear/Non-linear n-gram): This representation is the most causally effective across both "belief-dominated" (slow mixing, high noise) and "observation-dominated" (fast mixing, low noise) regimes.
   • Decodability vs. Causal Use: The study highlights a critical dissociation: algorithmic information (e.g., bigram statistics) may be decodable from early layers but is bypassed by downstream computation. Causal intervention is required to confirm that a representation actually drives the output.
   • Layer Specialization: Distinct computational stages are localized to different layers. Early layers often encode information that is not causally used, while late layers contain the representations that directly determine predictions.

Results

• Convergence to Optimality: LLM predictions consistently approach the Bayes-optimal Oracle as context length grows, regardless of model size, though convergence speed varies.
• Regime-Dependent Strategies:
  • In observation-dominated regimes (fast mixing, low emission entropy), simple bigram-like statistics suffice, and the model behaves similarly to a Markov chain predictor.
  • In belief-dominated regimes (slow mixing, high emission entropy), the model requires integrating observations over long contexts to maintain a latent state belief. Here, the Soft n-gram representation (a learned finite-window feature) emerges as the primary causal driver, outperforming simple bigrams and spectral methods in causal intervention tests.
• Small Transformer Validation: A small Transformer trained on a fixed HMM recovers the theoretical finite-window Linear n-gram predictor almost exactly, with attention mechanisms aggregating recent tokens into a linearly decodable feature representation.
• PAP Findings: The probe confirms that the "belief" (posterior over hidden states) is linearly decodable from the residual stream, particularly in later layers. However, causal patching reveals that in some regimes (e.g., early layers of belief-dominated tasks), high decodability does not imply causal usage.

Significance and Claims

The paper claims to advance the understanding of how LLMs perform ICL on HMMs by connecting external performance to internal mechanisms. The central conclusion is that LLMs do not implement statistically optimal, iterative inference (such as running the forward algorithm or Baum-Welch). Instead, they appear to learn finite-window, gradient-descent-like approximations well-described by Soft n-gram statistics.

This holds even in regimes where longer-range integration would theoretically yield lower cross-entropy loss. The study emphasizes that the "algorithm" an LLM implements is not a fixed, monolithic procedure but a flexible strategy that shifts based on the difficulty and structure of the underlying data distribution. The introduction of PAP provides a rigorous framework for distinguishing between representations that are merely present in the model's activations and those that are causally necessary for prediction.

The authors note limitations, restricting their analysis to synthetic HMMs with small finite alphabets and relying on linear probing, which may not capture fully nonlinear or distributed representations. They suggest extending PAP to natural language tasks with latent structure and performing circuit-level analysis as future directions.