Language Models use Lookbacks to Track Beliefs
This work addresses the challenge of understanding belief tracking in language models, which is crucial for advancing Theory of Mind reasoning in AI, though it appears incremental as it builds on existing causal mediation techniques.
The paper tackles the problem of how language models represent characters' beliefs, especially when they differ from reality, by analyzing their Theory of Mind capabilities using causal mediation and a new dataset called CausalToM. It uncovers a lookback mechanism where the model binds character-object-state triples using Ordering IDs in low-rank subspaces to retrieve correct information when reasoning about beliefs.
How do language models (LMs) represent characters' beliefs, especially when those beliefs may differ from reality? This question lies at the heart of understanding the Theory of Mind (ToM) capabilities of LMs. We analyze LMs' ability to reason about characters' beliefs using causal mediation and abstraction. We construct a dataset, CausalToM, consisting of simple stories where two characters independently change the state of two objects, potentially unaware of each other's actions. Our investigation uncovers a pervasive algorithmic pattern that we call a lookback mechanism, which enables the LM to recall important information when it becomes necessary. The LM binds each character-object-state triple together by co-locating their reference information, represented as Ordering IDs (OIs), in low-rank subspaces of the state token's residual stream. When asked about a character's beliefs regarding the state of an object, the binding lookback retrieves the correct state OI and then the answer lookback retrieves the corresponding state token. When we introduce text specifying that one character is (not) visible to the other, we find that the LM first generates a visibility ID encoding the relation between the observing and the observed character OIs. In a visibility lookback, this ID is used to retrieve information about the observed character and update the observing character's beliefs. Our work provides insights into belief tracking mechanisms, taking a step toward reverse-engineering ToM reasoning in LMs.