Decomposing Prediction Mechanisms for In-context Recall

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 to explore challenges with long context learning and associative recall in transformer models. Our setup involves interleaved segments of observations from randomly drawn linear deterministic dynamical systems. Each system is associated with a discrete symbolic label that must be learned in-context since these associations randomly shuffle between training instances.

Via out-of-distribution experiments we find that learned next-token prediction for this toy problem involves at least two separate mechanisms. One "label-based" mechanism uses the discrete symbolic labels to do the associative recall required to predict the start of a resumption of a previously seen system's observations. The second ``observation-based'' mechanism largely ignores the discrete symbolic labels and performs a prediction based on the state observations previously seen in context. These two mechanisms have different learning dynamics: the second mechanism develops much earlier than the first.

The behavior of our toy model suggested concrete experiments that we performed with OLMo training checkpoints on an ICL translation task. We see a similar phenomenon: the model learns to continue a translation task in-context earlier than it decisively learns to in-context identify the meaning of a symbolic label telling it to translate.

Primary Area: interpretability and explainable AI

Submission Number: 23149