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

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.

• The paper shows that a sharp collapse in epistemic uncertainty accurately indicates the transition from memorization to generalization in transformers.
• It employs IVON and Last-Layer Laplace methods for Bayesian inference to decouple epistemic and aleatoric uncertainty in modular arithmetic tasks.
• The results reveal that tracking uncertainty via spectral analysis provides a practical, label-free diagnostic for robust OOD generalization and grokking onset.

Bayesian Epistemic Uncertainty as a Diagnostic for Delayed Generalization in In-Context Learning

Introduction and Problem Setting

This work provides a Bayesian framework to analyze the link between epistemic uncertainty and the phenomenon of grokking—an abrupt transition from memorization to generalization—in transformer models exhibiting in-context learning (ICL) (2604.12434). The study focuses on modular arithmetic tasks, where a transformer must infer a latent linear function solely from in-context examples, and investigates how predictive uncertainty—decomposed into epistemic and aleatoric components—changes across training and evaluation regimes. A central claim is that a sharp collapse in epistemic uncertainty tracks the grokking transition, making it a practical, label-free diagnostic of generalization. The paper further supports these empirical findings with rigorous results for a simplified Bayesian linear model, attributing both the delayed generalization and uncertainty peak to a shared spectral mechanism.

Experimental Design and Bayesian Methods

The task design involves modular arithmetic of the form $z = ax + by \pmod{p}$: the underlying “task vector” $(a, b)$ is not provided to the transformer, which must infer it from observed input-output triplets within the prompt. Four evaluation splits are constructed: ID Train (in-distribution tasks, in-distribution inputs), ID Val, OOD Train (out-of-distribution tasks, in-distribution inputs), and OOD Val (OOD tasks, OOD inputs). Only ID Train is used for training, allowing precise localization of model behavior (memorization vs generalization) as training data, task diversity, context size, and context noise are systematically varied.

For Bayesian inference, two effective approximations are employed:

• IVON: Improved Variational Online Newton, maintaining a Gaussian variational posterior over weights with Newton-like updates.
• Last-Layer Laplace: Posterior is Gaussian only over the last-layer weights, estimated post-hoc via Kronecker-factored Hessian approximation.

Both support predictive uncertainty quantification by Monte Carlo sampling from weight posteriors. The Bayesian total predictive uncertainty for each example decomposes as

$$\mathrm{EU}(\mathbf{x}) = \mathbb{H}\left[\frac{1}{S}\sum_{s=1}^{S} p(\mathbf{y}|\mathbf{x}, \theta^{(s)})\right] - \frac{1}{S} \sum_{s=1}^S \mathbb{H}\left[p(\mathbf{y}|\mathbf{x}, \theta^{(s)})\right]$$

where $\mathbb{H}$ denotes entropy, the first term is the total uncertainty, and the second is the expected aleatoric uncertainty.

Main Empirical Findings

Emergence of OOD Generalization and Accuracy/Uncertainty Dynamics

Increasing the number of pretraining tasks and training data fraction enables transformers to achieve high OOD Train and OOD Val accuracy, but only for large values (e.g., $T\geq256$ for tasks). Laplace demonstrates broader and more consistent OOD generalization than IVON. Specifically:

Figure 1: Both IVON and Laplace achieve high OOD Train accuracy for $T\geq256$ and high OOD Val accuracy only at high task counts and data fractions, with Laplace being more robust.

Decomposition of predictive uncertainty reveals that at low task diversity, epistemic uncertainty dominates; as task diversity increases, epistemic uncertainty sharply collapses at the grokking transition points for both methods, whereas aleatoric uncertainty remains elevated, indicating persistent label ambiguity regardless of the learnt hypothesis.

Figure 2: Epistemic uncertainty drops sharply at the transition to generalization for both IVON and Laplace, providing a label-free uncertainty-based signal of grokking.

This collapse in epistemic uncertainty is tightly aligned with, or precedes, the rise in OOD accuracy, providing an operational, label-agnostic diagnostic: a strong claim is established that epistemic uncertainty collapse demarcates the grokking phase transition, extending characterizations beyond accuracy trajectories.

Sensitivity to Context Size and Noise

Varying the number of in-context examples at test time monotonically increases accuracy and diminishes aleatoric uncertainty, in line with accumulating evidence. However, epistemic uncertainty exhibits a non-monotonic response: it increases with small context sizes (due to initial ambiguity), then decreases rapidly as sufficient context is accumulated.

Figure 3: Adding more in-context examples monotonically boosts accuracy and reduces aleatoric uncertainty, while epistemic uncertainty rises and then falls.

Introducing random noise to the context labels smoothly degrades accuracy and increases both aleatoric and epistemic uncertainty. Notably, Laplace’s epistemic uncertainty is more sensitive to context corruption than IVON’s, suggesting distinct robustness profiles in these approximate inference schemes.

Figure 4: Adding noise to context labels increases both uncertainty types and degrades accuracy, highlighting the sensitivity of epistemic uncertainty as a diagnostic.

Uncertainty Peaks and Delays in Simplified Models

A simplified Bayesian linear regression model—with analytical tractability—exhibits dynamics mirroring grokking: training loss declines before generalization error does, and epistemic uncertainty peaks around the transition from memorization to generalization. The time delay between training (empirical) and generalization (test) risk decay scales with the model aspect ratio parameter ($\lambda = d/n$), directly inherited by the uncertainty peak timings:

Figure 5: As the aspect ratio $\lambda$ increases, the delay between interpolation and generalization extends, with epistemic uncertainty peaks precisely tracking the transition.

Both accuracy and epistemic uncertainty inherit the same spectral underpinnings, dominated by the slowest mode of the sample covariance spectrum, explaining the alignment (with nonzero offset) between grokking time and the uncertainty peak.

Theoretical Analysis and Implications

The spectral analysis attributes the observed delay to the edge of the Marchenko-Pastur spectrum of the data covariance, which controls the late-stage convergence of both risk and uncertainty. The accuracy threshold for grokking corresponds (via a tolerance $\epsilon$) to a specific value of excess risk in the Bayesian linear model. Notably, the epistemic uncertainty peak, computed from the entropy of the success probability under $q_t(w)$, does not exactly coincide with the crossing of the accuracy threshold, but precedes it by a quantifiable offset determined by spectral characteristics.

This result substantiates the empirical observation: both label-based and label-free diagnostics (epistemic uncertainty) localize the delayed generalization transition, with uncertainty enabling early, unsupervised detection.

Broader Implications and Future Directions

These findings have major implications for training detection and analysis in overparameterized models, particularly transformers exhibiting in-context learning:

• Uncertainty as a Diagnostic: Epistemic uncertainty can be monitored in a fully unsupervised manner as a robust indicator of generalization onset, potentially replacing reliance on test labels for grokking detection.
• OOD Robustness Analysis: Comparative robustness of inference schemes (IVON vs Laplace) to context corruption is revealed via uncertainty decomposition, informing the design of uncertainty-aware architectures and training protocols.
• Refinement of Uncertainty Scaling Laws: Non-monotonic scaling of epistemic uncertainty with context size and noise challenges assumptions derived from standard Bayesian theory, motivating updated theoretical frameworks explicitly recognizing in-context adaptation.
• Theoretical Generalization: The demonstrated alignment of accuracy and uncertainty dynamics in simplified linear models suggests spectral analyses are central for understanding delayed generalization in more complex, nonlinear architectures.

Further research is required to validate these observations in naturalistic tasks and large-scale pretraining, extending the theoretical analysis beyond the tractable linear setting. The transferability of uncertainty-based diagnostics to other domains and their potential integration into active learning or early-stopping criteria are promising avenues.

Conclusion

This paper demonstrates that in transformers trained for in-context learning of modular arithmetic, epistemic uncertainty—estimated via scalable Bayesian approximations—collapses synchronously with the onset of generalization (grokking). This provides a practical, label-free probe of generalization applicable in both in-distribution and out-of-distribution regimes. Theoretical analysis reveals that both delayed generalization and uncertainty peaks are manifestations of a unified spectral mechanism governing Bayesian learning dynamics. The results support a broader adoption of uncertainty-based diagnostics in the study and deployment of deep learning models, and imply the potential universality of spectral mechanisms underlying abrupt transitions in overparameterized learning systems.

(2604.12434)