Mixing Mechanisms: How Language Models Retrieve Bound Entities In-Context
Yoav Gur-Arieh, Mor Geva, Atticus Geiger
2025-10-08
Summary
This paper investigates how large language models (like the ones powering chatbots) remember and use information presented to them in a conversation or text. Specifically, it looks at how they keep track of relationships between things mentioned, like who loves what.
What's the problem?
Language models are pretty good at remembering simple connections between two things. However, when you give them a longer text with many different relationships to track, their ability to recall those connections accurately breaks down, especially for items mentioned in the middle of the text. It's like trying to remember everything from a long list – the middle items get fuzzy.
What's the solution?
The researchers found that language models don't rely on just *one* way to remember things. They use a combination of three methods: remembering based on where something was mentioned (position), remembering by linking it to what it's related to (lexical connection), and directly remembering it on its own (reflexive pointer). They built a model that combines all three of these methods and accurately predicts how the language model will behave, even with long and complex texts.
Why it matters?
Understanding how language models remember information is crucial for improving their performance. This research gives us a more complete picture of the memory process, which can help us build better and more reliable AI systems that can handle complex conversations and understand longer pieces of text more effectively.
Abstract
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.