💻 computer science

Mixing Mechanisms: How Language Models Retrieve Bound Entities In-Context

This paper reveals that while language models initially rely on a positional mechanism to retrieve bound entities in-context, this approach becomes unreliable as context complexity increases, prompting models to supplement it with lexical and reflexive mechanisms to form a robust, causal model that accurately predicts behavior across diverse and longer inputs.

Original authors: Yoav Gur-Arieh, Mor Geva, Atticus Geiger

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yoav Gur-Arieh, Mor Geva, Atticus Geiger

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 Picture: The "Lost in the Middle" Problem

Imagine you are reading a story where a character introduces a long list of friends and their favorite foods:

Ann loves pie.
Joe loves jam.
Pete loves tea.
Tim loves ale.
...and so on, for 20 different people.

Then, the story asks: "Who loves ale?"

If you are a human, you might scan the list. If you are an AI (Language Model), you have to figure out how to find the answer inside your memory of that story.

For a long time, researchers thought AI solved this by simply counting positions. It would think: "The question is about the 4th person mentioned, so I'll look at the 4th name." This is like a librarian who only knows books by their shelf number.

This paper says that librarian is wrong. The AI doesn't just use shelf numbers. In fact, when the list gets long, the "shelf number" method gets confused, especially for people in the middle of the list. To fix this, the AI actually uses three different tricks at the same time to find the right answer.

The Three Tricks (Mechanisms)

The authors discovered that the AI mixes three specific strategies to remember who likes what. Think of it like a detective solving a case using three different tools:

1. The Positional Mechanism (The "Shelf Number" Trick)

How it works: The AI looks at where the information appeared in the text. If "Ann" was the first person mentioned, the AI remembers "First person."
When it works: This is great for the very first and very last items in a list. It's like remembering the first and last books on a shelf perfectly.
When it fails: As the list gets longer, the "middle" of the shelf gets blurry. The AI gets confused about whether "Tim" is the 5th or 6th person. The signal becomes "noisy" and unreliable.

2. The Lexical Mechanism (The "Word Association" Trick)

How it works: Instead of counting, the AI looks at the words themselves. If the question is "Who loves pie?", the AI searches for the word "pie" in its memory and grabs the name attached to it.
The Analogy: It's like a waiter who doesn't remember table numbers but remembers: "Oh, the person who ordered the pie is the one in the red hat."
When it helps: This is very sharp and precise. It saves the day when the "shelf number" gets blurry in the middle of the list.

3. The Reflexive Mechanism (The "Self-Pointer" Trick)

How it works: This is the most unique one. When the AI reads "Tim loves ale," it creates a secret, invisible arrow pointing from the word "ale" back to "Tim." Later, when asked about "ale," it follows that arrow directly.
The Analogy: Imagine "ale" has a tiny GPS tracker attached to it that points directly to "Tim." The AI doesn't need to search; it just follows the tracker.
Why it's needed: Sometimes the question comes before the answer in the sentence structure (e.g., "Who loves ale? Tim does."). The "Word Association" trick can't work backwards easily, so the AI needs this direct "pointer" to jump straight to the answer.

The "Middle" Problem and the Solution

The paper found a fascinating pattern:

At the start and end of the list: The AI relies mostly on the Positional trick (counting). It's strong and confident here.
In the middle of the list: The Positional trick gets weak and fuzzy (this is called the "Lost in the Middle" effect).
The Fix: In the middle, the AI switches to a mix of the Lexical (word association) and Reflexive (pointer) tricks. These are much sharper and help the AI ignore the noise of the long list.

The authors proved this by "patching" the AI's brain. They took the AI's memory from one story and swapped it with the memory from a different story.

If they swapped the Positional memory, the AI guessed based on the new list's order.
If they swapped the Lexical memory, the AI guessed based on the new list's words.
If they swapped the Reflexive memory, the AI followed the new pointers.

By doing this, they showed that the AI isn't using just one method; it's constantly blending all three to get the right answer.

The "Causal Model" (The Recipe)

The researchers built a simple mathematical model that combines these three tricks.

They found that if you tell the model to use all three (Positional + Lexical + Reflexive), it can predict what the AI will say next with 95% accuracy.
If they tried to use only the old "Positional" method (the way everyone thought it worked), the accuracy dropped to about 45%—basically guessing.

Does this work in the real world?

The paper tested this not just on simple lists, but on stories with "filler" text (random sentences about the weather or traffic) inserted between the facts.

Even with all that extra noise, the AI still used the same three tricks.
However, as the text got very long, the "Word Association" (Lexical) trick got a little weaker, and the "Shelf Number" (Positional) trick got a little noisier. This explains why AI sometimes struggles with very long documents—it's not that the AI forgot; it's that its three tools are getting slightly less precise as the list grows.

Summary

Language models don't just count their way through a story. They use a hybrid system:

Counting (Positional) for the edges.
Matching words (Lexical) for the middle.
Following pointers (Reflexive) for direct connections.

By understanding this mix, we get a much clearer picture of how AI actually thinks and remembers things in long conversations.

Technical Summary: Mixing Mechanisms in In-Context Entity Retrieval

Problem Statement

Language models (LMs) excel at in-context reasoning, a capability fundamentally reliant on the ability to bind related entities within a text (e.g., binding "Ann" to "pie" in "Ann loves pie") and subsequently retrieve them based on a query (e.g., answering "Who loves pie?"). While prior research established that LMs utilize a positional mechanism—retrieving entities based on their index in the context sequence—this view was largely derived from studies involving short lists of entity groups (e.g., $n=2$ or $3$).

The central problem addressed in this work is the reliability of this positional mechanism in more complex, realistic settings. As the number of bound entity groups increases, LMs struggle with reasoning over long contexts, exhibiting phenomena such as the "lost-in-the-middle" effect. The authors investigate whether the positional mechanism generalizes to longer sequences or if LMs employ alternative strategies to maintain retrieval accuracy.

Methodology

The authors employ mechanistic interpretability techniques, specifically interchange interventions, to probe the internal causal structure of LMs.

Task Setup: The study utilizes templatic in-context reasoning tasks involving "Entity Groups" (tuples of entities filling specific roles, e.g., Name + Food). The model is queried to retrieve a target entity based on a query entity within a specific group.
Counterfactual Design: To distinguish between competing hypotheses, the authors construct paired datasets of original and counterfactual inputs. These pairs are designed such that three hypothesized mechanisms—Positional, Lexical, and Reflexive—would predict distinct outputs if an interchange intervention were performed on their respective intermediate variables.
- Positional Mechanism ( $P$ ): Retrieves the entity at the counterfactual query's group index.
- Lexical Mechanism ( $L$ ): Retrieves the entity bound to the counterfactual query entity (e.g., if the query changes from "Tim" to "Ann", it retrieves what "Ann" loves).
- Reflexive Mechanism ( $R$ ): Retrieves the target entity via a direct pointer originating from the target itself, dereferenced through the query.
Intervention Protocol: The authors perform interchange interventions on the last token residual stream at specific layers ( $\ell$ ) identified as the point just before retrieval occurs. They patch the residual stream vector from the counterfactual run into the original run and observe the resulting next-token distribution.
Causal Modeling: A causal model ( $M$ $M$ ) is developed to simulate the LM's behavior. This model combines the three mechanisms as a weighted mixture:
- The positional term is modeled as a Gaussian distribution centered at the query index, with variance dependent on the index.
- The lexical and reflexive terms are modeled as one-hot distributions.
- The model is trained to predict the LM's logit distribution using Jensen–Shannon Divergence (JSD) as the loss metric.
Scope: Experiments cover nine models across three families (Llama, Gemma, Qwen) ranging from 2B to 72B parameters, evaluated on ten distinct binding tasks.

Key Contributions

1. Identification of Three Distinct Mechanisms

The paper challenges the prevailing view that retrieval is solely positional. It demonstrates that LMs employ a mixture of three mechanisms:

Positional: Effective for entities at the beginning and end of the context but becomes noisy and diffuse for middle positions as the number of groups ( $n$ ) increases.
Lexical: Retrieves the bound counterpart of the query entity (e.g., "pie" $\to$ "Ann"). This mechanism provides a sharp, one-hot signal.
Reflexive: Uses a direct pointer from the target entity to itself, which is retrieved via the query. This is particularly crucial when the target appears before the query in a clause (e.g., "Tim loves tea" $\to$ "Who loves tea?"), a scenario where lexical copying is architecturally difficult for autoregressive models.

2. The "U-Shaped" Reliability Pattern

Through extensive intervention experiments, the authors uncover a consistent pattern:

First and Last Positions: Retrieval is dominated by the positional mechanism, which is highly concentrated and reliable.
Middle Positions: The positional signal degrades significantly. To compensate, LMs rely heavily on the lexical and reflexive mechanisms, which provide sharper signals to disambiguate the correct entity.
Interaction: The mechanisms exhibit "competitive synergy." When indices align, signals amplify; when they diverge, they suppress one another.

3. Causal Model of Next-Token Prediction

The authors construct a causal model that combines these three mechanisms. This model achieves 95% agreement (measured by Jensen–Shannon Similarity) with the actual next-token distributions of the LMs across various tasks and model sizes. In contrast, models relying solely on the positional mechanism (the prevailing view) achieve significantly lower scores (~45%), performing worse than a uniform baseline in some configurations.

4. Generalization to Naturalistic Inputs

The study validates that these findings hold in more naturalistic settings where entity groups are interleaved with "filler" text (open-ended sentences). While the addition of filler tokens increases noise and slightly weakens the lexical mechanism, the fundamental reliance on the three-mechanism mixture persists, offering a mechanistic explanation for the "lost-in-the-middle" effect.

Results

Mechanism Prevalence: In middle positions of long contexts, the positional mechanism accounts for only ~20% of the model's behavior, whereas the lexical and reflexive mechanisms become the primary drivers of correct retrieval.
Model Robustness: The mixing pattern is consistent across model families (Llama, Gemma, Qwen) and scales (2B to 72B parameters).
Reflexive Validation: Through specific counterfactual designs where the counterfactual answer does not exist in the original context, the authors prove the existence of a reflexive pointer mechanism distinct from the answer token itself.
Ablation Studies: Removing any single mechanism from the causal model significantly degrades performance, confirming that all three are necessary for an accurate abstraction of LM behavior.

Significance and Claims

The paper claims to establish a more complete picture of how LMs bind and retrieve entities in-context. By moving beyond the simplistic positional view, the authors provide a mechanistic explanation for both the strengths (robust retrieval via multiple signals) and fragilities (noise in middle positions leading to "lost-in-the-middle" errors) of LLMs in long-context settings.

The work suggests that the "lost-in-the-middle" phenomenon is not merely a failure of attention but a specific degradation of the positional signal, which is only partially compensated for by lexical and reflexive mechanisms. The authors release their code and data to facilitate further research into the causal mechanisms of in-context reasoning.