We propose Bayesian optimal sequential prediction as a new principle for understanding in-context learning (ICL). Unlike interpretations framing Transformers as performing implicit gradient descent, we formalize ICL as meta-learning over latent sequence tasks. For tasks governed by Linear Gaussian State Space Models (LG-SSMs), we prove a meta-trained selective SSM asymptotically implements the Bayes-optimal predictor, converging to the posterior predictive mean. We further establish a statistical separation from gradient descent, constructing tasks with temporally correlated noise where the optimal Bayesian predictor strictly outperforms any empirical risk minimization (ERM) estimator. Since Transformers can be seen as performing implicit ERM, this demonstrates selective SSMs achieve lower asymptotic risk due to superior statistical efficiency. Experiments on synthetic LG-SSM tasks and a character-level Markov benchmark confirm selective SSMs converge faster to Bayes-optimal risk, show superior sample efficiency with longer contexts in structured-noise settings, and track latent states more robustly than linear Transformers. This reframes ICL from "implicit optimization" to "optimal inference," explaining the efficiency of selective SSMs and offering a principled basis for architecture design.
Figures & tables
Figure 1: The Two Paradigms of ICL. (Left) GD/Transformer: The context is pooled into an empirical loss landscape (grey surface). The forward pass performs gradient descent (red path) to find a minimizer y^GD . This ignores the temporal structure. (Right) Selective SSM/Bayesian Filtering: Each observation updates an internal belief state (blue cloud) over latent variables. Prediction is the mean of the evolved belief. This process naturally accounts for temporal correlations and uncertainty. Theorem 2 shows the right paradigm achieves strictly lower risk for tasks with structured temporal noise.
Figure 2: Convergence to Bayes Optimality. (Left) Excess risk versus number of meta-training tasks. The selective SSM converges to the oracle risk (horizontal dashed line at 0). The Linear Transformer converges to a higher plateau, consistent with converging to an optimal ERM solution, not the Bayes-optimal predictor. (Right) Histogram of prediction errors for 1000 test sequences at the end of training. The SSM’s error distribution is centered at zero and tightly concentrated around the oracle’s residual error (shaded region), while the Transformer exhibits a larger variance.
Figure 3: Performance on HMM-Generated Text. (Top) Next-character prediction accuracy versus context length k . Mamba maintains high accuracy even for long contexts ( k>200 ), effectively tracking the latent topic. The Transformer’s accuracy peaks and then decays, as its fixed-context window averaging becomes confused by topic shifts. (Bottom) A qualitative example. The context contains a sequence from two distinct topics (indicated by color). Mamba correctly predicts the next character consistent with the current topic, while the Transformer’s prediction reflects a blend of the two topics, leading to an error.
Transformers and state-space models (SSMs) are two prominent sequential learning architectures, yet their comparison remains largely empirical and existing theoretical analyses are typically task-specific or architecturally restricted. In this paper, we develop belief geometry, a unified analytical framework for comparing the representational capabilities of broad classes of attention and SSMs. Starting from a generalized formulation of in-context linear regression and using cumulative Bayes regret as our measure, we abstract three capabilities required by many sequential learning problems in our belief geometry: evidence assembly, belief maintenance, and addressing. We then study three cases of our formulation that isolate these capabilities and yield sharp architectural lessons: For belief maintenance, SSMs attain the optimal regret over stationary aggregation kernels; for positional assembly, SSMs have a memory advantage; and for content addressing, softmax attention has an exponential width advantage over sigmoid-selective SSMs. Experiments with LLaMA-type Transformers and Mamba-2 show that these architectural insights extend beyond our analytically tractable classes and linear-regression testbed.
Enes Arda, Semih Cayci, Atilla Eryilmaz
The Ohio State University · RWTH Aachen University
In-context learning (ICL) allows Transformers to adapt to novel tasks without weight updates, yet the underlying algorithms remain poorly understood. We adopt a statistical decision-theoretic perspective by investigating simple binary hypothesis testing, where the optimal policy is determined by the likelihood-ratio test. Notably, this setup provides a mathematically rigorous setting for mechanistic interpretability where the target algorithmic ground truth is known. By training Transformers on tasks requiring distinct geometries (linear shifted means vs. nonlinear variance estimation), we demonstrate that the models approximate the Bayes-optimal sufficient statistics from context up to some monotonic transformation, matching the performance of an ideal oracle estimator in nonlinear regimes. Leveraging this analytical ground truth, mechanistic analysis via logit lens and circuit alignment suggests that the model does not rely on a fixed kernel smoothing heuristic. Instead, it appears to adapt the point at which decisions become linearly decodable: exhibiting patterns consistent with a voting-style ensemble for linear tasks while utilizing a deeper sequential computation for nonlinear tasks. These findings suggest that ICL emerges from the construction of task-adaptive statistical estimators rather than simple similarity matching.
Faris Chaudhry, Siddhant Gadkari
Department of Computer Science Imperial College London
State Space Models (SSMs) have emerged as a compelling alternative to Transformers, enabling sequence modeling with constant memory and linear compute. Although SSMs exhibit reasonable performance and favorable computational characteristics, they continue to lag behind Transformers on tasks that require in-context learning and precise retrieval, slowing their adoption for large-scale language modeling. In this work, we demonstrate that both the success and failure of SSMs in these domains can be explained by studying the role of the gating mechanism, a prevalent component in modern recurrent networks. Specifically, we show through theory and experiments that this gating mechanism causes SSMs to first learn an in-weights "memorization" solution, while delaying, or even preventing, convergence to a correct in-context learning solution. Importantly, this happens even in cases where there are no fundamental limitations due to the architecture or its memory capacity. On the other hand, we find that gating is often beneficial for improving generalization to long sequence lengths. Our results illuminate the crucial role of the gating mechanism in shaping both the training dynamics and generalization of SSMs, and provide a basis for understanding and improving linear-time models.