Fine-Tuning Dynamics of In-Context Factual Recall in Transformers
Abstract
In-context learning — performing tasks based on examples given in the prompt — is an important capability that has emerged in large language models and has received significant attention in both theory and practice. Existing theoretical work on in-context learning often focuses on settings where learning uses information purely from the prompt. However, many practical cases also require recalling factual knowledge stored in model parameters, with the context identifying which knowledge is relevant. We study this behavior through the in-context factual recall (IC-recall) task, where a transformer receives a context of (subject, answer) pairs generated from a hidden relation and must infer the relation to answer a query subject. We model factual knowledge by the transformer having access to a simple pre-constructed MLP associative memory that stores (subject, relation, answer) triplets. We analyze the supervised fine-tuning dynamics of a one-layer transformer on IC-recall data, and prove that the model successfully performs IC-recall by converging to a particular pairwise attention pattern. This fine-tuning stage requires a very small number of samples — only polylogarithmic in the number of stored knowledge triplets. Experiments verify our theoretical predictions and show that the pairwise attention pattern emerges even when the MLP layer is pretrained instead of constructed.