Do Language Models Track Entities Across State Changes?

Entity tracking (ET), the ability to keep track of states, is a fundamental skill that underlies complex reasoning. An increasing amount of work investigates how transformer language models (LMs) solve entity binding without state changes. However, there is limited understanding of how non-toy LMs address ET problems of realistic difficulties expressed in natural language. To this end, we investigate the mechanisms underlying ET in more complex scenarios featuring multiple state-changing operations. We find that LMs do not incrementally track world states across tokens or query-relevant states across layers, but simply aggregate relevant information in parallel at the last token when the query becomes evident. We further investigate mechanisms of individual operations (PUT, REMOVE, MOVE) to characterize this non-incremental ET mechanism. Surprisingly, LMs implement the REMOVE operation with a fragile global suppression tag; this global removal mechanism predicts various failure modes that we confirm behaviorally. We provide a mechanistic solution of nullifying this tag to partially address this issue. Overall, our findings reveal that LMs solve a fundamentally sequential task using a non-sequential strategy. More broadly, our work illustrates how behavioral and mechanistic analyses can fruitfully interact. Behavioral results inform mechanistic hypotheses, and insights from mechanistic analyses help build stronger behavioral evaluations by predicting failure modes missing from existing evaluations.

When we read stories or hold conversations, we need to keep track of how entities (e.g., people or objects) change their states over time. This entity tracking capability is an important skill that underlies many tasks that humans and machines are expected to tackle, from maintaining coherent discourse to solving mathematical problems. We ask how AI models solve this task by looking at the their inner computations. We use a controlled tracking problem setting involving boxes and their contents (e.g., "The apple is in Box 1. The orange is in Box 2. Remove the orange from Box 2. What is in Box 2?"). While this problem has to be solved by considering one state change after another, we find that models do not solve this problem sequentially; they only look back at relevant descriptions when asked. In particular, how models track removed objects deviates from our expectation. The models often remove objects from the world entirely rather than from a particular box that they were in, which leads to incorrect tracking. Our work shows that with better understandings of models' inner computations, we are able to predict when models will fail and can fix them.