A Bayesian Perspective on the Role of Epistemic Uncertainty for Delayed Generalization in In-Context Learning

2026-04-14

AI Summary

This study investigates why grokking—sudden task generalization after prolonged training in in-context learning—emerges abruptly and reveals the pivotal role of epistemic uncertainty in this phenomenon. Adopting a Bayesian perspective, the authors analyze the evolution of predictive uncertainty through modular arithmetic tasks, establishing the first explicit link between epistemic uncertainty and grokking. Leveraging approximate Bayesian inference to estimate posterior distributions, they conduct experiments varying context length and noise levels, complemented by spectral analysis. Their findings demonstrate that epistemic uncertainty sharply declines precisely at the onset of generalization, serving as an unsupervised diagnostic indicator. Further theoretical analysis shows that both delayed generalization and the peak in uncertainty arise from a shared spectral mechanism.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Game Theory and Economic Paradigms: Adversarial Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI

Abstract

In-context learning enables transformers to adapt to new tasks from a few examples at inference time, while grokking highlights that this generalization can emerge abruptly only after prolonged training. We study task generalization and grokking in in-context learning using a Bayesian perspective, asking what enables the delayed transition from memorization to generalization. Concretely, we consider modular arithmetic tasks in which a transformer must infer a latent linear function solely from in-context examples and analyze how predictive uncertainty evolves during training. We combine approximate Bayesian techniques to estimate the posterior distribution and we study how uncertainty behaves across training and under changes in task diversity, context length, and context noise. We find that epistemic uncertainty collapses sharply when the model groks, making uncertainty a practical label-free diagnostic of generalization in transformers. Additionally, we provide theoretical support with a simplified Bayesian linear model, showing that asymptotically both delayed generalization and uncertainty peaks arise from the same underlying spectral mechanism, which links grokking time to uncertainty dynamics.

Problem

• in-context learning
• grokking
• epistemic uncertainty
• delayed generalization
• Bayesian perspective

Innovation

• epistemic uncertainty
• in-context learning
• grokking
• Bayesian inference
• delayed generalization