Saved in:
Bibliographic Details
Main Authors: Qchohi, Abdessamed, Rossi, Simone
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2604.12434
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908962330247168
author Qchohi, Abdessamed
Rossi, Simone
author_facet Qchohi, Abdessamed
Rossi, Simone
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_12434
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle A Bayesian Perspective on the Role of Epistemic Uncertainty for Delayed Generalization in In-Context Learning
Qchohi, Abdessamed
Rossi, Simone
Machine Learning
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.
title A Bayesian Perspective on the Role of Epistemic Uncertainty for Delayed Generalization in In-Context Learning
topic Machine Learning
url https://arxiv.org/abs/2604.12434