Different simultaneous mechanisms for in-context recall have distinct learning dynamics

Sultan Daniels, Dylan Davis, Dhruv Gautam, Wentinn Liao, Gireeja Ranade, Anant Sahai

Keywords: emergence, in-context learning, time-series, associative recall, learning dynamics

TL;DR: We introduce a new family of toy problems that combine features of linear-regression-style continuous in-context learning (ICL) with discrete associative recall and find distinct learning dynamics for different prediction mechanisms.

Abstract:

We introduce a new family of toy problems that combine features of linear-regression style continuous in-context learning (ICL) with discrete associative recall. We pretrain transformer models on sample traces from this toy, specifically symbolically labeled interleaved observations from randomly drawn linear deterministic dynamical systems, and study if these transformer models can recall the state of a process previously seen in its context when prompted to do so with its in-context label. Training dynamics reveal the emergence of classic recall ability well into training, but surprisingly, well before this recall ability has emerged, a closely related task — predicting the second token in a recalled sequence given the first — shows clear evidence of seemingly recall-related behavior.

Through out-of-distribution experiments, and a mechanistic analysis on model weights via edge pruning, we find that next-token prediction for this toy problem involves two separate mechanisms. One mechanism uses the discrete labels to do the associative recall required to predict the start of a resumption of a previously seen sequence, and the second mechanism, which is largely agnostic to the discrete labels, performs a Bayesian-style prediction based on the previous token and the context. These two mechanisms have different learning dynamics.

To confirm that this two-mechanism (manifesting as separate emergence) phenomenon is not just an artifact of our toy setting, we used OLMo training checkpoints on an ICL translation task to see a similar phenomenon: a decisive gap in the emergence of good first-task-token performance vs second-task-token performance.

Student Paper: Yes

Submission Number: 69