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
Apr 14, 2026

Distinct mechanisms underlying in-context learning in transformers

Modern distributed networks, notably transformers, acquire a remarkable ability (termed `in-context learning') to adapt their computation to input statistics, such that a fixed network can be applied to data from a broad range of systems. Here, we provide a complete mechanistic characterization of this behavior in transformers trained on a finite set SS of discrete Markov chains. The transformer displays four algorithmic phases, characterized by whether the network memorizes and generalizes, and whether it uses 1-point or 2-point statistics. We show that the four phases are implemented by multi-layer subcircuits that exemplify two qualitatively distinct mechanisms for implementing context-adaptive computations. Minimal models isolate the key features of both motifs. Memorization and generalization phases are delineated by two boundaries that depend on data diversity, K=∣S∣K = |S|. The first (K1∗K_1^\ast) is set by a kinetic competition between subcircuits and the second (K2∗K_2^\ast) is set by a representational bottleneck. A symmetry-constrained theory of a transformer's training dynamics explains the sharp transition from 1-point to 2-point generalization and identifies key features of the loss landscape that allow the network to generalize. Put together, we show that transformers develop distinct subcircuits to implement in-context learning and identify conditions that favor certain mechanisms over others.

Cole Gibson, Wenping Cui, Gautam Reddy

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

In-Context LearningTransformer Training DynamicsTransformer InterpretabilityNeural Network Training DynamicsMechanistic InterpretabilityStatistical Physics of Learning
AskPDF
Mar 19, 2026

Transformers Learn Robust In-Context Regression under Distributional Uncertainty

Recent work has shown that Transformers can perform in-context learning for linear regression under restrictive assumptions, including i.i.d. data, Gaussian noise, and Gaussian regression coefficients. However, real-world data often violate these assumptions: the distributions of inputs, noise, and coefficients are typically unknown, non-Gaussian, and may exhibit dependency across the prompt. This raises a fundamental question: can Transformers learn effectively in-context under realistic distributional uncertainty? We study in-context learning for noisy linear regression under a broad range of distributional shifts, including non-Gaussian coefficients, heavy-tailed noise, and non-i.i.d. prompts. We compare Transformers against classical baselines that are optimal or suboptimal under the corresponding maximum-likelihood criteria. Across all settings, Transformers consistently match or outperform these baselines, demonstrating robust in-context adaptation beyond classical estimators.

Hoang T. H. Cao, Hai D. V. Trinh, Tho Quan +1

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

In-Context LearningTransformerDistribution Shift RobustnessRobust Regression
AskPDF
May 15, 2026

Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets

Why does a Transformer that has memorized its training set wait thousands of steps before it generalizes? Existing accounts locate this delay in norm minimization, feature emergence, or the late discovery of sparse subnetworks. These explanations capture important parts of the transition, but ignore a constraint unique to attention-based models: if attention discards an informative token, no bounded downstream computation can recover it. We formalize attention as an implicit Bayesian posterior over the task dependency graph and prove that generalization requires two separable conditions: a familiar Goldilocks bound on MLP capacity, coinciding with norm-based theories of grokking, and a novel Bayesian structural condition requiring attention to place sufficient mass on every informative token. This decoupling explains delayed generalization as delayed structural inference. Early in training, the MLP memorizes through unaligned features, drives the cross-entropy loss near zero, and thereby starves attention of structural gradient. Weight decay must then erode memorization before the missing graph becomes learnable, yielding the known inverse-weight-decay delay, which we derive as a structural waiting time. We then prove that this explaining-away delay can be bypassed by a KL-based structural intervention, yielding an inverse-intervention-strength scaling law for the grokking time. Experiments on algorithmic sequence tasks isolate structure from capacity and show that this Bayesian ticket matches or outperforms lottery-ticket transfer.

Kai Hidajat, Solden Stoll, Joseph An

Department of Computer Science University of Washington Seattle, WA 98195

Neural Network GeneralizationLottery Ticket HypothesisTransformerSelf-AttentionAttention MechanismsGrokking
AskPDF