💻 computer science

Decomposing Prediction Mechanisms for In-Context Recall

This paper introduces a novel toy problem combining continuous in-context learning with discrete associative recall to demonstrate that transformer models employ two distinct mechanisms with different learning dynamics: an early-emerging, label-agnostic "Bayesian-style" predictor for continuing sequences and a later-emerging, label-dependent mechanism for recalling and resuming specific prior states.

Original authors: Sultan Daniels, Dylan Davis, Dhruv Gautam, Wentinn Liao, Gireeja Ranade, Anant Sahai

Published 2026-06-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Sultan Daniels, Dylan Davis, Dhruv Gautam, Wentinn Liao, Gireeja Ranade, Anant Sahai

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

The Big Idea: Two Different Brains in One Model

Imagine you have a very smart robot assistant. You give it a notebook filled with stories about different people (let's call them "systems"). Sometimes the story is about Alice, sometimes Bob, and sometimes Charlie. The notebook uses special colored brackets to tell the robot when the story switches from Alice to Bob.

The robot's job is to read the story and predict what happens next. The researchers wanted to see how the robot learns to do this. They discovered something surprising: The robot doesn't use one single way of thinking. It uses two completely different "modes" or "brains" depending on which word it is trying to predict next.

The Setup: The "Haystack" Game

To test this, the researchers created a game called "Needle in a Haystack."

The Haystack: A long list of mixed-up story segments. One minute it's Alice's story, the next it's Bob's, then back to Alice, then Charlie.
The Labels: Each story segment starts with a unique "Open" label (like a red bracket [) and ends with a "Close" label (like a red bracket ]).
The Test: The robot reads the whole haystack. Then, the researchers show it a specific "Open" label (e.g., "Bob's bracket") and ask, "What comes next in Bob's story?"

The Two "Brains" (Mechanisms)

The paper found that the robot solves this puzzle using two distinct methods that turn on at different times during its training.

1. The "Label Reader" (The First Brain)

When it kicks in: This brain wakes up late in the robot's training.
How it works: It looks at the "Open" label (the red bracket) and says, "Ah, this is Bob's story! I remember Bob's rules from earlier in the notebook." It then uses those specific rules to predict the very first word of the new segment.
The Analogy: Imagine you are at a party. You see a name tag that says "Bob." You instantly remember, "Oh, Bob always tells jokes." You use that memory to predict Bob's first sentence. This is Associative Recall.

2. The "Pattern Follower" (The Second Brain)

When it kicks in: This brain wakes up early in the training.
How it works: It ignores the name tags (labels) almost entirely. Instead, it looks at the last few words it just read and says, "Based on the pattern of the last few words, the next one should be X." It uses a "Bayesian" style of guessing (looking at probabilities based on recent history).
The Analogy: You are walking down a street. You don't need to know the person's name to guess what they will do next; you just watch their footsteps. If they are walking fast, you guess they will keep walking fast. This is Pattern Prediction.

The Surprising Discovery: They Don't Turn On Together

The most important finding is that these two brains develop at different speeds and use different parts of the robot's "wiring."

The Pattern Follower (Brain 2) learns first. The robot gets really good at predicting the second, third, and fourth words of a story very early on. It does this by just looking at the immediate previous words, ignoring the labels.
The Label Reader (Brain 1) learns later. It takes a long time for the robot to learn that the "Open" label is a key to unlocking the first word of a new story.

The "Phase Transition":
Imagine a light switch. For a long time, the robot is bad at predicting the first word after a label. Then, suddenly, around a specific point in training, it flips a switch and gets perfect at it. But the ability to predict the second word had already been working smoothly for a long time.

How They Proved It (The Experiments)

The researchers played tricks on the robot to see which brain it was using:

The "Wrong Label" Trick: They showed the robot a label for "Bob" but the story was actually "Charlie's."
- Result: The robot failed at predicting the first word (because it was trusting the wrong label). But, it was still good at predicting the second word (because it was just following the pattern of the words it saw).
- Conclusion: The first word relies on the Label; the second word relies on the Pattern.
The "Confusing Label" Trick: They showed the robot a label for a story that didn't exist in the notebook.
- Result: At first, the robot ignored the fake label and kept predicting the story it was already in. But later in training, once the "Label Reader" brain woke up, the robot got confused and tried to start a new story based on the fake label, even though it was wrong.
- Conclusion: The robot learned to trust the labels too much, even when they were nonsense.
The "Circuit Breaker" Test: They physically "pruned" (cut) parts of the robot's neural network.
- Result: They found that the wires used to predict the first word were completely different from the wires used to predict the second word. They are two separate machines working inside one body.

Does This Happen in Real AI?

The researchers checked if this "two-brain" phenomenon happens in real Large Language Models (like the ones used for translation). They found the same thing:

The model gets good at continuing a translation (predicting the second word) before it gets good at starting a new translation based on a prompt (predicting the first word).
This suggests that even our most advanced AI models are juggling multiple, separate strategies to solve a single task.

Summary

This paper shows that when an AI learns to recall information from a context, it doesn't just "remember" things. It uses a Label-Based strategy to start a new task (which takes a long time to learn) and a Pattern-Based strategy to continue a task (which it learns quickly). These two strategies are so different that they live in separate parts of the AI's brain and turn on at different times.

Technical Summary: Decomposing Prediction Mechanisms for In-Context Recall

Problem Definition

This paper introduces a new family of toy problems designed to bridge the gap between continuous in-context learning (ICL) tasks (such as linear regression) and discrete associative recall tasks (such as Multi-Query Associative Recall, or MQAR). The core task involves predicting the state of a discrete, interleaved time-series system.

The underlying data consists of noise-free trajectories generated by random orthogonal linear dynamical systems ( $x_{i+1} = U x_i$ , where $U$ is a random orthogonal matrix). These trajectories are segmented and interleaved within a single context window, demarcated by discrete "symbolic punctuation labels" (open and close tokens). The model is trained to perform next-token prediction on these interleaved traces. The challenge requires the model to:

Identify which specific system's state should be recalled based on a symbolic label.
Apply the correct underlying linear dynamics to that state to predict subsequent observations.
Continue the prediction for the resumed sequence.

The authors posit that successful performance requires distinguishing between two distinct capabilities: learning the dynamics of a system as it is first seen (ICL) and recalling the dynamics of a system seen previously in the context (associative recall).

Methodology

Data Generation

The authors generate a training library of 40,000 orthogonal matrices ( $U \in \mathbb{R}^{5 \times 5}$ ) and corresponding initial states. Training traces are constructed by interleaving segments of these time-series.

Interleaving: Segments are cut and reassembled with random "open" and "close" symbolic labels assigned to each system instance. Crucially, the mapping between a system and its symbolic label is randomized for every training example, preventing the model from memorizing static label-to-system associations.
Distribution: The number of systems per trace follows a Zipf distribution, and the number of cuts follows a Poisson distribution, creating a "needle-in-a-haystack" scenario where the model must retrieve a specific system from a context containing multiple others.

Model Architecture and Training

The authors train GPT-2 style Transformers (ranging from 212K to 10.7M parameters) on next-token prediction using Mean Squared Error (MSE) loss. The input embedding dimension is 57 (50 for symbolic labels, 1 for start, 1 for payload flag, 5 for state vectors). The output layer predicts the 5-dimensional state vector.

Evaluation Protocols

Uninterleaved Test: Evaluates the model's ability to learn the dynamics of the first system seen in the context.
Needle-in-a-Haystack (NiH) Test: Evaluates associative recall. A "haystack" of $N$ systems is presented, followed by a query symbol indicating which system to resume ("the needle"). Performance is measured on the first token after the query (1-after) and subsequent tokens (2-after, etc.).
Out-of-Distribution (OOD) Experiments: Four specific experiments are designed to test the mechanism of recall:
- Misdirection to Incorrect Sequence: Swapping the query label to point to a different system in the haystack.
- Synchronized Rotations: Forcing all systems in the haystack to have the same state at the first index of the test segment, removing state-based disambiguation.
- Misdirection to Unseen Sequence: Using a query label that corresponds to a system not present in the haystack.
- Misdirection to Seen Sequence (Restart): Using a label pointing to an old system while the actual test segment is a new unseen system.
Circuit Analysis: Edge pruning is applied to the trained weights to isolate the specific computational circuits used for predicting the "1-after" token versus the "2-after" token.

Key Contributions and Results

1. Distinct Learning Dynamics for ICL and Recall

The training dynamics reveal a clear separation in when different capabilities emerge:

First System Learning: The ability to learn the dynamics of the first system seen in the context develops gradually throughout training, converging toward the optimal pseudoinverse predictor.
Restarting ICL: The ability to restart ICL for a new system (after seeing other systems) emerges relatively early but still precedes full associative recall.
Associative Recall: The ability to recall a previously seen system based solely on its symbolic label emerges abruptly (a phase transition) much later in training, specifically for the first token after the query.

2. The Multi-Mechanism Hypothesis (C3)

The central finding is that the model does not use a single mechanism for the entire task. Instead, it employs two distinct mechanisms that emerge at different times and utilize different information:

Mechanism 1 (Label-Based): Used exclusively for predicting the first token after a query symbol. This mechanism relies on the discrete symbolic label to retrieve the correct system state. Evidence: When the label is misdirected to an incorrect system, the error for the first token spikes immediately.
Mechanism 2 (Observation-Based/Bayesian): Used for predicting the second token and beyond. This mechanism largely ignores the symbolic label and instead uses the observed state of the first token to infer the underlying system dynamics (a "Bayesian-style" prediction). Evidence: When the label is misdirected, the error for subsequent tokens remains low because the model infers the correct system from the observed state.

3. Circuit Disentanglement

Edge pruning confirms that these two mechanisms are physically distinct in the model's weights. The circuits responsible for predicting the "1-after" token and the "2-after" token have 0% edge overlap, indicating that the model learns separate sub-networks for initiating a task episode versus continuing it.

4. Generalization to Natural Language

The authors validate this phenomenon on a natural language task (English-to-Spanish translation) using OLMo checkpoints. They observe a similar phase transition: the model's ability to initiate a translation based on a symbolic task label (in-context associative recall) emerges later and more abruptly than its ability to continue the translation once started. When semantic labels (e.g., "Spanish:") are used instead of symbolic ones (e.g., "X:"), the gap disappears, confirming that the delay is specific to the in-context learning of the task definition.

Significance and Claims

The paper claims to provide a mechanistic explanation for how Transformers perform complex in-context tasks, challenging the view of ICL as a monolithic capability.

Decomposition of ICL: The authors argue that "in-context learning" is not a single skill but a composite of distinct abilities (task recognition vs. task continuation) that emerge sequentially.
Local Underspecification: The paper suggests that because autoregressive training minimizes average loss over tokens, the model can succeed by using a "good enough" mechanism (observation-based inference) for the majority of tokens (2nd, 3rd, etc.), which reduces the gradient pressure to learn the more complex mechanism (label-based retrieval) required for the first token.
Phase Transitions: The results support the existence of distinct phase transitions in learning dynamics, where the emergence of one capability (continuation) does not guarantee the immediate emergence of another (initiation), even when the information theoretically required is similar.
Mechanistic Interpretability: The work demonstrates that edge pruning can successfully isolate these distinct functional circuits, providing a tool for understanding how different parts of a task are handled by different parts of the network.

The authors conclude that their findings offer a new dimension of nuance to the understanding of ICL, suggesting that for multi-token tasks, models may simultaneously operate in different modes (label-based vs. observation-based) depending on the position in the sequence.