A Bayesian Perspective on the Role of Epistemic Uncertainty for Delayed Generalization in In-Context Learning Abdessamed Qchohi [email protected] EURECOM, France
Simone Rossi [email protected] EURECOM, France
arXiv:2604.12434v1 [stat.ML] 14 Apr 2026
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.
1
Introduction
In-Context Learning (ICL) has emerged as a powerful paradigm in which Large Language Models (LLMs) can learn to perform new tasks by conditioning on a few examples provided in the input prompt (Brown et al., 2020). However, despite its empirical success, the underlying mechanisms that enable LLMs to effectively learn from in-context examples remain poorly understood. Recent studies have begun to explore the theoretical foundations of ICL, proposing various hypotheses ranging from implicit Bayesian inference (Xie et al., 2022; Akyurek et al., 2023) to meta-learning frameworks (Von Oswald et al., 2023). Yet, a comprehensive Bayesian perspective on how LLMs process and learn from in-context examples is still lacking. Furthermore, during training of LLMs, models often exhibit a phenomenon known as grokking, where they suddenly transition from poor to near-perfect performance on a task after extended training, despite having already fit the training data (Power et al., 2022). In this work, we aim to study what enables grokking in ICL from a Bayesian perspective and how learning dynamics (and uncertainty measures) evolve during training. Bayesian Deep Learning (BDL) provides a principled framework for modeling uncertainty in deep learning models by treating model parameters as random variables and inferring their posterior distributions given the data (Gal & Ghahramani, 2016; Blundell et al., 2015; Lakshminarayanan et al., 2017). While exact inference is intractable in large-scale models like LLMs, various approximate inference techniques, such as Variational Inference (VI) and Laplace approximations, have been successfully applied to capture uncertainty in deep learning (Papamarkou et al., 2024). By treating the model in a Bayesian framework, we experimentally analyze how uncertainty quantification and posterior distributions over model parameters change as the model groks with the in-context examples. Following the classic setting, we approch this phenomenon by analyzing the training dynamics of transformers on modular arithmetic tasks (Power et al., 2022). Additionally, we provide a theoretical support for our 1
EU (Generalization)
EU (OOD Val)
Acc. (Generalization)
Acc. (OOD Val)
empirical findings by analyzing the uncertainty dynamics in a simplified model, showing how the transition to generalization corresponds to a collapse in epistemic uncertainty. (a) Real model (b) Theory Our contributions are as follows: (1) We 0.08 1.00 0.8 EU peak study grokking in ICL in a Bayesian frame0.6 EU peak 0.06 0.75 0.6 work, using modular arithmetic tasks to an0.4 0.04 0.50 alyze how predictive uncertainty evolves 0.4 0.2 as a transformer learns to infer a latent task 0.02 0.25 0.2 from in-context examples. (2) We show em0.0 0.00 0.00 0.0 10,000 100,000 10 10 10 pirically that a sharp collapse in epistemic Step Time / iteration t uncertainty coincides with the grokking transition, making it a label-free diagnostic Figure 1: Epistemic uncertainty collapses of generalization. We further compare Im- sharply at the grokking transition. On a proved Variational Online Newton (IVON) decoder-only transformer trained on moduand Last-Layer Laplace across training, task lar arithmetic tasks, the epistemic uncertainty diversity, context length, and context noise, collapses sharply before the generalization highlighting both shared uncertainty pat- transition. This pattern follows the theoretterns and method-specific differences in out- ical prediction for a simplified linear model, of-distribution (OOD) generalization. (3) where the grokking transition corresponds to We provide theoretical support with a sim- a spectral shift in the posterior distribution plified Bayesian linear model, showing that that causes a collapse in epistemic uncertainty. delayed generalization and delayed epistemic-uncertainty peaks arise from the same underlying spectral mechanism, which links grokking time to uncertainty dynamics. A snapshot of our main empirical and theoretical findings is shown in Figure 1. −8
−5
−2
Related Works. ICL has been studied both empirically and theoretically since its emergence in large language models (Brown et al., 2020). A common theoretical view is that transformers can implement learning or estimation procedures in their forward pass, including implicit Bayesian inference, ridge regression, and gradient-based adaptation (Xie et al., 2022; Akyurek et al., 2023; Von Oswald et al., 2023; Bai et al., 2023; Xie et al., 2025; Zhang et al., 2025). Complementary work examines reliability under prompting: Zhang et al. (2024) study calibration across shot counts, Ling et al. (2024) decompose ICL uncertainty into epistemic and aleatoric terms, and Wang et al. (2025) analyze uncertainty in long-context prompting. A Bayesian framing provides a natural language for connecting these threads: pre-training can be interpreted as learning an implicit prior (or an implicit family of priors) over tasks/concepts, while ICL corresponds to approximate Bayesian inference (or Bayesian model averaging) over latent task variables given the prompt (Xie et al., 2022; Zhang et al., 2025). In parallel, BDL treats uncertainty as first-class, via priors and posteriors over parameters and predictive distributions as marginalization, and offers tools that make uncertainty trajectories during training and inference measurable (e.g. Graves, 2011; Blundell et al., 2015; Gal & Ghahramani, 2016; Lakshminarayanan et al., 2017; Maddox et al., 2019; Ritter et al., 2018). On the grokking side, Power et al. (2022) introduced the canonical delayed-generalization phenomenon on modular arithmetic. Follow-up work characterized its learning phases and representation dynamics (Liu et al., 2022), explained aspects of modular addition through the transition out of the Neural Tangent Kernel (NTK) regime (Mohamadi et al., 2024), and proposed information-theoretic or structural progress measures that can anticipate grokking (Clauw et al., 2024; Golechha, 2024). Closest to our setup, He et al. (2024) extend the classic task family to modular linear functions and study both in-distribution and out-of-distribution learning. While we build on the same experimental set up, our focus is on the Bayesian learning dynamics and uncertainty trajectories during grokking, rather than characterizing the phases of learning or the representational changes. Notably, we identify a the condition under which the model can generalize in-domain and out-of-domain, and we analyze how the model’s uncertainty evolves as it transitions from overfitting to generalization, providing a novel perspective on the grokking phenomenon. A more detailed review of the related work is available in Appendix A. 2
2
Preliminaries
Bayesian inference. Rather than treating model parameters as fixed values, Bayesian methods model them as random variables. Given a dataset D and model parameters θ, Bayes’ rule gives the posterior: p(θ | D) =
p(D | θ) p(θ) p(D)
where p(θ ) is the prior encoding beliefs about the parameters before observing data, p(D | θ) is the likelihood describingRhow probable the observed data is under a particular choice of parameters, and p(D) = p(D | θ) p(θ) dθ is the marginal likelihood, which integrates out the parameters and acts as a normalizing constant ensuring the posterior is a valid probability distribution. Predictions for a new input x∗ are obtained by marginalizing over the posterior: Z p(y∗ | x∗ , D) =
p(y∗ | x∗ , θ) p(θ | D) dθ
(1)
This predictive distribution captures uncertainty in the model parameters, which is relevant when data is limited or when the model is evaluated on out-of-distribution inputs. Exact inference is intractable for neural networks since the posterior is high-dimensional and the marginal likelihood has no closed form. We therefore need to rely on approximate inference methods. While the toolbox of approximate inference methods is large and rapidly evolving, we focus on two complementary approaches that are well-suited for our setting: variational inference and Laplace approximation. IVON (Shen et al., 2024) integrates uncertainty estimation into training by maintaining a Gaussian variational posterior over all weights, updated via Newton-like steps at a computational cost comparable to Adam. The Last-Layer Laplace Approximation (Daxberger et al., 2021) is a post-hoc method that fits a Gaussian approximation to the posterior of the final linear layer after standard MAP training, using a Kronecker-factored Hessian approximation, leaving all other parameters fixed. For both methods, predictions are obtained by averaging over S weight samples drawn from q(θ). Additional details on these methods are available in Appendix B. Uncertainty decomposition. Both approximate inference methods yield a distribution over predictions, enabling a principled decomposition of predictive uncertainty (Wimmer et al., 2023; Rosso et al., 2025; Hüllermeier & Waegeman, 2021). Given S weight samples θ(s) from the approximate posterior, total uncertainty (TU) is the entropy of the mean prediction; aleatoric uncertainty (AU) is the mean entropy of individual predictions and captures irreducible label noise; and epistemic uncertainty (EU) is their difference: " # h i 1 S 1 S (s) (s) EU( x) = H p ( y | x, θ ) − H p ( y | x, θ ) S s∑ S s∑ =1 =1 {z } | AU( x)
Epistemic uncertainty reflects the model’s lack of knowledge about the task and is expected to decrease as the model acquires more information. In the grokking setting, we use EU as a probe for the generalization transition.
3
In-Context Learning, Grokking and Uncertainty
In this section, we describe the main experimental results on the relationship between grokking and uncertainty in in-context learning. We first introduce the experimental setup, including the task distribution, model architecture, and evaluation protocol. We then present our main empirical findings on how accuracy and epistemic uncertainty behave. Finally, we analyze how task diversity, context length, and context noise affect both generalization and uncertainty dynamics. 3
3.1
Experimental Setup
Tasks. We follow He et al. (2024) and focus on linear modular arithmetic tasks of the form z = ax + by (mod p), where ( a, b) is the task vector and ( x, y) is the input vector, with all values in {0, . . . , p − 1}. Modular arithmetic provides a controlled setting for studying grokking, as the tasks have a clear underlying structure that the model can either memorize or learn to generalize. The task vector is never shown to the model, which receives only input-output triples ( x, y, z) and must infer the underlying task from the in-context examples. Each equation is tokenized into three numerical. An example of an input sequence is: ( x1 , y1 , z1 ), . . . , ( xnctx , ynctx , znctx ), ( xt , yt , ?) During training, the model predicts every third token zi as well as the final target zt , with the cross-entropy loss computed only on these positions. At evaluation time, we measure accuracy on zt only, which directly tests whether the model has inferred the correct task from context. We partition the set of task vectors and input examples into seen and unseen subsets, Tid , Tood and Xtrain , Xtest , respectively. Crossing these two partitions gives four splits: ID Train (Tid , Xtrain ), used for training; ID Val (Tid , Xtest ), which tests generalization to unseen inputs on seen tasks; OOD Train (Tood , Xtrain ), which tests whether the model adapts to unseen tasks on seen inputs; and OOD Val (Tood , Xtest ), the hardest split, which tests full out-of-distribution generalization to both unseen tasks and unseen inputs. We train exclusively on ID Train and use the remaining three splits to evaluate generalization. This construction lets us track whether the model memorizes, generalizes in-distribution, or generalizes out-of-distribution, and at which point grokking occurs in each regime. Model architecture and training. We use a decoder-only Transformer with Rotary Positional Embeddings (RoPE) (Su et al., 2024) with d = 6 layers, embedding dimension 512, 4 attention heads, and an MLP widening factor of 4. We tie input embedding and output projection weights and apply weight decay for regularization. Task vectors are sampled via a parallelogram rule, varying one component at a time so the model can build on previously learned structure, and each batch is constructed so that all tasks in Tid share the same input sequences, limiting inter-task interference. Unless otherwise specified, we train four random seeds and we report medians with 95% quantile intervals across seeds for all metrics and configurations. Detailed descriptions of the training procedure, hyperparameters, and approximate inference methods are provided in Appendix D.
4
Empirical Results and Analysis
Grokking across task diversity. To study how grokking emerges as a function of task diversity and training data, we train the model across a grid of configurations varying the number of pre-training tasks T ∈ {32, 64, 128, 256, 512} and the fraction of training data α ∈ {0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, running four random seeds per configuration. We measure accuracy on Val, OOD Train, and OOD Val, and compare IVON and Laplace to assess whether the choice of optimizer affects the generalization landscape. Figure 2 shows the results. Both IVON and Laplace achieve high OOD Train accuracy at T = 256 and T = 512. For OOD Val, Laplace exhibits stronger and more consistent generalization across task counts and data fractions, whereas IVON reaches competitive OOD Val accuracy only at the highest training data fractions. OOD Val generalization emerges only at the highest task counts for both methods and requires careful tuning of training data fraction. Epistemic uncertainty tracks grokking. To test whether epistemic uncertainty tracks the grokking transition, we use IVON and Laplace to decompose predictive uncertainty into epistemic and aleatoric components across the same hyperparameter grid. Figure 3 shows the results in each evaluation set. Both methods show that at low task diversity, total uncertainty is dominated by the epistemic component across all splits. As T increases, epistemic uncertainty drops sharply at T = 256 and T = 512 for both IVON and Laplace, while aleatoric uncertainty remains high. This reduction closely tracks the generalization transition in the accuracy heatmaps, providing a principled uncertainty-based signature of grokking consistent across approximate inference methods. This is an important finding: as 4
IVON OOD Val 1.0 0.5
Accuracy
Train %
OOD Train
30 40 50 60 70 80
Val
0.0 32
64
128
256
512
32
64
Tasks
128
256
512
32
64
Tasks
128
256
512
Tasks
Laplace OOD Val 1.0 0.5
Accuracy
Train %
OOD Train
30 40 50 60 70 80
Val
0.0 32
64
128
256
512
32
64
Tasks
128
256
512
32
64
Tasks
128
256
512
Tasks
Figure 2: Accuracy heatmap with IVON and Laplace. For OOD Train, both IVON and Laplace achieve high accuracy starting at T = 256 and T = 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.
Uncertainty (nats)
IVON EU + AU on Val
EU + AU on OOD Train
EU + AU on OOD Val EU α = 0.3 EU α = 0.4 EU α = 0.5 EU α = 0.6 EU α = 0.7 EU α = 0.8
2
AU α = 0.3 AU α = 0.4 AU α = 0.5 AU α = 0.6 AU α = 0.7 AU α = 0.8
1 0
32
64
128
256
512
32
64
T (tasks)
128
256
512
32
64
T (tasks)
128
256
512
T (tasks)
Uncertainty (nats)
Laplace EU + AU on Val
EU + AU on OOD Train
EU + AU on OOD Val
3
EU α = 0.3 EU α = 0.4 EU α = 0.5 EU α = 0.6 EU α = 0.7 EU α = 0.8
2
AU α = 0.3 AU α = 0.4 AU α = 0.5 AU α = 0.6 AU α = 0.7 AU α = 0.8
1 0
32
64
128
T (tasks)
256
512
32
64
128
T (tasks)
256
512
32
64
128
256
512
T (tasks)
Figure 3: Uncertainty decomposition with IVON and Laplace. Decomposition of predictive uncertainty into aleatoric and epistemic components across the pretraining hyperparameter grid. of this writing, grokking has been primarily characterized in terms of accuracy dynamics (which implies the need for labels to detect it), while here we show that it also has a clear signature in the model’s posterior uncertainty, which can be measured without access to labels and provides a more direct window into the model’s learning dynamics. Effect of context length on generalization. To study how the number of in-context examples affects the model’s ability to infer the underlying task, we take the models trained with nctx = 32 and evaluate them by varying the number of in-context examples from 1 to 32 at inference time, without any additional training. The task vector ( a, b) is never provided to the model, which must infer it solely from the input-output triples in the context. We compare IVON and Laplace on Val, OOD Train, and OOD Val. Figure 4 shows accuracy, 5
Val
OOD Train
OOD Val
Accuracy
1.0 0.5 IVON Laplace 0.0
EU
100
10−2
AU
100 10−1 10−2 1
8
16
24
32
1
8
Context size
16
24
32
1
8
Context size
16
24
32
Context size
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. Val
OOD Train
Accuracy
1.0
OOD Val
IVON Laplace
0.5
EU
100 10−1 10−2
AU
100 10−1
0.0
0.2
0.4
0.6
Noise probability
0.8
1.0
0.0
0.2
0.4
0.6
Noise probability
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Noise probability
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. epistemic uncertainty, and aleatoric uncertainty as a function of context size. Accuracy improves monotonically with context size for both IVON and Laplace across all splits, confirming that more examples indeed help the model infer the task (and thus generalize better on OOD splits). As expected, the aleatoric uncertainty decreases with more context and converges to a floor, reflecting the reduced ambiguity in the task given more evidence. The behavior of epistemic uncertainty is more complex: it increases at small context sizes, likely due to the model’s initial uncertainty about the task, then decreases sharply as more examples are added. Notably, this behavior is present in both IVON and Laplace, albeit with different magnitudes, suggesting that it isn’t an artifact of the approximate inference method but rather a property of the model’s (meta-)learning dynamics as it processes in-context evidence. Also, note that this evidence does not contradict the known fact that uncertainty scales monotonically with data in Bayesian inference (see, e.g., Rosso et al., 2025), but it rather highlights that the current scaling laws of uncertainty should be refined to account for the effect of test-time adaptation via in-context learning. Robustness to context noise. As a final test of the model’s uncertainty estimates, we evaluate how they respond to noise in the in-context examples. For this experiment, we corrupt the context by replacing each label zi in the context triples with a random token 6
sampled uniformly from {0, . . . , p − 1} with probability pnoise ∈ [0, 1]. The input tokens x and y are never corrupted, and the final query triple is always clean. As we did before, we evaluate both IVON and Laplace on Val, OOD Train, and OOD Val, measuring accuracy, aleatoric uncertainty, and epistemic uncertainty as a function of the noise level (Figure 5). Accuracy degrades smoothly as noise increases for both methods and all splits. Aleatoric uncertainty rises proportionally to the noise level, correctly capturing the increased label ambiguity in the context. Epistemic uncertainty also increases with noise, but more markedly for Laplace than for IVON, suggesting that the variational posterior learned by IVON is more robust to corrupted context. This is slightly surprising, as one might expect that corrupted context would primarily increase epistemic uncertainty, since it makes the task harder to infer, while aleatoric uncertainty should remain stable as it reflects irreducible noise in the data rather than lack of context. This demonstrates that recent analysis on disentangling aleatoric and epistemic uncertainty (e.g., Mucsányi et al., 2024; Kirchhof et al., 2025) do not necessarily apply to the in-context learning setting, and that further work is needed to understand how these uncertainty components interact when the model is performing test-time adaptation.
5
Why Epistemic Uncertainty Can Also Define Grokking Time
Our empirical results suggest that grokking is visible not only in accuracy curves but also in epistemic uncertainty. Making this claim rigorous for the full transformer on modular arithmetic is currently out of reach, so in this section we drastically simplify the problem. Following the linear-regression setup of Levi et al. (2024), we replace the sequence model by Bayesian linear regression. Although this is a very simple model, Levi et al. (2024) show that it already captures the key phenomenon of delayed generalization. Here we show that it also captures the key phenomenon of delayed EU peaks, when the model is trained with Gaussian variational inference. This abstraction is clearly much simpler than the model studied in the rest of the paper, but it gives fully tractable dynamics and still matches the qualitative empirical picture: training error decays before generalization error, the epistemic uncertainty is non-monotone, and the generalization-side transition is delayed. While full derivations are deferred to Appendix C, we sketch the main results here. Simplified problem. Let X ∈ Rn×d be the design matrix, with xi ∼ N (0, I d ), and y ∈ Rn be the target vector, generated from w⋆ with isotropic Gaussian prior. Consider the linearGaussian model p(y | w) = N (y; Xw, σ2 I n ),
p(w) = N (w; 0, τ 2 I d ).
(2)
We study the regime n, d → ∞ with fixed aspect ratio λ = d/n < 1. Our goal is to analyze the dynamics of Gaussian variational inference in this model, and to understand how the risk and uncertainty evolve along the variational trajectory. Following the results of Lambert et al. (2022), since the model is linear and the variational family is Gaussian, the gradient flows on the Bures–Wasserstein space of Gaussian measures qt (w) = N (w; µt , Σt ) can be written in closed form as a system of ODEs: µ̇t = b − Aµt ,
(3)
Σ̇t = 2I d − AΣt − Σt A.
(4)
with A = σ−2 X ⊤ X + τ −2 I d and b = σ−2 X ⊤ y. This system has exact solutions, as shown in the Appendix. Along the Gaussian variational trajectory, we are interested in the expected risks
Lemp (t) = Eqt
1 n ℓ(yi , xi⊤ w) n i∑ =1
and
Lgen (t) = Eqt (w) p(x,y) ℓ(yi , xi⊤ w)
(5)
where ℓ is the squared loss, i.e. ℓ(y, ŷ) = 12 (y − ŷ)2 . We can also define the excess risks, which are the differences between the risks at time t and their limiting values as t → ∞. 7
0.0
10−10 10−8 10−6 10−4 10−2
Time / iteration t
0.6 0.4 0.2 0.0
Generalization milestones
10−10 10−8 10−6 10−4 10−2
10−5 10−6
10−4
Numerical MP delay Asymptotic late-time formula
10−6 10−8
10−7 0.00
0.25
0.50
0.75
Aspect ratio λ
Time / iteration t
EU peak delay
95% acc. time EU peak time
∆tEU,peak
Empirical General.
EU behavior during training Time / iteration t
Accuracy
0.5
λ = 0.20 λ = 0.60 λ = 0.99
EU (nats)
Accuracy during training 1.0
1.00
0.00
0.25
0.50
0.75
1.00
Aspect ratio λ
Figure 6: Training and generalization dynamics of the simplified model. Increasing the aspect ratio λ 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, with the asymptotic formula of Theorem 5.1 providing a good proxy for the observed delay (right). Risk dynamics and spectral decomposition. In the well-specified model regime, the fixed point µ∞ coincides with the exact parameters w⋆ and the population covariance is isotropic, so the relevant objects are the mean gap δµ0 ≡ µ0 − µ∞ and the covariance gap ∆Σ0 ≡ Σ0 − Σ∞ . Writing Et ≡ e− At , the excess empirical and generalization risks are 1 1 ⊤ Leemp (t) = δµ0 Et X ⊤ X Et δµ0 + Tr X ⊤ X Et ∆Σ0 Et , (6) 2n 2n 1 1 Legen (t) = δµ0⊤ E2t δµ0 + Tr( Et ∆Σ0 Et ). (7) 2 2 The first term tracks the residual bias of the variational mean, while the second tracks the residual uncertainty from the covariance. Crucially, both risks are driven by the same operator A, but they are read out in two different quantities: the empirical risk uses the b = X ⊤ X/n, whereas the population risk uses the identity. In the training covariance Σ b converge in distribution high-dimensional limit, the eigenvalues of the sample covariance Σ to the Marchenko-Pastur law MP(λ), and the excess risks can be expressed as expectations over this distribution: 1 n Leemp (t) ≈ C0 Eν∼MP(λ) ν exp −2 ν + t , (8) σ2 τ2 n 1 Legen (t) ≈ C0 Eν∼MP(λ) exp −2 ν + t . (9) σ2 τ2 where C0 collects the initial mean and covariance gaps (details in the Appendix). Under the assumption of high-dimensional self-averaging regime, √ 2 the late-time behavior is controlled by the lower Marchenko-Pastur edge ν− = (1 − λ) : √ 2 n 1 −3/2 −κt e Lemp,gen (t) ∼ Kemp,gen t e , κ=2 (1 − λ ) + 2 , (10) σ2 τ √ with ratio Kgen /Kemp = (1 − λ)−2 . The common exponential rate κ reflects the shared edge mode, while the different factors is due to the extra ν in the empirical risk and it is responsible for the train–test delay. This asymptotic mismatch is the only ingredient needed for the uncertainty result below. Threshold-model uncertainty and delayed EU peaks. Following the construction used by Levi et al. (2024), we turn regression into a binary success event by choosing a tolerance ϵ > 0 and declaring a prediction correct whenever the expected squared residual is smaller than ϵ. For a datapoint ( x, y), define the thresholded success variable Zϵ ( x, y; w) = 1{(y − x⊤ w)2 ≤ ϵ}. We can define the empirical and generalization accuracies Aemp and Agen as the expected value of Zϵ under the variational distribution and the data distribution (see the Appendix for details). By analogy with the accuracy threshold defined in Levi et al. (2024), we can define the grokking time as the time at which the accuracy curve crosses a certain threshold, which is equivalent as the time at which the excess risk crosses a certain 8
e tgrok ) ≈ ϵ/8. Now our current goal is to analyze threshold; for 95% accuracy we have L( the epistemic uncertainty of this model. Lemma 5.1 (Epistemic uncertainty of the classification-equivalent model). For fixed ( x, y), the residual under qt is Gaussian and the thresholded success event is a Bernoulli random variable. Hence, its epistemic uncertainty is exactly the binary entropy of its success probability, which in the e t ): n, d → ∞ limit can be written in terms of the excess risk L( !! s ϵ EU(t) ≈ h2 erf , (11) e t) 4L( where h2 ( p) = − p log p − (1 − p) log(1 − p) is the binary entropy. It is a non-monotone function of the excess risk, and its peak occurs when ϵ e tpeak ) = L( ≈ 1.10 ϵ, (12) −1 4(erf (1/2))2 i.e., when the thresholded success event is maximally ambiguous. By reversing this relation, we can also derive the tpeak as a function of ϵ and the risk dynamics. Remark 5.1. Note that the peak of the epistemic uncertainty does not occur at the same time as the accuracy threshold, but rather before it, when the excess risk is still around 1.10 ϵ (instead of ϵ/8); this offset is illustrated in the center-right panel of Figure 6. Finally, we can analyze the delay between the EU peaks on train and test, which is a proxy for the grokking time defined through epistemic uncertainty. Theorem 5.1 (Delay between threshold-uncertainty peaks). Assume the construction of Lemma 5.1 and the asymptotics described above. If the empirical and generalization uncertainty curves peak in this regime, then the delay between the two EU peaks satisfies 1√ log √1− λ 1 . ∆tEU,peak (λ) = tpeak,gen − tpeak,emp ≃ n (13) 2+ ( 1 − λ ) 2 2 σ τ This scaling is also visible in Figure 6, where the separation between the empirical and generalization EU peaks grows markedly as λ → 1, and it follows the asymptotic prediction.
6
EU peak tracks generalization 0.8
0.10
Accuracy EU
0.08
0.6
0.06
0.4
0.04
0.2
0.02
0.0 1,000
10,000
100,000
EU (OOD Val)
Accuracy (OOD Val)
Takeaway. Theorem 5.1 is the main theoretical reason to say that grokking time can also be defined through epistemic uncertainty. The point is not that the EU peak occurs at the same absolute time as an accuracy threshold (indeed, the EU peak is happing earlier than the accuracy threshold), but rather that both observables inherit the same train-generalization delay from the same slow edge-of-spectrum modes. In this simplified model, accuracy and epistemic uncertainty are therefore two different probes of the same underlying mechanism. This is exactly the qualitative behavior we observe in the empirical analysis. Indeed, even in the full transformer model on the modulo arithmetic dataset, the EU peak occurs before the accuracy threshold is exceeded, showing that the theoretical picture of the simplified model is compatible with the empirical behavior (Figure 7).
0.00
Step
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.
Concluding Remarks
We presented an uncertainty quantification perspective on delayed generalization in incontext learning and showed that epistemic uncertainty provides a sharp, label-free signature of grokking. In a controlled modular linear arithmetic setting with explicit indistribution and out-of-distribution splits, we found that grokking coincides with a sudden 9
collapse in epistemic uncertainty, and that this behavior remains informative under changes in task diversity, context length, and context noise. For both IVON and Last-Layer Laplace, this uncertainty signal is robust, while still revealing meaningful differences in out-ofdistribution behavior. We further supported these findings with theorethical results on a simplified Bayesian linear model, showing that delayed generalization and uncertainty peaks are driven by the same underlying spectral mechanism. To conclude, we want to emphasize that while our findings provide a novel perspective on the grokking phenomenon and its relationship to uncertainty dynamics, they are based on a specific synthetic setting and simplified theoretical model. Extending these insights to larger models and other language tasks remains an important direction for future work, and the interpretation of uncertainty dynamics in more complex settings should be approached with caution.
Acknowledgments This project was provided with AI computing and storage resources by GENCI at IDRIS thanks to the grant 2025-AD011016811 on the supercomputer Jean Zay’s A100/H100 partition. SR SR acknowledges the support of CIRCALIS AI-HPC facility at EURECOM, with partial funding from French Region Sud.
References Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/ forum?id=0g0X4H8yN4I. Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/ forum?id=liMSqUuVg9. Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight Uncertainty in Neural Networks. In International Conference on Machine Learning (ICML), May 2015. URL http://arxiv.org/abs/1505.05424. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. Kenzo Clauw, Sebastiano Stramaglia, and Daniele Marinazzo. Information-Theoretic Progress Measures reveal Grokking is an Emergent Phase Transition. In ICML 2024 Workshop on Mutual Information (MI), 2024. doi: 10.48550/arXiv.2408.08944. Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace Redux - Effortless Bayesian Deep Learning. In neurips, volume 34, 2021. Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In International Conference on Machine Learning, 2016. Satvik Golechha. Progress Measures for Grokking on Real-world Tasks. In ICML 2024 Workshop on High-dimensional Learning Dynamics (HiLD), 2024. doi: 10.48550/arXiv.2405. 12755. Alex Graves. Practical Variational Inference for Neural Networks. In Advances in Neural Information Processing Systems, 2011. Tianyu He, Darshil Doshi, Aritra Das, and Andrey Gromov. Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks. In Advances in Neural Information Processing Systems, volume 37, 2024. URL https://openreview.net/ forum?id=2406.02550. 10
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: an introduction to concepts and methods. Machine Learning, 3(110):457–506, 2021. Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam. In International Conference on Machine Learning (ICML), 2018. URL http://arxiv.org/abs/ 1806.04854. Michael Kirchhof, Gjergji Kasneci, and Enkelejda Kasneci. Position: Uncertainty Quantification Needs Reassessment for Large Language Model Agents. In Forty-second International Conference on Machine Learning Position Paper Track, 2025. URL https: //openreview.net/forum?id=Lrv20S5RZV. Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. In Advances in Neural Information Processing Systems 30, 2017. URL https://papers.nips.cc/paper/2017/hash/ 9ef2ed4b7fd2c810847ffa5fa85bce38-Abstract.html. Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=K2PTuvVTF1L. Noam Itzhak Levi, Alon Beck, and Yohai Bar-Sinai. Grokking in Linear Estimators – A Solvable Model that Groks without Understanding. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=GH2LYb9XV0. Chen Ling, Xujiang Zhao, Xuchao Zhang, Wei Cheng, Yanchi Liu, Yiyou Sun, Mika Oishi, Takao Osaki, Katsushi Matsuda, Jie Ji, Guangji Bai, Liang Zhao, and Haifeng Chen. Uncertainty Quantification for In-Context Learning of Large Language Models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 3357–3370, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.naacl-long.184. URL https://aclanthology.org/2024.naacl-long.184/. Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651–34663, 2022. Wesley J. Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for Bayesian uncertainty in deep learning. In Advances in Neural Processing Systems (NeurIPS), volume 32, 2019. Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica J. Sutherland. Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 35155–35173. PMLR, 2024. Bálint Mucsányi, Michael Kirchhof, and Seong Joon Oh. Benchmarking uncertainty disentanglement: Specialized uncertainties for specialized tasks. Advances in neural information processing systems, 37:50972–51038, 2024. Theodore Papamarkou, Maria Skoularidou, Konstantina Palla, Laurence Aitchison, Julyan Arbel, David Dunson, Maurizio Filippone, Vincent Fortuin, Philipp Hennig, José Miguel Hernández-Lobato, Aliaksandr Hubin, Alexander Immer, Theofanis Karaletsos, Mohammad Emtiyaz Khan, Agustinus Kristiadi, Yingzhen Li, Stephan Mandt, Christopher Nemeth, Michael A Osborne, Tim G. J. Rudner, David Rügamer, Yee Whye Teh, Max Welling, Andrew Gordon Wilson, and Ruqi Zhang. Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI. In International Conference on Machine Learning, 2024. 11
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv, 2022. Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation for neural networks. In International conference on learning representations, 2018. Mattia Rosso, Simone Rossi, Giulio Franzese, Markus Heinonen, and Maurizio Filippone. Scaling Laws for Uncertainty in Deep Learning. arXiv preprint arXiv:2506.09648, 2025. Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Bazan Clement Emile Marcel Raoul, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Möllenhoff. Variational Learning is Effective for Large Deep Networks. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 44665–44686. PMLR, 21–27 Jul 2024. Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024. Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151–35174. PMLR, 2023. Yifei Wang, Yu Sheng, Linjing Li, and Daniel Dajun Zeng. Uncertainty Unveiled: Can Exposure to More In-context Examples Mitigate Uncertainty for Large Language Models? In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (eds.), Findings of the Association for Computational Linguistics: ACL 2025, pp. 20659–20678, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176256-5. doi: 10.18653/v1/2025.findings-acl.1062. URL https://aclanthology.org/2025. findings-acl.1062/. Lisa Wimmer, Yusuf Sale, Paul Hofman, Bernd Bischl, and Eyke Hüllermeier. Quantifying aleatoric and epistemic uncertainty in machine learning: Are conditional entropy and mutual information appropriate measures? In Conference on Uncertainty in Artificial Intelligence, 2023. Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An Explanation of In-context Learning as Implicit Bayesian Inference. In International Conference on Learning Representations, 2022. Shifeng Xie, Rui Yuan, Simone Rossi, and Thomas Hannagan. The Initialization Determines Whether In-Context Learning Is Gradient Descent. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https://openreview.net/forum?id=fvqSKLDtJi. Hanlin Zhang, YiFan Zhang, Yaodong Yu, Dhruv Madeka, Dean Foster, Eric Xing, Himabindu Lakkaraju, and Sham Kakade. A Study on the Calibration of In-context Learning. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 6118–6136, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.naacl-long.340. URL https://aclanthology.org/2024.naacl-long.340/. Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and How Does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization. In Proceedings of the 28th International Conference on Artificial Intelligence and Statistics (AISTATS), 2025.
12
A
Extended Related Works
Grokking and ICL are training-inference phenomena that expose a common theme in deep learning: generalization can emerge abruptly, late, and in ways not well predicted by early optimization signals. Grokking, originally characterized on small algorithmic datasets (notably modular arithmetic), is a delayed generalization regime where test performance remains poor long after training performance saturates, and then transitions sharply to near-perfect generalization (Power et al., 2022). ICL, popularized by large transformer language models (e.g., GPT-3), is the ability to adapt to a task at inference time by conditioning on a prompt containing input-output examples, without gradient-based parameter updates (Brown et al., 2020). Explanations of ICL increasingly characterize transformers as implementing recognizable statistical estimators or learning procedures in their forward pass, including ridge regression, gradient descent, and in certain regimes Bayesian estimators (Akyurek et al., 2023; Von Oswald et al., 2023; Bai et al., 2023; Xie et al., 2025). A Bayesian framing provides a natural language for connecting these threads: pretraining can be interpreted as learning an implicit prior (or an implicit family of priors) over tasks/concepts, while ICL corresponds to approximate Bayesian inference (or Bayesian model averaging) over latent task variables given the prompt (Xie et al., 2022; Zhang et al., 2025). In parallel, BDL treats uncertainty as first-class, via priors and posteriors over parameters and predictive distributions, and offers tools that make uncertainty trajectories during training and inference measurable (Graves, 2011; Blundell et al., 2015; Gal & Ghahramani, 2016; Lakshminarayanan et al., 2017; Maddox et al., 2019; Ritter et al., 2018). In-context learning. While the theoretical understanding of ICL is still developing, a growing body of work has empirically studied various aspects of ICL, including calibration and uncertainty measures under prompting. Zhang et al. (2024) systematically study calibration under ICL across NLP tasks and report that increasing the number of ICL examples can initially worsen miscalibration before later improving calibration, and that miscalibration is prominent in low-shot settings. Ling et al. (2024) propose a formulation that decomposes predictive uncertainty in ICL into epistemic and aleatoric components and introduce estimation methods aimed at plug-and-play uncertainty decomposition for ICL. Wang et al. (2025) study long-context ICL and report that adding more in-context examples can reduce total uncertainty and epistemic uncertainty, while also analyzing confidence evolution across layers. Grokking. Empirical studies of grokking have primarily characterized loss and accuracy dynamics, changes in representations, and mechanistic interpretations of the phenomenon. Power et al. (2022) established the canonical grokking setting on small modular arithmetic tasks, and observed that generalization emerges suddenly after a long period of overfitting. Liu et al. (2022) provided an extensive phase-diagram evidence and identified multiple learning phases, emphasizing structured representation emergence as a driver of generalization. Mohamadi et al. (2024) focused specifically on modular addition and propose a theoretical explanation based on the NTK regime, showing that the transition to generalization corresponding to leaving the kernel regime happens only after initial overfitting. A subset of grokking-adjacent work introduces information-theoretic progress measures, which is conceptually close to uncertainty quantification but not identical to predictive uncertainty. Clauw et al. (2024) analyze higher-order mutual information (synergy/redundancy) among neurons through training and report distinct phases that can anticipate grokking, framing grokking as an emergent phase transition caused by synergistic interactions. Golechha (2024) investigates grokking on real-world classification tasks and proposes progress measures such as activation sparsity, absolute weight entropy, and approximate local circuit complexity that correlate with grokking more strongly than weight norms in that study. Difference from prior work. Closer to our work is the work of He et al. (2024), which extended the classic grokking setting to a larger class of modular linear functions, allowing for a more detailed analysis of the learning dynamics not only in-distribution but also out-ofdistribution. While we build on the same experimental set up, our focus is on the Bayesian learning dynamics and uncertainty trajectories during grokking, rather than characterizing 13
the phases of learning or the representational changes. Notably, we identify a the condition under which the model can generalize in-domain and out-of-domain, and we analyze how the model’s uncertainty evolves as it transitions from overfitting to generalization, providing a novel perspective on the grokking phenomenon.
B
Description of the methods used
B.1
Improved Variational Online Newton (IVON)
IVON (Shen et al., 2024) is a scalable, second-order optimizer that makes variational inference tractable for large-scale deep neural networks. Instead of minimizing the empirical risk, IVON directly optimizes the Evidence Lower Bound (ELBO):
L(q) = λEq(w) [ℓ̄(w)] + DKL (q(w)|| p(w))
(14)
where, q(w) = N (w; µ, diag(Σ)2 ) is the diagonal Gaussian posterior approximation, p(w) is the prior, ℓ̄(w) is the empirical loss, and λ scales the objective. Previous methods like VON (Khan et al., 2018) relied on a computationally expensive per-example Gauss-Newton Hessian estimate. IVON circumvents this by using the reparameterization trick to estimate the Hessian ĥ using only the standard minibatch gradient: ˆ ℓ̄(w) ⊙ w − µ ĥ ← ∇ Σ2
(15)
where ⊙ is the element-wise product and w ∼ q(w). By requiring only a single vector multiplication, IVON’s computational cost matches Adam’s. To ensure the Hessian remains positive-definite, IVON applies a Riemannian gradient descent update: 1 (h − ĥ)2 h ← β 2 h + (1 − β 2 )ĥ + (1 − β 2 )2 2 h+δ
(16)
Here, β 2 is the Hessian momentum, and δ > 0 acts as a damping parameter corresponding to the prior precision (weight decay). The mean µ is updated via a preconditioned Newton-like step, and the standard deviation p is tied to the Hessian via Σ = 1/ λ(h + δ). At inference, IVON quantifies uncertainty through Bayesian Model Averaging (BMA) by sampling w(s) ∼ N (µ, diag(Σ)2 ). B.2
Laplace approximation (LA)
Laplace approximation (Daxberger et al., 2021) is a post-hoc framework that converts a deterministically trained neural network into a Bayesian model for uncertainty quantification, without modifying the training objective. LA interprets the final weights trained with L2 regularization as a Maximum A Posteriori (MAP) estimate, wMAP , obtained by minimizing the regularized negative log-posterior: N
wMAP = arg min − ∑ log p(yn | f w ( xn )) − log p(w) w
! (17)
n =1
where p(w) = N (w; 0, τ 2 I d ) is the isotropic Gaussian prior induced by weight decay. LA constructs a local Gaussian approximation of the intractable true posterior centered at the MAP solution: 14
p(w | D) ≈ N (w; wMAP , Σ)
(18)
The covariance matrix Σ = H −1 is the inverse of the Hessian H, derived via a second-order Taylor expansion evaluated at wMAP . Computing and inverting the full exact Hessian H ∈ Rd×d is computationally infeasible and often yields an indefinite matrix. Practical LA implementations resolve this through three core approximations: 1. Fisher Information Matrix (FIM): Replaces the exact Hessian to guarantee positive semi-definiteness. 2. Last-Layer Laplace (LLLA): Restricts Bayesian inference to the final linear layer while keeping preceding layers fixed at wMAP , drastically reducing dimensionality. 3. Kronecker-Factored Approximate Curvature (KFAC): Approximates the Fisher matrix as a Kronecker product F ≈ A ⊗ B (input and pre-activation gradient covariances) for efficient inversion. The predictive distribution for a new data point x∗ is computed by marginalizing over this approximate posterior: p(y | x∗ , D) ≈
Z
p(y | f w ( x∗ ))N (w; wMAP , Σ) dw
(19)
In practice, this integral is approximated via Monte Carlo sampling, producing a distribution of outputs to quantify epistemic uncertainty.
C
A Simplified Model for Grokking under Variational Inference
This section rewrites the derivation in a self-contained form for a linear-Gaussian teacherstudent model. The goal is to isolate the mechanism that produces distinct training and generalization timescales under Gaussian variational inference, while keeping all quantities analytically tractable. Note that while the proof style and some of the intermediate results are similar to those in Levi et al. (2024), the final results and the interpretation are different, as we focus on the variational learning dynamics rather than the gradient flow of the population risk. Furthermore, the analysis in Levi et al. (2024) is mostly focused on loss and accuracy dynamics, while we make a specific focus on the uncertainty dynamics and how they relate to the generalization transition, which is a novel perspective on the grokking phenomenon. C.1
Bayesian linear regression
Definition C.1 (Linear-Gaussian model). Let X ∈ Rn×d be the design matrix, let w ∈ Rd be the parameter vector, and let y ∈ Rn be the observed responses. We consider the Bayesian linear model p(y | w) = N (y; Xw, σ2 I n ), 2
p(w) = N (w; 0, τ I d ),
(20) (21)
with likelihood variance σ2 > 0 and prior variance τ 2 > 0. When needed, we also assume the true data-generating process is linear-Gaussian: y = x⊤ w⋆ + ε,
E[ε] = 0,
with population covariance Σ x ≡ E[ xx⊤ ]. 15
E[ε2 ] = σε2 ,
(22)
Definition C.2 (Posterior precision and linear term). For the model above, define X⊤X I + d2 , σ2 τ X⊤y b≡ . σ2
A≡
(23) (24)
Since τ 2 > 0, the matrix A is symmetric positive definite. The posterior distribution is for this model is the following Gaussian: p(w | y) = N (w; µ, Σ),
(25)
Σ = A −1 , µ = Σb.
(26) (27)
where
C.2
Gaussian variational inference as a Wasserstein flow
Definition C.3 (Gaussian variational family and Bures-Wasserstein flow). We approximate the posterior by a time-dependent Gaussian qt (w) = N (w; µt , Σt ),
(28)
and study the Bures-Wasserstein gradient flow of the variational objective (Lambert et al., 2022). In the Gaussian family, the flow can be written as qt (wt ) µ̇t = −Eqt ∇wt log , (29) p(y, wt ) " # qt (wt ) ⊤ qt (wt ) ⊤ (wt − µt ) + (wt − µt ) ∇wt log . Σ̇t = −Eqt ∇wt log p(y, wt ) p(y, wt ) Proposition C.1 (Closed ODEs for the variational mean and covariance). The Gaussian variational parameters obey µ̇t = b − Aµt ,
(30)
Σ̇t = 2I d − AΣt − Σt A.
(31)
1 1 log qt (w) = − (w − µt )⊤ Σ− t ( w − µt ) + const, 2
(32)
Proof. We have
so 1 ∇w log qt (w) = −Σ− t ( w − µ t ), ∇w log p(y, w) = b − Aw.
(33) (34)
Subtracting these two gradients yields
−∇w log
qt (w) 1 = b − Aw + Σ− t ( w − µ t ). p(y, w)
Taking the expectation of Equation (35) under qt and using Eqt [w] = µt gives qt (w) −Eqt ∇w log = b − Aµt , p(y, w) which proves Equation (30). 16
(35)
(36)
For the covariance dynamics, set d = w − µt so that Eqt [d] = 0 and Eqt [dd⊤ ] = Σt . Then i h i h qt (w) ⊤ 1 Eq t −∇w log d⊤ = Eqt (b − Aw)d⊤ + Eqt Σ− t dd p(y, w) = − AΣt + I d . (37) Indeed, the first term becomes h i Eqt (b − A(µt + d))d⊤ = − AEqt [dd⊤ ] = − AΣt ,
(38)
and the second simplifies to i h ⊤ 1 1 dd = Σ− Eq t Σ − t Σt = I d . t
(39)
Plugging this identity and its transpose into Equation (29) yields Σ̇t = ( I d − AΣt ) + ( I d − AΣt )⊤ = 2I d − AΣt − Σt A,
(40)
because A is symmetric. Theorem C.1 (Explicit variational trajectory and convergence). Let µ∞ ≡ A−1 b and Σ∞ ≡ A−1 . Then the solution of Equations (30) and (31) is µt = µ∞ + e− At (µ0 − µ∞ ), Σt = Σ∞ + e
− At
( Σ0 − Σ ∞ ) e
(41) − At
.
(42)
In particular, (µt , Σt ) converges exponentially fast to the exact posterior parameters from Equation (26). Proof. The fixed point of Equation (30) is µ∞ = A−1 b. Writing δµt ≡ µt − µ∞ , we obtain δµ˙ = − Aδµ , (43) t
t
hence δµt = e− At δµ0 , which proves Equation (41). Likewise, Σ∞ = A−1 solves the stationary equation 2I d − AΣ∞ − Σ∞ A = 0. Set ∆t ≡ Σt − Σ∞ . Then ∆˙ t = − A∆t − ∆t A. (44) Since A is symmetric, the function e− At ∆0 e− At satisfies the same ODE and initial condition, proving Equation (42). Because A ≻ 0, the matrix exponential decays exponentially, so µt → µ∞ and Σt → Σ∞ . Tese limits are exactly the posterior mean and covariance. C.3
Empirical and generalization risks
Definition C.4 (Empirical and population risks). For a deterministic predictor w, define the squared-loss risks 1 ∥y − Xw∥2 , 2n 1 Rgen (w) = E p(x,y) (y − x⊤ w)2 . 2 Along the Gaussian variational trajectory, we are interested in the expected risks
Remp (w) =
Lemp (t) ≡ Eqt Remp (w), Lgen (t) ≡ Eqt Rgen (w).
(45) (46) (47) (48)
Proposition C.2 (Exact risks under a Gaussian predictor). For any Gaussian predictor qt (w) = N ( µ t , Σ t ), 1 Lemp (t) = (49) ∥y − Xµt ∥2 + Tr X ⊤ XΣt , 2n 1 Lgen (t) = (µt − w⋆ )⊤ Σ x (µt − w⋆ ) + Tr(Σ x Σt ) + σε2 . (50) 2 17
Proof. Let w ∼ N (µt , Σt ). Then Eqt ∥y − Xw∥2 = ∥y − Xµt ∥2 + Tr X ⊤ XΣt ,
(51)
which proves Equation (49). For the population risk, write y = x⊤ w⋆ + ε and use independence of ε, x, and w: 2 E(y − x⊤ w)2 = E x⊤ (w⋆ − w) + σε2 h i = Ex (w⋆ − µt )⊤ xx⊤ (w⋆ − µt ) + Ex Tr xx⊤ Σt + σε2
= (µt − w⋆ )⊤ Σ x (µt − w⋆ ) + Tr(Σ x Σt ) + σε2 .
(52)
Multiplying by 1/2 gives Equation (50). Corollary C.1 (Exact risk decomposition along the VI flow). Define Et ≡ e− At ,
δµ0 ≡ µ0 − µ∞ ,
∆Σ0 ≡ Σ0 − Σ∞ .
(53)
Let r ∞ ≡ y − Xµ∞ and b∞ ≡ µ∞ − w⋆ . Then 1 1 1 ⊤ Lemp (t) = Lemp,∞ − r ⊤ δµ0 Et X ⊤ X Et δµ0 + Tr X ⊤ X Et ∆Σ0 Et , ∞ X Et δµ0 + n 2n 2n 1 1 ⊤ ⊤ Lgen (t) = Lgen,∞ + b∞ Σ x Et δµ0 + δµ0 Et Σ x Et δµ0 + Tr(Σ x Et ∆Σ0 Et ), 2 2
(54) (55)
where 1 ∥r ∞ ∥2 + Tr X ⊤ XΣ∞ , 2n 1 ⊤ Lgen,∞ ≡ b∞ Σ x b∞ + Tr(Σ x Σ∞ ) + σε2 . 2
Lemp,∞ ≡
(56) (57)
Proof. Substitute Equations (41) and (42) into Equations (49) and (50) and expand the resulting quadratic forms. The empirical and generalization fixed-point terms are the parts independent of Et ; the linear and quadratic time-dependent contributions are exactly those displayed above. Corollary C.2 (Matched interpolation regime). Assume Σ x = I d and that the stationary predictor is well-specified in the sense that µ∞ = w⋆ ,
r ∞ = 0.
(58)
Define the excess risks
Leemp (t) ≡ Lemp (t) − Lemp,∞ , Legen (t) ≡ Lgen (t) − Lgen,∞ .
(59) (60)
Then 1 ⊤ 1 Leemp (t) = δµ0 Et X ⊤ X Et δµ0 + Tr X ⊤ X Et ∆Σ0 Et , 2n 2n 1 1 ⊤ 2 Legen (t) = δµ0 Et δµ0 + Tr( Et ∆Σ0 Et ). 2 2
(61) (62)
Proof. In Equations (54) and (55), the linear terms vanish because r ∞ = 0 and b∞ = 0, while Σ x = I d simplifies the quadratic and trace terms. 18
C.4
Spectral reduction and random-matrix asymptotics
Lemma C.1 (Spectral form of the excess risks). Let b ≡ 1 X ⊤ X = U diag(ν1 , . . . , νd )U ⊤ . Σ n
(63)
Then A = U diag( a1 , . . . , ad )U ⊤ ,
ai =
n 1 νi + 2 . 2 σ τ
(64)
d̄0 ≡ U ⊤ δµ0 ,
S̄0 ≡ U ⊤ ∆Σ0 U,
(65)
Writing
the excess risks in Equations (61) and (62) become 1 d Leemp (t) = ∑ νi d¯20,i + (S̄0 )ii e−2ai t , 2 i =1
(66)
1 d Legen (t) = ∑ d¯20,i + (S̄0 )ii e−2ai t . 2 i =1
(67)
b and A are diagonalized by the same orthogonal matrix U. In that basis, Proof. Σ Et = e− At = U diag(e−a1 t , . . . , e− ad t )U ⊤ .
(68)
Substituting this into Equations (61) and (62) immediately yields the stated sums. Assumption C.1 (Self-averaging initialization). In the large-system limit, assume that the b so that initial mode weights are asymptotically uncorrelated with the eigenbasis of Σ, 2
∥δµ0 ∥ + Tr(∆Σ0 ) d¯20,i + (S̄0 )ii ≈ d
(69)
uniformly over i. Proposition C.3 (Self-averaged loss formulas). Under Assumption C.1, if 1 C0 ≡ ∥δµ0 ∥2 + Tr(∆Σ0 ) , 2
(70)
then 1 d n 1 Leemp (t) ≈ C0 ∑ νi exp −2 ν + t , i d i =1 σ2 τ2 1 d n 1 e Lgen (t) ≈ C0 ∑ exp −2 νi + 2 t . d i =1 σ2 τ If, moreover, d, n → ∞ with aspect ratio λ = d/n < 1, then n 1 e Lemp (t) ≈ C0 Eν∼MP(λ) ν exp −2 ν+ 2 t , σ2 τ 1 n e Lgen (t) ≈ C0 Eν∼MP(λ) exp −2 ν+ 2 t . σ2 τ
(71) (72)
(73) (74)
Proof. Replace the mode-dependent prefactors in Equations (66) and (67) by the common self-averaged value from Assumption C.1. This yields the finite sums. The MarchenkoPastur formulas follow by replacing the empirical spectral average with its large-d, n limit.
19
Theorem C.2 (Late-time asymptotics and grokking delay). Assume the setting of Proposition C.3 with λ < 1, and let √ √ 2 n 1 ν± = (1 ± λ)2 , κ=2 ( 1 − λ ) + . (75) σ2 τ2 Then, as t → ∞, Leemp (t) ∼ Kemp t−3/2 e−κt , (76)
Legen (t) ∼ Kgen t−3/2 e−κt ,
(77)
with prefactor ratio Kgen 1 1 √ . = = Kemp ν− (1 − λ )2 If the grokking time is defined by the 95% accuracy threshold e t⋆ ) = ϵ , L( 8 then the leading train-generalization delay is 1√ log ⋆ ⋆ √1− λ 1 . ∆tgrok ≡ tgen − temp ≃ n ( 1 − λ )2 + τ 2 σ2 Proof. The Marchenko-Pastur density on [ν− , ν+ ] behaves near its lower edge as √ pMP (ν) ∝ ν − ν− .
(78)
(79)
(80)
(81)
2 + u. In Equations (73) and (74), the exponential factor is e−κt e−2(n/σ )ut . Hence
Write ν = ν− both integrals are governed by the same edge contribution e−κt
Z ∞ 0
2
u1/2 e−2(n/σ )ut du ∝ t−3/2 e−κt ,
(82)
which proves Equations (76) and (77). The empirical loss contains an extra factor ν = ν− + u, which contributes only the multiplicative constant ν− at leading order; this yields Equation (78). To obtain the delay, solve Kt−3/2 e−κt = ϵ/8 asymptotically for t. The subleading t−3/2 correction is the same for the empirical and generalization losses, so the leading difference between the two solutions is log Kgen /Kemp ∆tgrok ≃ . (83) κ Using Equation (78) and the definition of κ gives Equation (80). C.5
Accuracy and uncertainty of the sampled predictor
Proposition C.4 (Exact accuracy of the sampled VI predictor). Let w ∼ qt . For a fixed datapoint ( x, y), define mt ( x, y) ≡ y − x⊤ µt ,
vt ( x) ≡ x⊤ Σt x.
(84)
Then the residual rt ( x, y; w) ≡ y − x⊤ w (85) is conditionally Gaussian under qt , and for any tolerance ϵ > 0 the empirical and population accuracies are " ! !# √ √ 1 n ϵ − mi ( t ) − ϵ − mi ( t ) p p Aemp (t) = ∑ Φ −Φ , (86) n i =1 vi ( t ) vi ( t ) " ! !# √ √ ϵ − x⊤ (w⋆ − µt ) − ϵ − x⊤ (w⋆ − µt ) p p Agen (t) = E p(x) Φ −Φ , (87) σε2 + x⊤ Σt x σε2 + x⊤ Σt x where mi (t) = yi − xi⊤ µt and vi (t) = xi⊤ Σt xi , and Φ is the standard Gaussian CDF. 20
Proof. Since w ∼ N (µt , Σt ), the scalar random variable x⊤ w is Gaussian with mean x⊤ µt and variance x⊤ Σt x. Therefore rt ( x, y; w) | ( x, y) ∼ N (mt ( x, y), vt ( x)).
(88) √ √ The formulas follow by evaluating the Gaussian probability of the interval [− ϵ, ϵ]. For the generalization expression, use y = x⊤ w⋆ + ε and absorb the label noise into the conditional variance. Proposition C.5 (Entropy decomposition of predictive uncertainty). Assume the predictive model y | x, w ∼ N ( x⊤ w, σε2 ).
(89)
Then the predictive entropy decomposes as H[y | x, D , t] = Eqt (w) H[y | x, w] + It (y; w | x, D).
(90)
Moreover, the predictive distribution is pt (y | x, D) = N y; x⊤ µt , σε2 + x⊤ Σt x ,
(91)
and the pointwise epistemic uncertainty is 1 x⊤ Σt x EU( x, t) = log 1 + 2 σε2
! .
Using the explicit covariance trajectory from Equation (42), this can also be written as ! 1 x⊤ Σ∞ x + x⊤ Et ∆Σ0 Et x EU( x, t) = log 1 + . 2 σε2
(92)
(93)
Proof. The entropy decomposition is the standard identity H (Y | X ) = EH (Y | X, W ) + I (Y; W | X ). Since both p(y | x, w) and qt (w) are Gaussian, marginalizing overw gives Equation (91). The differential entropy of a Gaussian N (0, s2 ) is 12 log 2πes2 , so the aleatoric term is 12 log 2πeσε2 and the epistemic contribution is the difference between the predictive and aleatoric entropies, which yields Equation (92). Remark C.1 (Late-time accuracy closure). In the matched interpolation regime, the late-time residual is approximately centered and its variance is controlled by the excess risk: e t ). E[rt2 ] ≈ 2L( Approximating the residual by a single effective Gaussian therefore gives s ϵ . Aemp/gen (t) ≈ erf 4Leemp/gen (t)
(94)
(95)
e t⋆ ) = ϵ/8 because The √ 95% grokking threshold used in Theorem C.2 then corresponds to L( erf( 2) ≈ 0.95. Lemma C.2 (Trace dynamics of the covariance). The covariance trace obeys d Tr(Σt ) = 2 Tr( I d − AΣt ). dt
(96)
In the eigenbasis of A, if Σt = Σ∞ + Et ∆Σ0 Et and ∆si denotes the i-th diagonal entry of U ⊤ ∆Σ0 U, then d d Tr(Σt ) = −2 ∑ ai ∆si e−2ai t . dt i =1
21
(97)
Proof. Taking the trace of Equation (31) gives d Tr(Σt ) = 2 Tr( I d ) − Tr( AΣt ) − Tr(Σt A) = 2 Tr( I d − AΣt ), dt
(98)
using cyclicity of the trace. The spectral formula follows by substituting the explicit solution from Equation (42). Lemma C.3 (Population-averaged epistemic uncertainty). The exact population-averaged epistemic uncertainty is ! 1 x⊤ Σt x EU(t) ≡ E p( x) EU( x, t) = E p( x) log 1 + . (99) 2 σε2 Under isotropic Gaussian design x ∼ N (0, I d ) and the same self-averaging closure used above, Tr(Σt ) 1 . (100) EU(t) ≈ log 1 + 2 σε2 If Σ0 ≺ Σ∞ , this approximate epistemic curve is monotonically increasing; if Σ0 ≻ Σ∞ , it is monotonically decreasing; and if Σ0 = Σ∞ , it is constant. Proof. The exact identity is just the average of Equation (92). Under isotropic Gaussian design, the quadratic form x⊤ Σt x concentrates around its trace in high dimension, yielding Equation (100). Because u 7→ 12 log 1 + u/σε2 is increasing, the monotonicity of the approximate entropy curve matches the monotonicity of Tr(Σt ). The explicit solution Equation (42) shows that each spectral mode evolves as ai−1 + ∆si e−2ai t , so the sign of ∆si determines whether that mode increases or decreases. If Σ0 ≺ Σ∞ , then all ∆si < 0 and Lemma C.2 implies that Tr(Σt ) increases monotonically; the opposite case is analogous, and equality gives a constant curve. Corollary C.3 (Isotropic initialization). If Σ0 = s0 I d , then 1 1 1 −2a(ν)t Tr(Σt ) ≈ Eν∼MP(λ) + s0 − e , d a(ν) a(ν)
a(ν) =
n 1 ν+ 2. σ2 τ
(101)
In particular, when Σ0 = 0, " # 1 − e−2a(ν)t 1 Tr(Σt ) ≈ Eν∼MP(λ) . d a(ν)
(102)
Proof. Under isotropic initialization, ∆Σ0 is diagonal in every eigenbasis, so the spectral representation of Equation (42) reduces to a scalar mode-by-mode formula. Averaging over the empirical spectrum and then passing to the Marchenko-Pastur limit gives the result. Proposition C.6 (Epistemic uncertainty of the threshold model). Define the thresholded success indicator Zϵ ( x, y; w) ≡ 1{(y − x⊤ w)2 ≤ ϵ}. For fixed ( x, y), the success probability under qt is ! √ ϵ − mt ( x, y) p pt ( x, y) = Φ −Φ vt ( x)
! √ − ϵ − mt ( x, y) p , vt ( x)
(103)
(104)
and Zϵ ( x, y; w) | ( x, y, D , t) ∼ Bernoulli( pt ( x, y)).
(105)
Hence the threshold-model epistemic uncertainty is exactly EUt ( Zϵ ; w | x, y, D) = h2 pt ( x, y) , 22
(106)
where h2 ( p) = − p log p − (1 − p) log(1 − p) is the binary entropy. In the late-time self-averaging regime, ! s ϵ pt ≈ erf , (107) e t) 4L( and therefore s EUth (t) ≈ h2 erf
ϵ e 4L(t)
!! .
(108)
Proof. The exact success probability is the Gaussian interval probability from Proposition C.4. Conditional on w, the variable Zϵ is deterministic, so its conditional entropy vanishes. The predictive distribution is therefore Bernoulli with entropy h2 ( pt ( x, y)), and the full uncertainty is epistemic. The closure formula follows from the same effective-Gaussian approximation that relates the residual variance to the excess risk at late times. Remark C.2 (Threshold uncertainty versus grokking). The binary entropy h2 ( p) is maximized at p = 1/2, so the threshold-model epistemic uncertainty typically follows a non-monotone trajectory: it rises while predictions are ambiguous and then falls as the model becomes confidently correct. Using Equation (107), the peak occurs when ϵ e tpeak ) = L( , (109) −1 4(erf (1/2))2 and numerically 1
≈ 1.10. (110) (1/2))2 Thus the peak occurs when the loss is still of order ϵ, which is much earlier than the 95% e t⋆ ) = ϵ/8. Thus the entropy peak and the grokking delay are distinct grokking threshold L( observables, even though both are controlled by the same slow edge-of-spectrum modes in the Marchenko-Pastur regime. Theorem C.3 (Delay between threshold-uncertainty peaks). Assume the setting of Theorem C.2 and Proposition C.6, and suppose the empirical and generalization threshold-model uncertainty peaks occur in the late-time regime where Equations (76), (77) and (108) apply. Then ∆tEU peak (λ) ≡ tpeak,gen − tpeak,emp 1√ log √1− λ 1 , ≃ n (111) (1 − λ )2 + τ 2 σ2 4(erf
−1
up to threshold-dependent subleading corrections. Proof. For α ∈ {emp, gen}, write s
h α ( t ) ≈ h2 p α ( t ) ,
pα (t) ≈ erf
as in Equation (108). Since h2′ ( p) = log p = 1/2. Moreover, the map
ϵ 4Leα (t)
! ,
(112)
1− p p , the binary entropy is uniquely maximized at
r
Le 7−→ erf
ϵ 4Le
(113)
is strictly decreasing in Le, while in the late-time regime Equations (76) and (77) imply that each loss is strictly decreasing for large t. Hence each late-time uncertainty curve reaches its maximum exactly when pα (t) = 1/2, namely when ϵ Leα (tpeak,α ) = cpeak ≡ , (114) −1 4(erf (1/2))2 23
which is Equation (109). Now use the late-time asymptotics from Theorem C.2. For each α ∈ {emp, gen},
Leα (t) ∼ Kα t−3/2 e−κt .
(115)
−3/2 −κtpeak,α cpeak = Kα tpeak,α e 1 + o (1) .
(116)
Therefore the peak time satisfies
Taking logarithms gives Kα
κtpeak,α = log
cpeak
!
−
3 log tpeak,α + o (1). 2
Subtracting the empirical and generalization versions yields ! tpeak,gen Kgen 3 − log , κ∆tEU peak ≃ log Kemp 2 tpeak,emp
(117)
(118)
up to o (1) terms. The logarithmic ratio of peak times is a subleading correction coming from the common t−3/2 factor, exactly as in the proof of Theorem C.2. Thus the leading contribution is log Kgen /Kemp ∆tEU peak ≃ , (119) κ with threshold dependence entering only through the discarded subleading terms. Finally, Equation (78) gives √ Kgen (120) log = −2 log 1 − λ , Kemp and κ=2
√ 2 n 1 ( 1 − λ) + 2 σ2 τ
.
(121)
Substituting these expressions gives Equation (111). Thus the threshold-model uncertainty peak and the grokking delay are different observables, but near λ ≃ 1 they are governed by the same edge-of-spectrum slowdown.
D
Reproducibility
We summarize here the full experimental specification needed to reproduce the reported results. Across all experiments, we keep the task family, the transformer architecture, and the evaluation protocol fixed, and vary only the quantities explicitly listed. D.1
Task distribution and sequence construction
Each example is a two-variable modular arithmetic rule z = ( ax + by) mod 29, where ( a, b) identifies the task and ( x, y) is the input pair. The arithmetic base, vocabulary size, and modulus are all 29. The input space therefore contains 292 = 841 distinct pairs ( x, y). For each experimental setting, we first sample ntask in-distribution coefficient pairs and then apply a parallelogram augmentation, which expands the in-distribution family to 4ntask related tasks, check He et al. (2024) for details on the augmentation procedure. Outof-distribution evaluation uses 256 additional tasks drawn from the complement of the in-distribution set. 24
Table 1: Task-generation and sequence-construction choices shared across the main experiments. Quantity
Value
Arithmetic rule Vocabulary size Task dimension Input-pair dimension Input space size Base in-distribution task count Parallelogram augmentation Effective in-distribution task count Out-of-distribution task count Training-point fraction Context length Token window Nominal batch size
z = ( ax + by) mod 29 29 2 2 292 = 841 ntask Enabled in all reported runs 4ntask 256 {0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, depending on sweep 32 triplets 96 tokens 1024
Table 2: Shared architecture for MAP, IVON, and Laplace evaluations. Quantity
Value
Architecture Layers Attention heads Embedding width Feed-forward width Vocabulary size Maximum token window Normalization Positional encoding Embedding tying Bias terms Numerical precision
Decoder-only transformer 6 4 512 2048 29 96 LayerNorm, ϵ = 10−5 Rotary Enabled Disabled Mixed bfloat16
Training and evaluation are performed on sequences of 32 triplets, corresponding to a token window of length 96. Only the answer token z in each triplet is supervised, so learning is next-token prediction on answer positions only. The nominal batch size is 1024 for both training and evaluation. Training uses a virtual in-distribution set multiplier of 1000 to increase the number of sampled batches per optimization run. Every evaluation reports three splits: in-distribution validation, out-of-distribution training-point evaluation, and out-of-distribution validation-point evaluation. D.2
Shared model architecture
All methods use the same decoder-only transformer. The backbone has 6 layers, 4 attention heads, embedding width 512, and feed-forward width 2048. The vocabulary size is 29 and the maximum token window is 96. The model uses LayerNorm with ϵ = 10−5 , full rotary positional encoding, tied input and output embeddings, and no bias terms in attention, MLP, or output projections. All runs use mixed bfloat16 precision and one accelerator per run. D.3
Training and evaluation protocol
The default optimization horizon is 100,000 update steps. For the largest base task setting, ntask = 128, the horizon is doubled to 200,000 steps. For every validation split we report accuracy, expected calibration error, mean log-likelihood, aleatoric uncertainty, and epistemic uncertainty. Accuracy, calibration, and uncertainty are reported both over all supervised answer positions in the sequence and on the final answer token only. 25
Table 3: Optimization settings for the two trained model families. Field
MAP (AdamW)
IVON
Peak learning rate Warmup steps Initial learning-rate scale Final learning-rate scale Schedule Weight decay Momentum parameters Stability / curvature init Gradient control Training samples per step Evaluation samples
1.5 × 10−4
0.5 2000 0.01 0.1 Linear warmup + cosine decay 10−6 (0.9, 0.99999) Hessian init = 1 Clip radius 10−3 1 Monte Carlo sample 16/64 Monte Carlo samples
1000 0.01 0.1 Linear warmup + cosine decay 1 (0.9, 0.98) ϵ = 10−8 Norm clip at 1.0 Deterministic model Deterministic model
Table 4: Laplace-specific settings. Field
Laplace approximation
Base solution Bayesian scope Curvature structure Prior precision Posterior samples Fit data Role
Final MAP model Last layer only Kronecker-factored Hessian Equal to MAP weight decay 16/64 Clean in-distribution training data Post-hoc predictive uncertainty estimation
D.4
Method-specific settings
The MAP baseline is trained with AdamW using peak learning rate 1.5 × 10−4 , linear warmup for 1000 steps from 1% of the peak rate, and cosine decay to 10% of the peak rate. The momentum parameters are (0.9, 0.98), the numerical stability constant is 10−8 , and the gradient norm is clipped at 1.0. IVON uses learning rate 0.5, 2000 warmup steps, the same 1% → 100% → 10% cosine schedule envelope, β 1 = 0.9, β 2 = 0.99999, Hessian initialization 1, clip radius 10−3 , one Monte Carlo weight sample during training. The effective sample size is set automatically to the number of in-distribution training point pairs multiplied by the effective in-distribution task count after augmentation. The Laplace model is fit post hoc on the final MAP solution. The approximation is restricted to the last layer, uses a Kronecker-factored Hessian, and sets the prior precision equal to the MAP weight decay (no marginal likelihood optimization). The Laplace fit is always performed on clean in-distribution training data, with full context sequences of 32 triplets. D.5
Experiment families
The main MAP study sweeps base task count, training-point fraction, and initialization seed. Specifically, it uses ntask ∈ {8, 16, 32, 64, 128}, ptrain ∈ {0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, and 4 seeds. The full Laplace study is aligned exactly with this grid, using the corresponding final MAP solutions. The primary IVON study uses the same sweep axes over base task count and training-point fraction, with 4 seeds. Unless otherwise stated, these runs use the default IVON weight decay 10−5 , automatic effective sample size resolution, and 100,000 optimization steps, with the same 200,000-step extension for ntask = 128. Inference-time context-length sweeps evaluate context sizes 1, 2, . . . , 32 while keeping the trained model fixed. These sweeps use ntask = 64 and ptrain = 0.8. 26
Inference-time context-noise sweeps also use ntask = 64 and ptrain = 0.8, with fixed context length 32. Noise probabilities range over {0.0, 0.1, . . . , 1.0}. Corruption is applied only to eligible context answer tokens; the query triplet and evaluation targets remain clean. D.6
Hardware allocation
All runs use one accelerator per experiment. The main training runs use NVIDIA H200/H100/A100 accelerators while lightweight validation use small H200 MIG partitions. The total computational budget used for the experimental campaign is approximately 12′ 000 accelerator-hours, splitted roughly equally between development and final runs.
E
LLM use disclosure
In accordance with the COLM policy on large language model usage, we disclose that LLMbased tools were used for software engineering tasks and editorial assistance, specifically for code implementation support and code review (beyond simple auto-completion), and rewriting or polishing text (beyond simple grammar checking). They were not used to originate research ideas or formulate scientific claims. The authors take full responsibility for the design of the study, the experiments, the analysis, and the final manuscript.
27