🤖 machine learning

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

This paper demonstrates that while pretraining alone equips transformers with factual knowledge, finetuning on tasks requiring implicit inference is necessary to trigger the emergence of contextual recall by forming low-dimensional latent encodings of shared attribute types, a mechanism validated through both synthetic experiments and a constructed attention-only transformer model.

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

Published 2026-03-24

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant student named Transformer. This student has spent years reading every book in a massive library (this is Pretraining). They have memorized millions of facts: "The Eiffel Tower is in Paris," "The Colosseum is in Rome," "Niagara Falls is in Canada."

However, there's a catch. When the student reads these facts, they are always presented with a specific "label" or "hint" in the text, like a grammar tag.

Text: "The Eiffel Tower [Location: Paris]."
Text: "The Colosseum [Location: Rome]."

The student learns to associate the word "Eiffel Tower" with "Paris" only because they see the tag [Location] right before it. They are like a parrot that only repeats a phrase if it hears a specific whistle first.

The Problem: The "Silent" Test

Now, imagine a teacher gives the student a new kind of test. The teacher writes on the board:

"Niagara Falls, Canada. Colosseum, Italy. Parthenon, ______."

The teacher asks: "What goes in the blank?"

The student panics. There are no tags like [Location] or [Country] here. The student has to figure out the pattern themselves. They have to look at the first two examples, realize, "Oh! The pattern is 'Landmark, Country'," and then apply that logic to the third one to guess "Greece."

This is called Contextual Recall. The student must use the context (the examples) to recall a fact, without any explicit hints.

The Discovery: Reading isn't Enough

The researchers in this paper found something surprising:

Just reading the library (Pretraining) isn't enough. Even though the student knows all the facts, they fail the "Silent Test." They can't figure out the pattern because they are so used to relying on the "whistle" (the grammar tags).
They need a specific kind of practice (Finetuning). To fix this, the researchers gave the student a special training session. They didn't just give them more facts; they gave them puzzles where the "whistle" was missing or hidden, forcing the student to guess the pattern based on the examples.

The Magic: Learning to "Think" in Patterns

Here is the most exciting part of the paper. The researchers didn't need to teach the student new facts. They only taught them on a small subset of subjects (maybe just 50 landmarks out of 256).

But here's the magic: Once the student learned how to solve the puzzle for those 50 landmarks, they could instantly solve it for the other 206 landmarks they had never seen in this specific format before!

How did they do it?
Think of the student's brain as having a "sticky note" system.

Before training: The student looked for a specific sticky note (the grammar tag) to know what to do.
After training: The student learned to create a new, invisible "sticky note" inside their mind. When they see "Niagara Falls, Canada," they create a mental tag that says "This is a Country Game." They then use that invisible tag to find the country for "Parthenon."

The paper shows that during this special training, the student's brain forms a low-dimensional "Task Vector."

Analogy: Imagine the student's brain is a giant map. Before, all the facts were scattered randomly. After the special training, the student draws a new, invisible highway on the map that connects all "Country" facts together. When they see the first example, they jump onto this highway and zoom straight to the answer for the next one.

Why This Matters

This explains how AI models (like the ones you chat with) learn to be smart.

Pretraining gives them the raw data (the encyclopedia).
Finetuning teaches them the skill of looking at a few examples and figuring out the rules, even if the rules change or the hints disappear.

The paper proves that you don't need to re-teach an AI everything from scratch to make it smarter at reasoning. You just need to show it how to look for patterns in a few examples, and it will generalize that skill to everything it already knows.

Summary in One Sentence

Pretraining fills the AI's brain with facts, but a specific type of practice (finetuning) teaches the AI how to ignore the obvious hints and figure out the hidden rules on its own, allowing it to solve new puzzles using old knowledge.

1. Problem Definition

The paper investigates Contextual Recall, a specific form of In-Context Learning (ICL) where a model must recall specific facts about a query subject based on in-context examples, without explicit cues indicating which attribute type (e.g., "country" vs. "year built") is relevant.

The Challenge: Standard pretraining on open-ended text allows models to learn factual associations (e.g., "Parthenon is in Greece"). However, in a contextual recall scenario, the prompt format changes (e.g., [Niagara Falls, Canada, Colosseum, Italy, Parthenon, ]). The model must implicitly infer that the pattern is "Subject-Country" and predict "Greece" for Parthenon, despite never seeing facts presented as implicit pairs during pretraining.
Core Question: Does this ability emerge naturally from pretraining, or does it require specific finetuning? What are the underlying mechanisms?

2. Methodology

The authors introduce a controlled synthetic framework to isolate and study these phenomena.

Data Generation

The vocabulary consists of $M$ subjects, $L$ attribute types, $G$ grammar tokens, and a separator.

Pretraining (PT) Sequences: Designed to instill factual knowledge. They consist of tuples: [Subject, Grammar Sequence, Attribute]. Crucially, the grammar sequence is generated by a Markov chain specific to the attribute type. Thus, the grammar statistics implicitly encode the attribute type (e.g., a specific bigram pattern implies "Country").
ICL Sequences (Evaluation): Designed to test contextual recall. They consist of pairs: [Subject, Attribute, Subject, Attribute, ...]. No grammar tokens are present. The model must infer the attribute type solely from the pattern of the in-context examples.
ICL+Grm Sequences (Finetuning): A hybrid used for finetuning. They contain subject-attribute pairs but include short, variable-length grammar sequences. This forces the model to rely on context rather than fixed positional offsets or long grammar statistics.

Experimental Setup

Model: A 2-layer, single-head, decoder-only Transformer (and a 1-layer attention-only model for mechanistic analysis).
Protocol:
1. Pretrain on PT sequences to learn facts.
2. Evaluate on ICL sequences (held-out subjects).
3. Finetune on a subset of subjects using either ICL or ICL+Grm sequences.
4. Re-evaluate on ICL sequences with held-out subjects to test generalization.

3. Key Contributions & Findings

Finding 1: Pretraining is Insufficient

Models pretrained on PT sequences achieve high accuracy on factual recall (predicting attributes when grammar cues are present) but fail completely on ICL sequences (near-zero accuracy).

Implication: Pretraining teaches the model to decode attributes using explicit grammar statistics, but it does not teach the model to implicitly infer the attribute type from context alone.

Finding 2: Finetuning Enables Generalization

Finetuning the pretrained model on ICL sequences (or ICL+Grm sequences) using a subset of subjects enables the model to generalize to held-out subjects on ICL tasks.

Mechanism: The model learns a new "access mechanism." It stops relying on grammar statistics and learns to infer the attribute type by aggregating information from the in-context attribute tokens. Because the structural relationship between subjects and their attributes is shared across all subjects, learning this inference mechanism on a subset transfers to the whole set.

Finding 3: Robustness to Distribution Shift

Finetuning on ICL+Grm sequences (where grammar tokens are present but short/variable) also induces contextual recall.

The model must bridge the format gap (grammar present in training vs. absent in evaluation) and rely on context. This confirms that the key driver is the need for implicit inference, not just the specific prompt format.

Finding 4: Data Parameters Matter

Performance improves with:

Higher diversity in the Markov chains used for different attribute types during pretraining (making the grammar cues more distinct).
Longer grammar sequences during pretraining.
A larger number of subjects in the pretraining corpus.

Finding 5: Emergence of Low-Dimensional Task Vectors

Representational analysis reveals that finetuning induces the formation of low-dimensional latent encodings (task vectors) representing the shared attribute type.

As the number of in-context demonstrations increases, representations of sequences with the same attribute type cluster tightly together in the attention layer.
The separability of these clusters correlates with the diversity of the pretraining grammar statistics.

4. Mechanistic Analysis (Theoretical Insight)

The authors provide a theoretical construction for a single-layer, attention-only model to explain the transition from factual recall to contextual recall.

Pretraining Construction: The model uses a "Relation Head" to attend to the explicit relation token (or grammar statistics) and a "Subject Head" to filter attributes.
Finetuning Construction:
- The Subject Head remains largely unchanged (attending to the query subject).
- The Relation Head undergoes a critical transformation. Since explicit relation tokens are missing in ICL, the head learns to attend to the attribute tokens in the context. It learns to map the observed attributes to the set of all attributes of that type.
Training Dynamics: The paper proves that a few steps of gradient descent on the finetuning data can update the weights to replicate this construction. Specifically, the $W_{OV}$ (Output-Value) matrix learns attribute-to-attribute associations, and the $W_{KQ}$ (Key-Query) matrix learns to attend to context attributes rather than relation tokens.

5. Significance

Decoupling Knowledge and Reasoning: The paper demonstrates that factual knowledge (acquired via pretraining) and in-context reasoning capabilities (acquired via finetuning) are distinct. Pretraining provides the "database," while finetuning provides the "query engine" capable of implicit inference.
Mechanism of ICL: It provides a concrete mechanistic explanation for how Transformers perform ICL on factual tasks: by forming task vectors that encode the input-output relationship, allowing the model to steer its output based on context.
Generalization: It shows that models can generalize in-context reasoning to unseen subjects if the underlying structural relationship (the "task") is learned on a subset of data.
Implications for LLMs: This suggests that the "reasoning" capabilities of Large Language Models in novel formats may not be inherent to pretraining alone but are often unlocked or refined through specific finetuning regimes that require implicit inference.

Summary

The paper establishes that while pretraining equips Transformers with factual knowledge, it does not equip them with the ability to reason about that knowledge in novel, implicit formats. Finetuning is the catalyst that transforms the model's internal representations, enabling it to form low-dimensional task vectors and implicitly infer attribute types from context, thereby achieving robust contextual recall and generalization to unseen data.