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

TL;DR

This paper investigates how differential learning rates of sub-circuits in transformers govern the shift from memorization to generalization during in-context learning, revealing a scaling law and mechanistic insights.

Contribution

It introduces a theory explaining the transition from memorization to generalization based on learning kinetics, independent of model capacity, supported by experiments on a synthetic task.

Findings

Memorization and generalization sub-circuits learn at different rates.

A memorization scaling law predicts the task diversity threshold for generalization.

The theory explains phenomena like ICL acquisition timing and solution bimodality.

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.