In-Context Learning Before Specialisation: Theory of Learning in a Solvable Topic Model

Francesco Cagnetta; Casper L. Christensen

cs.LGOct 15, 2026

In-Context Learning Before Specialisation: Theory of Learning in a Solvable Topic Model

Authors: Francesco Cagnetta, Casper L. Christensen

Organizations: Simplex · Simplex, Astera Institute

Abstract

Sequence models can generalise in fundamentally distinct ways: extrapolating structures seen in training, or using their prompt to adapt to novel tasks. The prevailing strategy determines how the model behaves beyond its training distribution, controlling potential capabilities and risks. Recent work has characterised many different strategies theoretically, and empirically described the transition between them as the quality and quantity of pretraining data vary. However, there is no learning theory for this competition, hence no fundamental understanding of why one solution might be favoured over another. Here we provide such a theory in a solvable topic model---documents drawn from a finite mixture of latent topics, mirroring the heterogeneity of modern pretraining data. The dataset comes with two distinct solutions: a specialised predictor, that stores the finite topic mixture, and an in-context one, optimal for topics never seen during training. Our theory predicts their emergence at distinct sample complexities: the in-context predictor is available early, while the specialised one emerges only after the training tokens outnumber the parameters of the data-generating process. Trained transformers reproduce both scales and the separating crossover in remarkable agreement with the theory, with no fitted parameters.

Figures & tables

Panel of Figure 1 — Figure 1: In-context learning and specialisation in a random topic model. Left: A random topic model consists of $D$ topics, i.e. unigram probabilities $\mathbf{p}_{z}\in\Delta^{K-1}$ over a vocabulary of size $K$ . Training sequences are obtained by sampling a topic $z$ uniformly, then $T$ tokens i.i.d. from $\mathbf{p}_{z}$ . Right: Learning curve of a one-layer transformer trained for next-token prediction on a random topic model with $(K,D,T)\,{=}\,(32,512,128)$ . Initially, both In-Distribution (ID, on the $D$ training topics, blue) and Out-Of-Distribution (OOD, on a new topic $\mathbf{p}_{\star}$ , orange) descend with $N$ towards the entropy of the Bayes-optimal OOD predictor (orange dotted): the add-one estimator $r_{t}(a\mid X_{\leq t})\,{=}\,(\#{\left\{i\leq t\mid X_{i}=a\right\}}+1)/(t+K)$ . Later, the predictor specialises to the training topics: the ID loss converges to the ID Bayes-optimal floor (blue dotted), while the OOD loss rises as in overfitting. Code: https://github.com/fracagnetta/random-topic-model .

Figure 2: Next-token loss profiles of the two Bayes-optimal predictors and of trained transformers , for

K\,{=}\,64

and

T\,{=}\,128

. Left: ID loss

\overline{\mathcal{L}_{t}[\mathbf{q}]}

of the specialised predictor for

D

from

32

8192

(colour), averaged over realisations of the topic matrix, and OOD loss

\mathcal{L}^{\star}_{t}[\mathbf{r}]

of the in-context predictor (black dashed), which does not depend on

D

. Dotted line:

\log K

. Markers: context length at which each specialised profile has covered half of its descent, growing as

\log D

(grey guide; Sec. 2.1 ). Middle and right: the same two references at

D\,{=}\,32

and

D\,{=}\,8192

, with the ID (blue) and OOD (orange) loss profiles of a one-layer transformer trained on

N

documents at that

D

, mean

\pm

s.e. over four runs with independent topic matrix, corpus and initialisation, at the checkpoint of minimal validation loss. At

D\,{=}\,32

the transformer matches the specialised predictor in distribution and stays above

\log K

out of distribution; at

D\,{=}\,8192

both profiles follow the in-context predictor.

Figure 3: Transformers match the Bayes-optimal specialisation learning curve. Left: Scaled ID autoregressive loss

(\mathcal{L}_{\rm AR}(N)-\mathcal{L}_{\rm AR}(\infty))/(\log{K}-\mathcal{L}_{\rm AR}(\infty))

versus number of training sequences for one-block transformers trained on random topic model data with varying

(K,D)

T\,{=}\,128

. The scaled loss is

1

at the uniform predictor

1/K

0

at the specialised predictor. Middle: Rescaling the

x

-axis as the number of tokens per parameter of the data-generating process collapses all the curves onto a single functional form. Right: The decay towards

0

(black dashed) shows remarkable agreement with the

x^{-1}

rate of the Bayes-optimal learning curve.

Figure 4: In-context learning gives way to specialisation at the parametric scale

NT\sim DK

. Left: Scaled ID (solid) and OOD (dashed) autoregressive losses, both subtracting the ID loss floor for comparison. Right: Rescaling the

x

-axis as the number of tokens per parameter of the data-generating process collapses both the descent of ID losses and the ascent of OOD losses.

Figure 5: Empirical OOD learning curves decay at the predicted sample scale. Left: Scaled OOD autoregressive loss

(\mathcal{L}^{\star}_{\rm AR}(N)-\mathcal{L}^{\star}_{\rm AR}(\infty))/(\log{K}-\mathcal{L}^{\star}_{\rm AR}(\infty))

versus number of training tokens

NT

for one-block transformers trained on random topic model data with varying

(K,D)

T\,{=}\,128

(solid), and with

K\,{=}\,64

D\,{=}\,1024

T\,{=}\,64

(dotted) and

T\,{=}\,256

(dashed). These losses are not recorded at early stopping (training step with the lowest ID validation loss), but at the training step with the lowest OOD validation loss. Hence, the OOD losses shown in this figure do not rise at

N\sim DK/T

as in Fig. 4 . Right: The early-

N

departure from the trivial loss

\log{K}

does not depend on

D

and collapse after dividing the number of training sequences by

K^{2}/G(K,T)

, in perfect agreement with the theory of Eq. 13 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Population spectrum of

M_{2}

for

D\leq K

. Left: ranked eigenvalues

\lambda_{r}\,DK

against the scaled rank

r/D

\alpha\,{=}\,D/K\,{=}\,1/2

, for

K\in\{256,512,1024,2048\}

and one topic realisation per

(K,D)

; dashed: the ranked Marchenko–Pastur law; the spike is the outlier at

\lambda_{1}DK\approx D

. Right: density of the bulk eigenvalues, spike removed, on the

\lambda\,DK

scale, against the Marchenko–Pastur density (dashed).

Figure 7: Population spectrum of

M_{4}

for

K<D\leq K(K+1)/2

. Left: ranked eigenvalues of

M_{4}

for one realisation at

K\,{=}\,192

D\,{=}\,K(K+1)/4\,{=}\,9264

: the spike at

K^{-2}

(dotted), the

K-1

eigenvalues of the linear block and, beyond

r\,{=}\,K

(dotted), the

D-K

eigenvalues of the quadratic block, against the Marchenko–Pastur predictions of the two blocks (dashed). Right: Frobenius norm of each block against

K

, for

D\,{=}\,K(K+1)/4

, with the parameter-free predictions (dashed).

Figure 8: Convergence of the empirical pair moment for

D\leq K

. Left: operator-norm error

\|\widehat{M}_{2}-M_{2}\|_{\mathrm{op}}

, geometric mean over eight replicas, against

NK

for eleven

(K,D)

configurations,

\alpha\,{=}\,D/K\in\{1/4,1/2,3/4\}

times

K\in\{128,256,512\}

plus

K\,{=}\,256

with

D\in\{16,32,64,128,192\}

; colour encodes

K

, line style

\alpha

; all curves collapse on the dotted

1.5/\sqrt{NK}

; circles mark where each curve crosses its own population

\lambda_{D}

, beyond which the lower bulk edge is resolved. Right: overlap

\frac{1}{D-1}\|U_{2:D}^{\top}\widehat{U}_{2:D}\|_{F}^{2}

between the bulk eigenvector subspaces of

M_{2}

and

\widehat{M}_{2}

, mean

\pm

one standard deviation over replicas, against

N/(D^{2}K)

; each curve rises from the random-subspace value

(D-1)/(K-1)

towards one around

N/(D^{2}K)=O(1)

(dashed line at one). One topic realisation per

(K,D)

, uniform topic prior.

Figure 9: Learning curves at different document lengths collapse on the number of tokens per parameter. Left: Scaled ID autoregressive loss

(\mathcal{L}_{\rm AR}(N)-\mathcal{L}_{\rm AR}(\infty))/(\log{K}-\mathcal{L}_{\rm AR}(\infty))

versus number of training sequences for one-block transformers trained on random topic model data with

K\,{=}\,128

D\,{=}\,16

and

T\,{=}\,128,256,512

(colour).

\mathcal{L}_{\rm AR}(\infty)

is the autoregressive loss of the specialised predictor at each

T

, averaged over

32

realisations of the topic matrix. Right: Against

NT/(DK)

the three curves coincide: at fixed

N

, longer documents bring the loss closer to its floor in proportion to the number of tokens they contribute.

Figure 10: Transformers match the Bayes-optimal specialisation learning curve at the longest context too.

T

-gram counterpart of Fig. 3 , for the same twelve

(K,D)

T\,{=}\,128

. Left: Scaled ID

T

-gram loss

(\mathcal{L}_{T}(N)-\mathcal{L}_{T}(\infty))/(\log{K}-\mathcal{L}_{T}(\infty))

versus number of training sequences, where

\mathcal{L}_{T}(\infty)

is the mean entropy

\frac{1}{D}\sum_{z}H(\mathbf{p}_{z})

of the topics of each run, the large-context floor of the specialised predictor. Middle: Rescaling the

x

-axis as the number of tokens per parameter of the data-generating process collapses all the curves onto a single functional form, more tightly than for the autoregressive loss. Right: The decay towards

0

(black dashed) follows the

x^{-1}

rate of the Bayes-optimal learning curve.

Figure 11: The ID loss descends to the in-context level at the predicted sample scale. Left: ID autoregressive loss at early stopping, scaled as

(\mathcal{L}_{\rm AR}(N)-\mathcal{L}^{\star}_{\rm AR}[\mathbf{r}])/(\log{K}-\mathcal{L}^{\star}_{\rm AR}[\mathbf{r}])

, where

\mathcal{L}^{\star}_{\rm AR}[\mathbf{r}]\,{=}\,\log K-G(K,T)

is the autoregressive loss of the in-context predictor, the same on ID and OOD documents ( App. C ). This scaled loss is

1

at the uniform prediction,

0

at the in-context solution, and negative at specialisation. Versus number of training tokens

NT

, for the

(K,D)

and

T

of Fig. 5 , with the same line styles. Right: When plotted against

NT\,G(K,T)/K^{2}\,{=}\,N/(2N_{\rm icl})

( Eq. 13 ) the descents to the in-context level coincide for all

K

D

and

T

, with no adjustable constant. The subsequent departure below zero is specialisation.

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

Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning
Sep 24, 2026Zhongjie Shi, Rongjie Lai, Alexander Cloninger +1Transformer Multi-Task Learning
School of Mathematics, Georgia Institute of Technology, Atlanta, GA 30332, United States · Department of Mathematics, Purdue University, West Lafayette, IN 47907, United States · Department of Mathematics and Halıcıoğlu Data Science Institute, University of California, San Diego, La Jolla, CA 92093, United States
Test time training enhances in-context learning of nonlinear functions
Sep 30, 2025Kento Kuwataka, Taiji SuzukiSingle-Index Models Test-Time Training
Department of Mathematical Engineering and Information Physics The University of Tokyo Hongo 7-3-1, Bunkyo-ku, Tokyo
When Context Sticks: Studying Interference in In-Context Learning
Apr 25, 2026Hanna Rød, Dagny Streit, Nils Valseth Selte +1In-Context Learning

Sep 24, 2026J/K move · Enter open · S save

Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning

Zhongjie Shi, Rongjie Lai, Alexander Cloninger, Wenjing Liao

School of Mathematics, Georgia Institute of Technology, Atlanta, GA 30332, United States · Department of Mathematics, Purdue University, West Lafayette, IN 47907, United States · Department of Mathematics and Halıcıoğlu Data Science Institute, University of California, San Diego, La Jolla, CA 92093, United States

Transformer Multi-Task Learning In-Context Learning

Open paper AskPDF arXiv

Transformers achieve remarkable performance by jointly learning broad families of tasks during pretraining and adapting to unseen tasks from only a short prompt. Yet a rigorous mathematical and statistical understanding of this phenomenon remains limited. This paper aims to study how Transformers exploit shared cross-task structure and how this structure affects the sample complexity of in-context learning (ICL). Specifically, we characterize task-space complexity through covering numbers under a prescribed metric, thereby quantifying the low-dimensional cross-task structure without requiring an explicit parametric representation. The resulting cover provides a set of anchor functions, which we use to introduce a task-identification-and-evaluation procedure: context observations localize an unseen task among the anchor functions, and the response at a query is predicted by aggregating the corresponding anchor function query evaluations. For approximation, we explicitly construct a Transformer with Softmax attention to approximate this procedure. For generalization, we derive an error bound that separates the effects of the number of pretraining tasks and the prompt length. The scaling with respect to the number of pretraining tasks is governed by the intrinsic dimensions of the task space and input domain; once sufficiently many tasks are available, the dependence on the prompt context length becomes dimension-free. To the best of our knowledge, this is the first work to quantify cross-task complexity for general nonlinear task families and explicitly construct a Transformer that exploits their low-dimensional structure to perform ICL. Our theory provides a quantitative explanation of how joint pretraining across related tasks improves in-context generalization.

Abstract

In-Context Learning

Explore similar work

Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning

Test time training enhances in-context learning of nonlinear functions

When Context Sticks: Studying Interference in In-Context Learning