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.