Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledge

Bhavya Vasudeva; Puneesh Deora; Alberto Bietti; Vatsal Sharan; Christos Thrampoulidis

cs.LGMar 21, 2026

Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledge

Authors: Bhavya Vasudeva, Puneesh Deora, Alberto Bietti, Vatsal Sharan, Christos Thrampoulidis

Organizations: University of Southern California · University of British Columbia · Flatiron Institute

Abstract

Transformer-based language models excel at in-context learning (ICL), where they can adapt to new tasks based on contextual examples, without parameter updates. In a specific form of ICL, which we refer to as contextual recall, models pretrained on open-ended text leverage pairwise examples to recall specific facts in novel prompt formats. We investigate whether contextual recall emerges from pretraining alone, what finetuning is required, and what mechanisms drive the necessary representations. For this, we introduce a controlled synthetic framework where pretraining sequences consist of subject-grammar-attribute tuples, with attribute types tied to grammar statistics. We demonstrate that while such pretraining successfully yields factual knowledge, it is insufficient for contextual recall: models fail to implicitly infer attribute types when the grammar statistics are removed in ICL prompts. However, we show that finetuning on tasks requiring implicit inference, distinct from the ICL evaluation, using a subset of subjects, triggers the emergence of contextual recall across all subjects. This transition is accompanied by the formation of low-dimensional latent encodings of the shared attribute type. For mechanistic insight, we derive a construction for an attention-only transformer that replicates the transition from factual to contextual recall, corroborated by empirical validation.

Figures & tables

Figure 1: Illustration of the data generation process and the evaluation metrics. The vocabulary consists of

M

subjects with

L

attributes each (where the attribute tokens may be shared across the subjects),

G

grammar tokens and a separator token. We sample

L

Markov chains, one for each attribute type. We consider three types of sequences (see Section 2 for details): PT sequences contain subject-grammar-attribute tuples, with attribute type information encoded in the grammar sequence statistics, for one subject, ICL sequences contain pairwise subject-attribute examples for a shared attribute type, but no grammar, and ICL+Grm sequences are analogous to ICL sequences, but contain subject-grammar-attribute tuples. We evaluate whether the model’s prediction (on PT sequences) matches the ground truth attribute (exact match), any attribute of the same type as the last subsequence, any attribute of the same subject as the sequence, or any attribute token.

Figure 2: Transformer trained on PT sequences performs well on PT sequences, but does not generalize to ICL sequences.

Figure 3: Finetuning the model pretrained on PT sequences, using ICL sequences with a subset of subject s, leads to good performance on ICL sequences with held-out query subject s.

Figure 4: Finetuning the model pretrained on PT sequences, using ICL+Grm sequences with short, variable grammar length with a subset of subject s, leads to good performance on ICL sequences with held-out query subject s.

Panel of Figure 5 — Figure 5: Increasing the number of finetuning subject s with ICL (left) or ICL+Grm (right) sequences improves the model’s performance on ICL sequences with held-out query subject s.

Figure 6: Increasing the sequence length $S$ and/or Markov chain diversity Div used while pretraining improves the model’s performance on PT sequences as well as on ICL sequences with held-out subject s after finetuning with ICL+Grm sequences.

Figure 7: Increasing the number of subject s used while pretraining improves the model’s performance on ICL sequences with held-out subject s after finetuning with ICL+Grm sequences.

Panel of Figure 8 — Figure 8: Comparison of accuracy on ICL sequences with held-out subjects, and layer-2 attention representation clustering strength for a finetuned model, as the number of demonstrations is increased (left), and ranks of the original and target attribute type when using those task representations for steering model output (right; see text for details), using $\texttt{Div}\approx 0.2$ (top) and $\texttt{Div}\approx 0.5$ (bottom). Inset figures visualize the averaged cosine similarity for inter- and intra-task representations. In both cases, both performance and clustering strength improve with number of demonstrations.

Figure 9: Attention scores for each head across a PT sequence at the end of pretraining. Head $2$ attends to the most recent subject, while head $0$ attends to the most recent relation token, as predicted by Proposition 1 .

Figure 10: Cosine similarity between the actual outputs of each head and the outputs $f_{\text{subj}}({\bm{X}})$ (left) and $f_{\text{rel}}({\bm{X}})$ (right) from our construction in Proposition 1 . Head $2$ output closely matches $f_{\text{subj}}({\bm{X}})$ , while head $0$ matches $f_{\text{rel}}({\bm{X}})$ .

Figure 11: Attention scores for each head across an ICL sequence at the end of finetuning. Head $2$ attends to the most recent subject, while head $0$ attends to the attribute tokens in the sequence, as predicted by Proposition 2 .

Figure 12: Cosine similarity between the actual outputs of each head and the outputs $f_{\text{subj}}({\bm{X}})$ (left) and $f_{\text{rel}}({\bm{X}})$ (right) from our construction in Proposition 2 . Head $2$ output closely matches $f_{\text{subj}}({\bm{X}})$ , while head $0$ matches $f_{\text{rel}}({\bm{X}})$ .

Figure 13: Effect of number of unique attributes seen during finetuning (with ICL sequences) on the accuracy on ICL sequences with held-out query subjects, across different number of PT and FT subjects. Performance improves as the number of unique attributes seen while finetuning increases.

Panel of Figure 14 — Figure 14: Using 256 (top row) or 2048 (bottom row) PT subjects, (left) after pretraining, the cosine similarities of the output of each head with the subject and relation head outputs predicted by the construction in Proposition 1 , and (right) after FT with 32 subjects, the performance on ICL sequences with held-out query subjects. In the first case, the cosine similarities for subject and relation head are high, and after FT, the subject-match accuracy is high, while attribute-type match accuracy is lower because all attributes were not seen during FT. In the second case, the cosine similarities for subject and relation head are low, and consequently, after FT, the subject-match accuracy is low, as well as the attribute-type match accuracy is low because all attributes were not seen during FT.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 15: Performance of the model when finetuning with ICL+Grm sequences with a subset of subjects (same setting as Fig. 4 ), on ICL+Grm sequences with $S_{\text{ft}}=1$ (left) and ICL sequences (right) with seen or held-out subjects. Finetuning with ICL+Grm sequences enables out-of-distribution generalization on ICL sequences with held-out subjects. Performance improves first on ICL+Grm and later on ICL sequences.

Panel of Figure 16 — Figure 16: Comparison of clustering strength (see Section 3.4 for details) using representations from different layers of the finetuned 2-layer 1-head transformer with ICL sequences as inputs, as the number of demonstrations is increased (same setting as Fig. 8 ) with $\texttt{Div}\approx 0.2$ (left) and $\texttt{Div}\approx 0.5$ (right). We find that layer-2 attention representations cluster most strongly, while layer-1 representations exhibit no clustering based on the attribute type information in the context. Further, using higher Div while pretraining leads to stronger representation clustering after finetuning.

Panel of Figure 17 — Figure 17: Each subfigure visualizes the cosine similarity for inter- and intra-task representations (from layer-2 attention layer of the finetuned model), $\bar{C}^{t}(\ell,\ell^{\prime})$ (see Section 3.4 for details) across attribute types $\ell,\ell^{\prime}\in[L]$ , averaged over $50$ sequences for each attribute type (same setting as Fig. 8 ; top: $\texttt{Div}\approx 0.2$ , bottom: $\texttt{Div}\approx 0.5$ ). We find that as the number of demonstrations is increased (from $0$ to $15$ ), the representations of ICL sequences with the same attribute type get clustered together, and the clustering is stronger for higher Div ( i.e. , more separated attribute-specific Markov shains).

Panel of Figure 18 — Figure 18: Comparison of rank of the original and target attribute type when steering the model output using the averaged layer-2 attention representations of the finetuned 2-layer 1-head transformer with ICL sequences as inputs at the last token position (see Section 3.4 for details) with $\texttt{Div}\approx 0.2$ (left) and $\texttt{Div}\approx 0.5$ (right), for seen and held-out FT subjects, for each target attribute type. We find that higher Div used while pretraining leads to stronger task vectors and better steering with lower $\beta$ values.

Figure 19: Finetuning the single-layer attention-only model with relative positional encoding pretrained on PT sequences using ICL sequences with a subset of subjects enables generalization on ICL sequences with held-out query subjects.

Figure 24: Finetuning the single-layer attention-only model with absolute positional encoding pretrained on PT sequences using ICL sequences with a subset of subjects enables generalization on ICL sequences with held-out query subjects.

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 404 papers

Fine-Tuning Dynamics of In-Context Factual Recall in Transformers
May 26, 2026Ruomin Huang, Eshaan Nichani, Jason D. Lee +1Supervised Fine-Tuning Transformer Attention
Duke University · Princeton University · UC Berkeley
When Thoughts Become Memory: Unifying In-Context Learning and Fine-Tuning
Oct 21, 2026Leon Chlon, Julie HuangIn-Context Learning Data Valuation
Hassana Labs.
When Context Sticks: Studying Interference in In-Context Learning
Apr 25, 2026Hanna Rød, Dagny Streit, Nils Valseth Selte +1In-Context Learning

May 26, 2026J/K move · Enter open · S save

Fine-Tuning Dynamics of In-Context Factual Recall in Transformers

Ruomin Huang, Eshaan Nichani, Jason D. Lee, Rong Ge

Duke University · Princeton University · UC Berkeley

Supervised Fine-Tuning Transformer Attention In-Context Learning Factual Knowledge in Language Models

Open paper AskPDF arXiv

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 often focuses on settings where the learning uses information purely from the prompt. However, many practical instances of in-context learning require the model to retrieve factual knowledge stored in the model's parameters, with the context serving to identify which knowledge is relevant. In this work, we study how in-context learning leverages factual knowledge recall. We formalize this behavior by introducing the in-context factual recall (IC-recall) task, where a transformer is provided a context of (subject, answer) pairs generated from a hidden relation, along with a query subject, and must both infer this hidden relation and retrieve the corresponding answer. Factual knowledge is modeled by the transformer having access to a simple pre-constructed MLP associative memory storing (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.

Abstract

In-Context Learning

Explore similar work

Fine-Tuning Dynamics of In-Context Factual Recall in Transformers

When Thoughts Become Memory: Unifying In-Context Learning and Fine-Tuning

When Context Sticks: Studying Interference in In-Context Learning