PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yoav Gur-Arieh, Mor Geva, Atticus Geiger

arXiv:2510.06182v2cs.CL

TL;DR

Language models bind related entities to support in-context reasoning, but positional retrieval becomes unreliable in middle positions as entity lists grow. This paper studies and combines positional, lexical, and reflexive mechanisms; across models and binding tasks, it finds that models mix them and develops a causal model of retrieval.

  • Problem

    Binding supports in-context reasoning, but positional retrieval becomes unreliable in middle positions as entity lists grow.

    Source: §1 · p. 1

  • Method

    The paper uses counterfactual interventions to distinguish retrieval mechanisms and combines positional, lexical, and reflexive signals in a causal model of next-token distributions.

    Source: §3.2 · p. 6, §4 · p. 8

  • Results

    Across 9 models from 2 to 72 billion parameters and 10 binding tasks, language models mix all three mechanisms, with positional retrieval strongest at context ends and diffuse in the middle.

    Source: §6 · p. 10

  • Takeaways & Limitations

    In middle groups, sharper lexical and reflexive signals refine diffuse positional information, helping language models retrieve the correct entity.

    Source: §3.3 · p. 7

    The experiments use binding matrices consisting of distinct entities.

    Source: §2 · p. 3

Abstract

from arXiv · show

A key component of in-context reasoning is the ability of language models (LMs) to bind entities for later retrieval. For example, an LM might represent "Ann loves pie" by binding "Ann" to "pie", allowing it to later retrieve "Ann" when asked "Who loves pie?" Prior research on short lists of bound entities found strong evidence that LMs implement such retrieval via a positional mechanism, where "Ann" is retrieved based on its position in context. In this work, we find that this mechanism generalizes poorly to more complex settings; as the number of bound entities in context increases, the positional mechanism becomes noisy and unreliable in middle positions. To compensate for this, we find that LMs supplement the positional mechanism with a lexical mechanism (retrieving "Ann" using its bound counterpart "pie") and a reflexive mechanism (retrieving "Ann" through a direct pointer). Through extensive experiments on nine models and ten binding tasks, we uncover a consistent pattern in how LMs mix these mechanisms to drive model behavior. We leverage these insights to develop a causal model combining all three mechanisms that estimates next token distributions with 95% agreement. Finally, we show that our model generalizes to substantially longer inputs of open-ended text interleaved with entity groups, further demonstrating the robustness of our findings in more natural settings. Overall, our study establishes a more complete picture of how LMs bind and retrieve entities in-context.

1 INTRODUCTION

Binding lets language models connect entities in context so they can answer queries, but positional retrieval becomes unreliable for middle positions in longer lists. The paper argues that models combine positional, lexical, and reflexive mechanisms, with their contributions depending on entity positions.

1 INTRODUCTION

  • Binding connects related entities in text to create context representations that language models can query for next-token prediction.The paper investigates this process because language models also struggle with reasoning over long contexts.Source: §1 · p. 1
  • The positional account retrieves a target by using the query entity to identify its group’s position in context.For example, the query “pie” identifies the group containing “Ann loves pie,” from which the model retrieves “Ann.”Source: §1 · p. 1
  • Position-based retrieval holds in simple settings but becomes unreliable for entities in the middle of longer lists.This pattern echoes the “lost-in-the-middle” effect in language models.Source: §1 · p. 1
  • Lexical retrieval uses the query entity’s bound counterpart, while reflexive retrieval uses a direct pointer to the queried entity.Figure 1 illustrates counterfactual patching that separates these signals after binding information is encoded in entity tokens and used jointly.Source: §1 · p. 2, Figure 1 · p. 2
  • Across 9 models from 2 to 72 billion parameters and 10 binding tasks, retrieval robustly mixes positional, lexical, and reflexive mechanisms.The mechanisms’ interplay depends on the positions of the query and retrieved entities.Source: §1 · p. 2

2 PROBLEM SETUP AND PRIOR WORK

The section formalizes templatic binding tasks and describes the intervention and causal-model approaches used to study retrieval. It situates this work against prior accounts of lookback and positional retrieval, including their narrower settings and lower causal faithfulness in longer contexts.

2 PROBLEM SETUP AND PRIOR WORK

  • Tasks pair entities from distinct roles into groups placed in templatic text, then query an entity and target position to form a question.A group is a tuple containing one entity per role; the task’s binding matrix records the groups.Source: §2 · p. 3
  • Experiments use binding matrices made up of distinct entities.This is an assumption of the experimental setup.Source: §2 · p. 3
  • Interchange interventions replace selected hidden states in an original run with states from a paired counterfactual run to identify causally relevant components.The paper uses paired examples and may patch residual-stream vectors.Source: §2 · p. 3
  • The causal abstraction predicts next-token distributions with a mixture of three mechanisms, evaluated against the language model using paired counterfactual interventions.The comparison measures similarity between the causal model’s and language model’s next-token distributions, averaged across a dataset.Source: §2 · p. 3
  • Prior work describes lookback retrieval as aggregating entity information at the last token of a group, then using attention to retrieve an entity conditionally.Earlier positional analyses either used narrow contexts or found low causal faithfulness in longer contexts.Source: §2 · pp. 3–4
  • For Gemma 2 2B with three entities per group, the plotted effects show a U-shaped pattern: positional retrieval is stronger at the first and last indices, while lexical and reflexive effects rise in the middle.The figure plots effects across layers for representative indices and across all indices at layer 18.Source: Figure 2 · p. 4

3 THREE MECHANISMS FOR RETRIEVING BOUND ENTITIES

The section proposes positional, lexical, and reflexive mechanisms for retrieving bound entities, then uses counterfactual interchange interventions to distinguish them. Experiments show that their contributions vary with entity-group position and target position, and validate reflexive retrieval as a pointer distinct from the answer entity.

3 THREE MECHANISMS FOR RETRIEVING BOUND ENTITIES

  • The section evaluates nine models, using two binding tasks broadly and all ten tasks for Gemma 2B and Qwen 7B, with interventions on last-token residual streams.The intervention layer is chosen just before retrieval starts and is consistent across tasks for each model.Source: §3 · p. 4

3.1 THE POSITIONAL MECHANISM AND TWO ALTERNATIVES

  • The proposed alternatives retrieve through the queried entity’s bound counterpart (lexical) or a direct pointer to the target (reflexive), alongside position-based retrieval.The reflexive route is relevant when the target precedes the query, because autoregressive attention looks only right to left and cannot copy the target backward to the query.Source: §3.1 · pp. 4–5

3.2 DESIGNING COUNTERFACTUAL INPUTS TO DISTINGUISH THE THREE MECHANISMS

  • Paired original and counterfactual inputs make positional, lexical, and reflexive mechanisms predict distinct outputs under interchange interventions on their respective variables.A remaining confound—whether the reflexive signal is a pointer or the answer itself—is addressed with a separate counterfactual design.Source: §3.2 · pp. 5–6, §3.4 · p. 7

3.3 INTERVENTION EXPERIMENTS

  • In 20-group contexts, positional retrieval accounts for only 20% of behavior in middle groups, where lexical and reflexive contributions vary with target position and interact with positional signals.The positional signal is diffuse and often links to nearby groups; lexical and reflexive signals are sharper, and the target’s within-group position modulates their use.Source: §3.3 · p. 7

3.4 VALIDATING THE EXISTENCE OF THE REFLEXIVE MECHANISM

  • A modified counterfactual with an answer entity absent from the original context shows that the reflexive signal is a pointer, not the answer entity itself.At layer ℓ the pointer cannot be dereferenced, while at layer ℓ+1 the already-retrieved answer entity can be patched and output.Source: §3.4 · pp. 7–8, Figure 4 · p. 8

4 A SIMPLE MODEL FOR SIMULATING ENTITY RETRIEVAL IN-CONTEXT

The section models retrieval as a position-weighted mixture of positional, lexical, and reflexive signals, then evaluates its fit and how the mechanisms shift across positions and padding.

4 A SIMPLE MODEL FOR SIMULATING ENTITY RETRIEVAL IN-CONTEXT

  • The causal model combines one-hot lexical and reflexive terms with a weighted Gaussian positional term whose standard deviation varies quadratically with position.It computes a logit for each entity from the three mechanism variables.Source: §4 · p. 8
  • 0.95 average JSS for the full model exceeds 0.44 for the prevailing positional view and 0.5 for the uniform baseline, while remaining slightly below the oracle.The reported comparison is for gemma-2-2b-it on the music task.Source: Figure 5 · p. 9, §4 · p. 9
  • 150 interchange interventions per index combination yield 8,000 probability distributions, split into 70% training and evenly divided validation and test sets.The model is trained using Jensen–Shannon divergence.Source: §4 · p. 9
  • 0.85 JSS for the one-hot positional variant is below the full model’s 0.95, showing that replacing the Gaussian positional term hurts performance.Source: §4 · p. 9
  • As padding increases, predicted indices become fuzzier for early tokens, while the positional mechanism strengthens at the expense of the lexical mechanism.Source: Figure 6 · p. 10
  • Ablations show that lexical retrieval matters less when the target is first, while reflexive retrieval matters less when it is last; the positional spread widens in middle indices.For the middle target setting, learned lexical and reflexive weights are similar, with reflexive slightly more dominant.Source: §4 · pp. 9–10, Figure 5 · p. 9

5 INTRODUCING FREE FORM TEXT INTO THE TASK

The paper tests retrieval in longer, more naturalistic prompts by inserting entity-less filler sentences between entity groups. The model initially remains consistent, but greater padding changes the balance and sharpness of the mechanisms.

5 INTRODUCING FREE FORM TEXT INTO THE TASK

  • Entity groups are interleaved with filler sentences, up to 500 tokens between groups, to test retrieval in longer and noisier prompts.The filler sentences are designed not to contain sequences signaling entity tracking or binding.Source: §5 · p. 10
  • With 10,000 filler tokens, positional information is nearly absent for first-half entity tokens other than the first group, while lexical effects decline and mixed effects become more diffuse.The model initially remains consistent across a ten-fold increase in tokens; the authors suggest the later shift may mechanistically explain the lost-in-the-middle effect.Source: §5 · p. 10

6 CONCLUSION

The conclusion describes retrieval as a mixture of positional, lexical, and reflexive mechanisms, validated across models and binding tasks. The appendix states that generalization is evaluated across values of n, models, and tasks.

6 CONCLUSION

  • Positional retrieval is effective at context beginnings and ends but diffuse in the middle, where lexical and reflexive signals provide sharper retrieval cues.The account is validated across 9 models ranging from 2–72B parameters and 10 binding tasks.Source: §6 · p. 10

A EVALUATING GENERALIZATION

A.1 BINDING TASKS

The appendix evaluates ten templatic binding tasks spanning domains, syntaxes, and entity-set sizes, and reports a consistent mechanism pattern across tasks and target positions.

A.1 BINDING TASKS

  • Ten binding tasks span domains, syntaxes, and entity sets of 23–80 items; one task uses m = 2 and nine use m = 3.Source: §A.1 · p. 13
  • For m = 3, the two query entities are the group members that are not the target; when the target is second, the queries occupy positions one and three.Source: §A.1 · p. 13
  • Across tasks and target positions, positional retrieval dominates for early and late groups, while lexical and reflexive mechanisms take over in the middle.Reflexive effects are more prominent for first-position targets, lexical effects for last-position targets, and both are balanced for middle-position targets.Source: §A.1 · p. 13

A.2 REPLICATING RESULTS ACROSS MODELS

Across models and tasks, retrieval mixes positional, lexical, and reflexive mechanisms, with position strongest at the ends and the other mechanisms taking over in the middle. As entity groups grow, middle positions become less reliable and harder to distinguish by position.

A.2 REPLICATING RESULTS ACROSS MODELS

  • Across 9 models in 3 families spanning 2–72B parameters, positional retrieval dominates at the first and last groups while lexical and reflexive mechanisms take over in the middle.In Qwen, positional efficacy strengthens with model size.Source: §A.2 · p. 13, Figure 10 · p. 15, Figure 11 · p. 16

A.3 EFFECT OF n

B ENCODING OF POSITIONAL INFORMATION

  • First and last entity groups are readily separable by position, whereas middle groups overlap increasingly as context size grows.The analysis uses PCA projections and linear-probe accuracy across context sizes.Source: B ENCODING OF POSITIONAL INFORMATION · pp. 14–15

C BINDING SIGNALS IN ENTITY TOKENS

Binding signals are present in entity tokens as well as at the last token, but their locations and routes differ by mechanism. Additional experiments show how mechanism agreement and layer choice affect retrieval.

C BINDING SIGNALS IN ENTITY TOKENS

  • Positional and lexical binding information in entity tokens is used in layers 12–19, while reflexive information is used in layers 6–12.PosSwap, LexSwap, and RefSwap achieve nearly identical intervention accuracies to TargetRebind at the last token.Source: C BINDING SIGNALS IN ENTITY TOKENS · p. 16
  • Blocking last-token attention to the query entity in layers 11–16 drops accuracy from 98% to 37%; blocking both it and the following token drops accuracy to 8%.Blocking the following token alone lowers accuracy by only one point, while blocking both routes in the described intervention drops it to 6%.Source: C BINDING SIGNALS IN ENTITY TOKENS · p. 16

D.1 MECHANISM AGREEMENT

  • When positional and reflexive or lexical mechanisms point to the same entity, the model consistently responds with that entity.The alignment experiments use t_entity = 1 for positional–reflexive agreement and t_entity = 3 for positional–lexical agreement.Source: §D.1 · p. 19

D.2 FINDING THE TARGET LAYER

  • At layer ℓ, counterfactual-answer responses are at or below chance, but become the majority at ℓ+1 and ℓ+2, marking a shift from binding information to retrieval.The target layer is consistent across the boxes and music tasks.Source: Figure 12 · p. 18, §D.2 · p. 20

D.3 REMOVING TARGETED ENTITY TOKENS

Removing entities targeted by lexical or reflexive pointers leaves the positional mechanism to produce noisy retrieval. The three-mechanism pattern also persists under varied phrasings, while supplementary causal-model evaluations report additional metrics and training details.

D.3 REMOVING TARGETED ENTITY TOKENS

  • When lexical or reflexive targets are absent from the original prompt, the model relies on the positional mechanism, yielding a noisy distribution around the positional index.At the retrieval layer, the model instead responds with the counterfactual answer, ruling out the proposed alternative that another mechanism simply prevents answers absent from context.Source: §D.3 · pp. 20–21
  • Positional and lexical binding information is localized to entity-token layers 12–19, while reflexive binding information is localized to layers 6–12.These are the reported PosSwap, LexSwap, and RefSwap intervention results.Source: Figure 19 · p. 21

D.4 LINGUISTIC VARIABILITY

  • With 12 randomized phrasing variations for each boxes and music task, results remain consistent with reliance on all three mechanisms.The pattern remains mediated by entity-group index and t_entity.Source: §D.4 · p. 21, Figure 21 · p. 22

E ADDITIONAL CAUSAL MODELS

  • Additional causal-model evaluations report KL divergence for Gemma on music, other metrics for Gemma on sports and Qwen on both tasks, and an Adam training setup.Training uses learning rate 0.05, batch size 512, and up to 2,000 epochs with early stopping after 200 epochs.Source: E ADDITIONAL CAUSAL MODELS · p. 22

F FURTHER VALIDATION OF THE REFLEXIVE MECHANISM

The reflexive-mechanism validation blocks attention from the last token to a target entity. The reported interventions support the mechanism's role in binding and retrieving entities in context.

F FURTHER VALIDATION OF THE REFLEXIVE MECHANISM

G CONTEXT LENGTH ABLATION STUDY

Padding controls show that model behavior is governed primarily by the number of bound entities rather than sequence length; separability also decreases as entity count rises.

G CONTEXT LENGTH ABLATION STUDY

  • Contexts with n = 3–19 were padded with entity-less sentences to match the sequence length for n = 20.If sequence length explained the effects of increasing n, the padded results would match those for n = 20 and each other.Source: G CONTEXT LENGTH ABLATION STUDY · p. 23
  • First and last entity groups remain linearly separable, while middle groups overlap substantially; separability decreases as n increases.Figure 22 presents PCA projections and multinomial logistic regression probes across layers and values of n.Source: Figure 22 · p. 24
  • Padding results preserve the overall patch-effect patterns, indicating that the number of entities in context primarily governs model behavior rather than sequence length.Padding slightly shifts the distribution, but trends align closely with unpadded results.Source: G CONTEXT LENGTH ABLATION STUDY · p. 23, Figure 27 · p. 26
Loading 2510.06182v2…