Technical Summary: Do Language Models Track Entities Across State Changes?

Problem Statement

Entity tracking (ET) is a fundamental skill required for complex reasoning, involving the ability to maintain and update the states of objects within a context as operations (e.g., moving, adding, removing) occur. While previous research has investigated how transformer language models (LMs) handle static entity binding (associating an object with a location without state changes) or solved ET problems using simplified synthetic languages and "toy" models, there is limited understanding of how non-toy, pretrained LMs handle realistic ET tasks expressed in natural language with multiple state-changing operations. Specifically, it remains unclear whether these models incrementally build a world state as the context unfolds or if they employ a different strategy to resolve queries.

Methodology

The authors investigate the mechanisms underlying ET in non-toy models (Gemma-2-2B, CodeLlama-13B, and Llama-3.1-70B) using a dataset derived from the "box" benchmark (Kim & Schuster, 2023). The dataset consists of an initial description of objects in boxes, followed by a sequence of state-changing operations (PUT, REMOVE, MOVE), and a query asking for the contents of a specific box.

The study employs a dual approach combining behavioral analysis and mechanistic interpretability:

1. Behavioral Probing:

   • Global vs. Local Probes: Linear probes are trained on the residual stream of the final token to determine if the model encodes a global state (location of all objects) or only local states (contents of the queried box).
   • Prior State Probing: Probes are trained to detect if prior states (states before the final operation) are sequentially decodable across layers, testing whether the model processes state changes incrementally or in parallel.

2. Mechanistic Analysis:

   • Circuit Tracing (Path Patching): Used to identify the specific attention heads responsible for the PUT operation, comparing them to the circuit for static entity binding (DESCRIPTION).
   • Subspace Analysis: Desiderata-based counterfactual masking (DCM) and subspace patching are used to determine if PUT and DESCRIPTION circuits share the same communication subspaces for positional information.
   • REMOVE Mechanism Investigation: Due to the difficulty of analyzing "absence" of signals, the authors employ logit/rank analysis, ternary probes (classifying states as {exists, removed, does not exist}), and causal interventions (nullifying probe weights) to characterize how REMOVE operations are implemented.

Key Findings

1. Non-Incremental Tracking Strategy

Contrary to the hypothesis that LMs incrementally build world states as they process tokens or layers, the results support a non-incremental strategy:

• Across Tokens: Global probes fail to decode the full world state in the final token's residual stream, while local probes (targeting the queried box) achieve high accuracy. This suggests models do not maintain a global state but dynamically consolidate relevant information only when the query is presented.
• Across Layers: Probing for prior states reveals that models do not process state changes sequentially across layers. Instead, prior states are decoded with similar accuracy to final states in the same layers, indicating parallel processing of all relevant operations at the final token.

2. Mechanisms of Operations

• PUT Operation: The PUT mechanism resembles the entity binding circuit found in prior work (Prakash et al., 2024). It utilizes a "look-back" mechanism where specific attention groups (A, B, C, D) bind new objects to box IDs using positional information (order IDs). Crucially, PUT and DESCRIPTION circuits share functional subspaces for transmitting positional information.
• REMOVE Operation: The REMOVE operation is implemented via a "global suppression tag."
  • Logit analysis shows that when an object is removed, its logit is suppressed, but this suppression is not strictly local; it affects the object globally regardless of the box context.
  • Ternary probes identify a "remove tag" signal primarily located at the object token (not the box ID token).
  • Causal interventions confirm that nullifying the remove tag at the object token prevents the model from predicting the removed object, whereas interventions at the box ID token have no effect. This implies the model uses a fragile, global heuristic to suppress objects rather than a precise, box-specific unbinding mechanism.

3. Predicted and Observed Failure Modes

The mechanistic insight regarding the "global remove tag" predicts specific failure modes that are not evident in standard behavioral tests:

• No-op Remove: Removing an object from a box that was never mentioned in the context (or removing an object that doesn't exist) leads to the global suppression of that object type, causing errors in other boxes.
• Shared-label Objects: If the same object type exists in multiple boxes, removing it from one box causes the model to suppress it globally, leading to incorrect predictions for the other box.
• Re-introduction: Removing an object and then putting it back in a different box fails because the global suppression tag persists.
  Behavioral tests on these specific scenarios confirmed high degeneration rates, validating the mechanistic hypothesis.

Contributions

1. Evidence of Non-Incremental Tracking: The paper provides empirical evidence that LMs solve ET tasks by retrieving target tokens in parallel at the final query step rather than through cumulative world building across tokens or layers.
2. Mechanistic Characterization of Operations: The study characterizes PUT as a binding circuit similar to static entity tracking and identifies REMOVE as a global suppression mechanism driven by object-level tags.
3. Bridging Behavioral and Mechanistic Analysis: The work demonstrates a feedback loop where mechanistic insights (the global remove tag) predict novel failure modes, which are then validated through new behavioral tests. Conversely, behavioral results inform the formulation of mechanistic hypotheses.
4. Interventional Solutions: The authors show that nullifying the global remove tag via intervention can partially fix these failure modes, offering a proof-of-concept for improving model robustness through mechanistic understanding.

Significance

The paper highlights a fundamental architectural limitation in current transformer LMs: they solve fundamentally sequential tasks (tracking state changes over time) using non-sequential, parallel strategies. This reliance on aggregating relevant context at the last token, rather than maintaining a persistent world state, leads to brittle behaviors when operations are ambiguous or when objects are shared across contexts. The work illustrates how combining behavioral evaluation with mechanistic interpretability can uncover hidden failure modes and guide the development of more robust reasoning capabilities in language models.