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

TL;DR

This paper investigates how epistemic uncertainty relates to delayed generalization in in-context learning, revealing that uncertainty collapse signals model grokking and linking this to spectral mechanisms.

Contribution

It introduces a Bayesian framework to analyze uncertainty dynamics in transformers, providing theoretical and empirical insights into delayed generalization and grokking phenomena.

Findings

Epistemic uncertainty sharply collapses at grokking.

Uncertainty dynamics are linked to spectral properties in a simplified model.

Uncertainty serves as a label-free diagnostic of generalization.

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.