Differential learning kinetics govern the transition from memorization to generalization during in-context learning

TL;DR

This paper tackles the mechanism behind the emergence of in-context learning (ICL) in transformers by introducing a minimal one-layer transformer that disentangles memorization (IWL) from generalization (ICL) into two largely independent sub-circuits. A dynamical theory shows that ICL is acquired via differential learning kinetics, not solely by capacity limits, and predicts a memorization scaling law $I_K(\\infty) \\sim K^{\\nu}$ with $\\nu \\approx 0.7$, yielding a task-diversity threshold $K^* \\sim N^{1/\\nu} e^{-\\beta_0/\\nu}$; the acquisition time scales exponentially with initial conditions, producing long tails and bimodality near the threshold. The theory further explains ICL transience under regularization and predicts a nontrivial relationship between the ICL and IWL losses after acquisition. Empirical validation on synthetic data and the original model confirms the key predictions, linking data distribution, context length, and dynamics to the sharp memorization-generalization transition and its transient nature, with implications for understanding learning in both artificial and biological systems.

Abstract

Transformers exhibit in-context learning (ICL): the ability to use novel information presented in the context without additional weight updates. Recent work shows that ICL emerges when models are trained on a sufficiently diverse set of tasks and the transition from memorization to generalization is sharp with increasing task diversity. One interpretation is that a network's limited capacity to memorize favors generalization. Here, we examine the mechanistic underpinnings of this transition using a small transformer applied to a synthetic ICL task. Using theory and experiment, we show that the sub-circuits that memorize and generalize can be viewed as largely independent. The relative rates at which these sub-circuits learn explains the transition from memorization to generalization, rather than capacity constraints. We uncover a memorization scaling law, which determines the task diversity threshold at which the network generalizes. The theory quantitatively explains a variety of other ICL-related phenomena, including the long-tailed distribution of when ICL is acquired, the bimodal behavior of solutions close to the task diversity threshold, the influence of contextual and data distributional statistics on ICL, and the transient nature of ICL.

Figures (14)

• Figure 1: (a) In the capacity-constrained model, the network's limited capacity to memorize favors ICL acquisition with increasing task diversity. (b) In the differential learning kinetics model, independent sub-circuits contribute towards IWL and ICL. IWL is slower for greater task diversity. The network acquires ICL before the network can significantly memorize the training set. IWL is significantly slowed down as ICL explains most of the loss, but does eventually memorize the training set. The network subsequently loses the ICL capability due to regularization.
• Figure 2: (a) Data generation process: We create a dataset $\mathcal{D}$ consisting of $K$ item-label $(x_i,\ell_i)$ pairs where each $x_i \sim \mathcal{N}(0,1/D)$ and $\ell_i$ is randomly sampled from $\{-1,+1\}$. The network receives a sequence of $N+1$ tokens. Each of the first $N$ tokens is a concatenation $x_i\oplus \ell_i$ of an $(x_i,\ell_i)$ pair sampled uniformly from $\mathcal{D}$ (details in main-text). The final $N+1$ target token consists of only the item $x_i$ as its label component is zero-ed out. The network is trained to correctly predict the label of the last item. Model architecture: The input is normalized using LayerNorm, then fed to our network, which consists of a one-layer attention network followed by a 3-layer MLP. (b) ICL performance demonstrates a sharp transition as a function of data diversity $K$. Further, at the transition threshold $K^*\approx 10^4$, we observe bimodality where the model either memorizes or generalizes. (c) IWL accuracy vs $K$. IWL and ICL accuracies follow opposite trends. (d) ICL accuracy curves show that ICL performance plateaus at the beginning of training but undergoes a rapid transition as ICL is acquired. (e) ICL is transient, i.e., ICL accuracy gradually decreases to chance levels when the parameters in the attention head are heavily regularized.
• Figure 3: Phenomenology of the minimal model. (a) ICL performance in the minimal model demonstrates a sharp transition as a function of data diversity $K$. (b) ICL acquisition is abrupt during training. (c) ICL is transient when $w$ is exclusively regularized.
• Figure 4: The approximate ICL loss landscape $\mathcal{L}$ (fixing $z_{\text{MLP}} = 0$) in the minimal model as a function of the key parameters $\beta, w$ exhibits a nearly flat region close to initialization, but the dynamics always leads to ICL acquisition ($w,\beta \gg 1$).
• Figure 5: $I_K(\infty)$ shows a power-law scaling with $K$ with exponent $\nu\approx 0.7$.