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 LLMs 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 \emph{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.

• The paper demonstrates that fine-tuning induces a novel pairwise attention pattern that enables precise retrieval of factual knowledge in in-context learning tasks.
• It employs chain-of-thought supervision where explicit intermediate relation decoding is crucial for achieving near-perfect answer prediction with minimal samples.
• The study uses a one-layer transformer with a fixed MLP associative memory, offering insights for optimizing model design in knowledge-intensive applications.

Fine-Tuning Dynamics of In-Context Factual Recall in Transformers: An Expert Analysis

Motivation and Problem Formulation

The paper "Fine-Tuning Dynamics of In-Context Factual Recall in Transformers" (2605.27774) systematically investigates how transformer models leverage parameterized factual knowledge during in-context learning (ICL), specifically in settings where the context alone cannot suffice for task resolution. Here, factual recall is operationalized via the novel IC-recall task: the model receives (subject, answer) pairs generated from a hidden relation and a query subject; it must infer the hidden relation and subsequently retrieve the corresponding answer.

The factual knowledge corpus is formalized as triplets (subject, relation, answer), mapped into distinct subject, relation, and answer sets. The task construction ensures that context alone is insufficient, requiring the model to combine prompt information with memorized knowledge for successful completion.

Figure 1: IC-recall data construction, showing subject, relation, and answer sets, and illustrating the prompt format required for the task.

Model Architecture and Factual Memory Encoding

The model comprises a one-layer transformer with a single self-attention head, followed by an MLP associative memory layer. The MLP is constructed (or pretrained) to encode the factual knowledge corpus, acting as a key–value store for triplet retrieval, and remains fixed during fine-tuning.

Key architectural assumptions include:

• Orthogonal embeddings for all vocabulary tokens (subjects, relations, answers, EoS).
• One-hot positional encodings concatenated to token embeddings.
• Quadratic activation in the MLP for analytical tractability.

The associative memory construction guarantees perfect retrieval accuracy via argmax decoding whenever any two triplet elements are provided, with MLP width scaling as $O(|\mathcal{S}| \cdot |\mathcal{A}|)$.

Fine-Tuning Dynamics and Pairwise Attention Emergence

The fine-tuning protocol involves a chain-of-thought (CoT) objective: in the first decoding step, the transformer predicts the hidden relation; in the second step, it predicts the answer given the decoded relation and query subject.

Theoretically, it is proven that after fine-tuning, the transformer converges to a pairwise attention pattern: the EoS token distributes attention equally across the two subject–answer pairs in the context but assigns different weights to the pairs, enabling distinction between the correct and "confusing" relations. Crucially, this pattern diverges from conventional induction head mechanisms, revealing a qualitatively distinct "pairwise" solution that optimally performs IC-recall. Two-stage dynamics are observed: initial convergence to a saddle point (uniform distribution across pairs), followed by symmetry-breaking via random perturbation and subsequent gradient descent, yielding the pairwise attention solution.

Figure 2: The emergence of pairwise attention in fine-tuned transformers for the IC-recall task.

Sample Efficiency and Scaling Laws

A principal numerical result is that the number of fine-tuning samples required–to achieve high test accuracy on IC-recall tasks–grows polylogarithmically with the number of stored knowledge associations. For models with $|\mathcal{S}| = n$ subjects, $8$ samples suffice for empirical accuracy approaching $99.9\%$. This scaling holds both when the MLP is constructed and when it is pretrained, as demonstrated by experimental verification.

Figure 3: Test accuracies for fine-tuning on the first decoding step as a function of the number of subjects.

Chain-of-Thought Supervision and Generalization

Chain-of-thought supervision (explicit intermediate relation decoding) is essential for effective learning dynamics and task solvability. Without strong supervision, accuracy for answer prediction cannot exceed $1/3$, even in minimal settings. This finding aligns with broader literature on the sample efficiency and expressive power improvements provided by intermediate reasoning steps in transformers.

The pairwise attention pattern generalizes empirically to cases with more than two in-context examples, though formal analysis is limited to the two-example scenario due to symmetry and loss landscape tractability.

Figure 4: Attention scores after fine-tuning with different numbers of in-context examples, highlighting the pairwise structure.

Implications and Future Directions

Practical Implications

The paper rigorously characterizes fine-tuning protocols that allow models to optimally leverage parameterized factual memory in few-shot settings. The findings underscore the critical importance of attention structure and CoT supervision in factual recall, with direct implications for prompt engineering and model design in knowledge-intensive applications.

Theoretical Implications

The results spotlight the nuanced mechanisms beyond induction heads that underpin in-context learning with stored knowledge. The emergence and optimality of pairwise attention call for further investigation into architectural biases and loss surface geometry, particularly as models scale and tasks become more compositional.

Future Directions

Extensions of this analysis to multi-layer/pretrained transformer architectures with non-orthogonal embedding spaces, dynamic associative memory formation, and end-to-end training are warranted. It remains to be explored how these dynamics manifest in large-scale, real-world LLMs, and whether pairwise attention can be explicitly targeted for enhanced factual recall or transfer learning across tasks.

Conclusion

This work provides a detailed theoretical and empirical analysis of fine-tuning dynamics in factual recall tasks for transformers equipped with associative memory. The emergence of pairwise attention patterns, the necessity of chain-of-thought supervision, and the strong sample efficiency results offer precise guidance for model design and training strategies in factual knowledge-intensive settings. The analysis sets the stage for deeper investigations into architectural mechanisms that govern in-context learning and knowledge retrieval in contemporary AI systems.