A Bayesian Perspective on the Role of Epistemic Uncertainty for Delayed Generalization in In-Context Learning
For researchers studying in-context learning and grokking in transformers, this work provides a practical diagnostic (epistemic uncertainty) and theoretical understanding of delayed generalization.
The paper studies delayed generalization (grokking) in in-context learning using a Bayesian perspective, finding that epistemic uncertainty collapses sharply when the model groks, providing a label-free diagnostic of generalization. Theoretical analysis with a Bayesian linear model shows that both delayed generalization and uncertainty peaks arise from the same spectral mechanism.
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.