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

A transformer can memorize training examples long before it learns to infer an unfamiliar task from a prompt. We study this delayed generalization, called grokking, through disagreement among approximate posterior weight samples. On modular arithmetic tasks, epistemic uncertainty rises and then collapses around the transition to generalization. A tractable Bayesian linear model links the delay in accuracy and uncertainty to slowly learned spectral modes, while distinguishing their transition times.

Presented at the ProbML 2026 Workshop Track as A Bayesian perspective on learning to grok in-context examples.

Task inference and input memorization

Each task hides a coefficient pair $(a,b)$ and supplies numerical examples of a modular linear function:

$$\begin{aligned} z_i &= (a x_i+b y_i)\bmod 29,\\ \text{prompt} &= (x_1,y_1,z_1),\ldots,(x_{t-1},y_{t-1},z_{t-1}),(x_t,y_t,?). \end{aligned}$$

All values belong to $\{0,\ldots,28\}$; the model must infer the coefficients from preceding triples to predict the final answer. We independently withhold coefficient pairs and input pairs. In-distribution (ID) validation tests unseen inputs on seen tasks. Out-of-distribution (OOD) Train tests unseen tasks on seen inputs, and OOD validation withholds both. This separation makes task inference and input memorization independently observable.

We use a six-layer decoder-only transformer with width 512, four attention heads, rotary positional embeddings, and tied embeddings. Training supervises answer positions in sequences of 32 triples. The experiments report medians and 95% quantile intervals over four seeds.

Disagreement over predictions

Improved Variational Online Newton (IVON) learns a diagonal Gaussian over all weights during training. Last-Layer Laplace fits a Gaussian to the final layer of an AdamW-trained model, using Kronecker-factored curvature and fixed preceding layers. For a prompt $\boldsymbol{x}$, predictions from $S$ posterior samples $\boldsymbol{\theta}^{(s)}$ define

$$\begin{aligned} \bar p(\boldsymbol{y}\mid\boldsymbol{x}) &=\frac1S\sum_{s=1}^{S}p(\boldsymbol{y}\mid\boldsymbol{x},\boldsymbol{\theta}^{(s)}),\\ \mathrm{EU}(\boldsymbol{x}) &=\mathbb H[\bar p] -\underbrace{\frac1S\sum_{s=1}^{S}\mathbb H[p(\boldsymbol{y}\mid\boldsymbol{x},\boldsymbol{\theta}^{(s)})]}_{\mathrm{AU}(\boldsymbol{x})}. \end{aligned}$$

Here $\mathbb H$ denotes entropy. Total uncertainty measures the spread of the averaged prediction; the aleatoric term averages individual entropies; their difference measures disagreement across weights. Computing these quantities requires the prompt and sampled predictions, with no query label. Their values depend on the posterior approximation.

Our uncertainty-scaling study uses the same entropy decomposition to examine how predictive disagreement changes with training-set size.

Task diversity changes generalization and uncertainty together

We sweep effective task counts $T\in\{32,64,128,256,512\}$ and training-input fractions $\alpha\in\{0.3,0.4,0.5,0.6,0.7,0.8\}$. Both methods achieve high accuracy on unseen tasks with seen inputs at $T=256$ and $512$. Laplace generalizes across a broader range of configurations when both tasks and inputs are unseen; IVON performs competitively there at high training fractions.

The inference-time sweeps fix 64 base tasks, expanded to 256 by parallelogram augmentation, and $\alpha=0.8$. Increasing context from one to 32 examples improves accuracy overall and reduces aleatoric uncertainty. Epistemic uncertainty initially rises before falling, exposing a non-monotonic response to additional context.

Replacing context answers with uniformly random tokens, with probability $p_{\mathrm{noise}}$, lowers accuracy and raises both uncertainty components. Inputs and evaluation targets remain clean. Laplace's epistemic uncertainty rises more strongly than IVON's, illustrating method-dependent sensitivity to misleading prompts.

Slow spectral modes explain a delay in a linear model

For Gaussian linear regression, let $\boldsymbol{X}\in\mathbb R^{n\times d}$ contain the training inputs, $\sigma^2$ be likelihood variance, and $\tau^2$ prior variance. Gaussian variational inference has an exact Bures-Wasserstein flow:

$$\begin{aligned} q_t(\boldsymbol{w})&=\mathcal N(\boldsymbol{\mu}_t,\boldsymbol{\Sigma}_t),\\ \dot{\boldsymbol{\mu}}_t&=\boldsymbol{b}-\boldsymbol{A}\boldsymbol{\mu}_t,\\ \dot{\boldsymbol{\Sigma}}_t&=2\boldsymbol{I}-\boldsymbol{A}\boldsymbol{\Sigma}_t-\boldsymbol{\Sigma}_t\boldsymbol{A},\\ \boldsymbol{A}&=\sigma^{-2}\boldsymbol{X}^{\top}\boldsymbol{X}+\tau^{-2}\boldsymbol{I}, \qquad\boldsymbol{b}=\sigma^{-2}\boldsymbol{X}^{\top}\boldsymbol{y}. \end{aligned}$$

Here $\boldsymbol{\mu}_t$ and $\boldsymbol{\Sigma}_t$ are the weight mean and covariance, and $\boldsymbol{y}$ contains responses. The same precision operator drives both dynamics. Training risk downweights small sample-covariance eigenvalues, while population risk retains their contribution. Those slow modes therefore delay generalization.

Under isotropic Gaussian design, matched interpolation, self-averaging initialization, and $n,d\to\infty$ with $\lambda=d/n<1$, the lower Marchenko-Pastur edge controls late-time excess risk $\widetilde{\mathcal L}(t)$. Matched interpolation assumes the stationary mean equals the teacher weights and the stationary training residual vanishes. Thresholding squared residuals at tolerance $\epsilon$ gives a binary success event whose entropy is entirely epistemic. Using the effective Gaussian residual approximation,

$$\begin{aligned} \mathrm{EU}_{\mathrm{th}}(t) &\approx h_2\!\left(\operatorname{erf}\sqrt{\frac{\epsilon}{4\widetilde{\mathcal L}(t)}}\right),\\ \Delta t_{\mathrm{EU,peak}}(\lambda) &\simeq\frac{\log\!\left(\frac{1}{1-\sqrt\lambda}\right)} {\frac{n}{\sigma^2}(1-\sqrt\lambda)^2+\frac1{\tau^2}}. \end{aligned}$$

Here $h_2$ is binary entropy and $\Delta t_{\mathrm{EU,peak}}$ separates population and training peaks, provided both occur in the late-time regime. The peak occurs at 50% thresholded success, before the 95% accuracy criterion; both delays share the leading spectral scaling.