In-Context Learning as Conditioned Associative Memory Retrieval

Introduction

Despite the massive success of Large Language Models (LLMs), In-Context Learning (ICL)—the ability of a model to learn a new task from a few examples in a prompt without permanent weight updates—remains somewhat of a "black box." Existing theories often view ICL abstractly, either as implicit gradient descent in the forward pass or as a form of Bayesian inference.

This paper provides a mechanistic, neuroscience-inspired explanation for ICL by framing it as "memory reshaping" within a Modern Hopfield Model (MHM). It addresses a fundamental question: How does a fixed-weight transformer physically implement a new behavior on the fly? By connecting transformer attention to associative memory, the authors show that prompt examples act as cues that temporarily warp the model's internal energy landscape, biasing retrieval toward task-relevant information.

A Mechanistic View of Attention

The researchers analyze an exactly solvable case using a one-layer attention model. Their approach hinges on bridging the mathematical equivalence between the softmax attention mechanism and the update dynamics of a Modern Hopfield Network. In this view, the "keys" and "values" in a transformer act as stored memories, while the "query" serves as a search cue.

Memory Reshaping

The core theoretical contribution is proving that providing $t$ in-context examples induces an implicit linear transformation on the model’s pre-trained attention weights. Essentially, the prompt "reshapes" the energy landscape of the model's memory. It moves task-relevant "attractors" (local minima in the energy function) closer to the query while pushing irrelevant ones away. ICL doesn't teach the model new facts; rather, it dynamically re-weights the importance of memories already stored in its weights, acting as a filter that highlights the right neighborhood of internal knowledge.

To fairly compare models with and without prompts, the authors introduce "ghost examples"—mathematical constructs representing the model's base state before it is conditioned by a real prompt. They then bridge this mechanistic view with the Bayesian model averaging framework, showing that memory reshaping is the physical manifestation of the model performing a conditional mean estimation of a hidden task variable.

Experimental Validation and Key Results

The authors validated their theory using synthetic tasks (linear regression, decision trees, and 2-layer neural networks) on a GPT-2 architecture, as well as a real-world sentiment classification task using GPT-J.

Theoretically, the authors recovered the known $\epsilon$-stability generalization bound for ICL within their one-layer attention framework, confirming that their memory-retrieval interpretation is mathematically consistent with established statistical learning theory.

Empirically, they successfully explained and validated three key ICL behaviors through sensitivity analysis:

1. Covariance Shift: ICL performance drops when the test query comes from a different distribution than the prompt examples, as the memory is reshaped incorrectly.
2. Label Accuracy: Incorrect labels in the prompt mislead the reshaping process, significantly degrading performance.
3. Prompt-Query Similarity: Higher similarity between prompt examples and the test query leads to better retrieval, as the query lands closer to the reshaped attractors.

Expand

Source, p. 7

As the number of prompt examples increases, the model's internal hidden state (the retrieved memory) becomes more stable and robust, converging on the correct task representation.

Limitations and Future Directions

While this framework provides a compelling look into the mechanics of ICL, it has a few limitations:

• One-Layer Limitation: The current mathematical proof is restricted to one-layer attention. Real LLMs have dozens of layers. The authors suggest that multi-layer transformers perform iterative retrieval, though the math for this is significantly more complex.
• Proxy Attention: The "Bayesian model averaging attention" used in the proofs is an effective proxy but doesn't perfectly match the exact implementation of attention in every production LLM.
• Perfect Pre-training Assumption: The theory assumes the model was perfectly pre-trained on the relevant task distribution, which may not always hold for out-of-distribution tasks.

Future work will aim to extend this framework to multi-layer architectures and explore how phenomena like "memory hijacking"—where prompts distract or mislead the model—can be explained through this energy-landscape lens.