stat.MLApr 14, 2026

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

Authors: Abdessamed Qchohi, Simone Rossi

Organizations: EURECOM, France

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: Epistemic uncertainty collapses sharply at the grokking transition. On a decoder-only transformer trained on modular arithmetic tasks, the epistemic uncertainty collapses sharply before the generalization transition. This pattern follows the theoretical prediction for a simplified linear model, where the grokking transition corresponds to a spectral shift in the posterior distribution that causes a collapse in epistemic uncertainty.
Figure 2: Accuracy heatmap with IVON and Laplace . For OOD Train, both IVON and Laplace achieve high accuracy starting at T=256T=256 and T=512T=512 for any training data fraction. For OOD Val (the most challenging split), Laplace generalizes more broadly across configurations, while IVON achieves competitive performance only at high training data fractions.
Figure 3: Uncertainty decomposition with IVON and Laplace . Decomposition of predictive uncertainty into aleatoric and epistemic components across the pretraining hyperparameter grid.
Figure 4: Accuracy, epistemic and aleatoric uncertainty across context sizes . Increasing the number of in-context examples monotonically improves accuracy and decreases alepatoric uncertainty, while the epistemic uncertainty follows a non-monotonic behavior.
Figure 5: Accuracy vs. epistemic and aleatoric uncertainty across context noise levels . Corrupting in-context examples with random label noise degrades accuracy and raises both aleatoric and epistemic uncertainty.
Figure 6: Training and generalization dynamics of the simplified model. Increasing the aspect ratio λ\lambda delays generalization relative to training ( left ). The epistemic uncertainty peaks coincide with the transition from interpolation to generalization ( center left ), although the EU peak occurs before the accuracy threshold is exceeded ( center right ). Consistently, the EU peak delay grows rapidly with λ→1\lambda\to 1 , with the asymptotic formula of Theorem 5.1 providing a good proxy for the observed delay ( right ).
Figure 7: Empirical trajectory of accuracy and epistemic uncertainty. Even in the full transformer model, the EU peak occurs before the accuracy threshold is exceeded, similarly to the asymptotic behavior of the simplified model.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

QuantityValue
Arithmetic rulez=(ax+by) mod 29z=(ax+by)\bmod 29
Vocabulary size2929
Task dimension22
Input-pair dimension22
Input space size292=84129^{2}=841
Base in-distribution task countntaskn_{\mathrm{task}}
Appendix
Table 1: Task-generation and sequence-construction choices shared across the main experiments.
QuantityValue
ArchitectureDecoder-only transformer
Layers66
Attention heads44
Embedding width512512
Feed-forward width20482048
Vocabulary size2929
Appendix
Table 2: Shared architecture for MAP, IVON, and Laplace evaluations.
FieldMAP (AdamW)IVON
Peak learning rate1.5×10−41.5\times 10^{-4}0.50.5
Warmup steps1000100020002000
Initial learning-rate scale0.010.010.010.01
Final learning-rate scale0.10.10.10.1
ScheduleLinear warmup + cosine decayLinear warmup + cosine decay
Weight decay1110−610^{-6}
Appendix
Table 3: Optimization settings for the two trained model families.
FieldLaplace approximation
Base solutionFinal MAP model
Bayesian scopeLast layer only
Curvature structureKronecker-factored Hessian
Prior precisionEqual to MAP weight decay
Posterior samples1616 / 6464
Fit dataClean in-distribution training data
Appendix
Table 4: Laplace-specific settings.
More in this topic

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 402 papers

Explore similar work

CardsList

  1. Distinct mechanisms underlying in-context learning in transformers

    Apr 14, 2026Cole Gibson, Wenping Cui, Gautam ReddyIn-Context LearningTransformer Training Dynamics

    Lewis-Sigler Institute for Integrative Genomics, Princeton University, Princeton, NJ 08544, USA · Department of Physics, Princeton University, Princeton, NJ 08544, USA

  2. Transformers Learn Robust In-Context Regression under Distributional Uncertainty

    Mar 19, 2026Hoang T. H. Cao, Hai D. V. Trinh, Tho Quan +1In-Context LearningTransformer

    Faculty of Computer Science and Engineering, Ho Chi Minh University of Technology (HCMUT), Vietnam National University Ho Chi Minh City (VNU-HCM), Vietnam

  3. Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets

    May 15, 2026Kai Hidajat, Solden Stoll, Joseph AnNeural Network GeneralizationLottery Ticket Hypothesis

    Department of Computer Science University of Washington Seattle, WA 98195