Incremental Learning in Transformers for In-Context Associative Recall

Aditya Varre, Nicolas Flammarion

Keywords: incremental learning, stage-wise learning, training dynamics, incontext learning, associative recall

Abstract: Transformers acquire in-context learning abilities in abrupt phases during training, often unfolding over multiple stages, during which certain keys circuits like induction heads emerge. In this work, we characterize the dynamics behind the emergence of such circuits during these stages. We focus on a synthetic associative recall task, where sequences are drawn from random maps between a permutation group and a vocabulary range and the model is required to complete the mapping of a permutation by retrieving it from the context. On this task, we study the trajectories of gradient flow of a simplified two-layer, attention-only transformer. Leveraging symmetries in both the transformer architecture and the data, we derive conservation laws that guide the dynamics of the parameters. These conservation laws crucially reveal how initialization —both in shape and scale— determines the order of learning as well as the timescales over which such circuits emerge revealing the implicit curriculum. Finally, we provide empirical evidence across different architectural choices, validating our simplifications and generalizing the insights from our analysis beyond the simple setting.

Primary Area: learning theory

Submission Number: 23386