Entity Tracking in Language Models (2023)

Abstract: Keeping track of how states of entities change as a text or dialog unfolds is a key prerequisite to discourse understanding. Yet, there have been few systematic investigations into the ability of large language models (LLMs) to track discourse entities. In this work, we present a task probing to what extent a language model can infer the final state of an entity given an English description of the initial state and a series of state-changing operations. We use this task to first investigate whether Flan-T5, GPT-3 and GPT-3.5 can track the state of entities, and find that only GPT-3.5 models, which have been pretrained on large amounts of code, exhibit this ability. We then investigate whether smaller models pretrained primarily on text can learn to track entities, through finetuning T5 on several training/evaluation splits. While performance degrades for more complex splits, we find that even when evaluated on a different set of entities from training or longer operation sequences, a finetuned model can perform non-trivial entity tracking. Taken together, these results suggest that language models can learn to track entities but pretraining on text corpora alone does not make this capacity surface.

## introduction a key prerequisite to long-context understanding and generating coherent text is the ability to accurately represent entities as the discourse unfolds @xcite @xcite @xcite . for example, consider the following example in the context of a recipe: q: box 1 contains the book. box 2 contains the apple. box 4 contains the brain. move the book into box 2. put the bell into box 4. move the bell and the brain into box 5. box 2 contains ____ a: the apple and the book in order to understand this instruction, several distinct abilities are necessary: new discourse entity recognition: recognizing when new discourse entities are introduced. e.g., a bowl introduces a new discourse entity but the final batter or any lumps of ... does not. coreference resolution: associating referring expressions with discourse entities. e.g., a light batter and the final batter refer to the same entity. discourse entity tracking: tracking the state changes made to each discourse entity. e.g., the eggs are put into the bowl and mixed with the other ingredients. there exist many datasets that aim to evaluate these abilities (e.g., @xcite @xcite @xcite @xcite and many nlp models that aim to solve these tasks (e.g., @xcite @xcite @xcite @xcite @xcite @xcite . in the context of large language models (llms), @xcite , @xcite , and @xcite found that representations of lstms and transformer-based models such as bert @xcite do capture coreference relations. @xcite and @xcite found that pretrained models are able to detect whether noun phrases introduce discourse entities, albeit not fully systematically. the question of whether llms can track the state of discourse entities, however, has mostly been indirectly evaluated. @xcite showed that gpt-2 @xcite can learn to predict valid chess moves based on a compact, nonlinguistic description of previous moves. similarly, @xcite showed that a gpt model trained on othello can predict valid next moves, and that these predictions are tied to the model's internal representations of the board states. still, these results do not tell us whether llms track state changes expressed in natural language discourses. the most relevant evaluation is @xcite , where they tested whether model representations encode entity states described in naturalistic text. using a probing classifier, they found that the states can be decoded from t5 @xcite and bart @xcite with high accuracy. however, as we show in a reanalysis of their results (section 2), they do not provide definitive evidence for entity tracking. hence, whether llms can track entities during the processing of natural language discourse remains an open question. contributions this work attempts to answer this question by developing a task targeted towards evaluating a language model's ability to track state changes of discourse entities (illustrated in figure 1 ). we use this novel task to evaluate @xcite , @xcite without any finetuning. we find that only models in the gpt 3.5 series, which have been trained on both text and code, are able to perform non-trivial entity tracking. we then show that a smaller language model (t5) can learn to perform non-trivial entity tracking and also demonstrates some capacity to generalize to state descriptions with more operations or with low lexical overlap. our results suggest that language models can learn to track entities but pretraining on text corpora alone does not make this capacity surface. more broadly, our task can also serve as a useful tool for investigations into emergent world models in lms (e.g., @xcite . 1 2 reanalysis of @xcite we start from examining @xcite , the most relevant work to ours. they adapted two exist-ing datasets, alchemy @xcite and textworld @xcite , to test a model's ability to track state changes of an entity. the input to the model is a text description of the initial world state followed by state-changing instructions. based on this description, the model is expected to identify the correct final state of each entity. for example, for alchemy, the model receives formulaic descriptions of 7 beakers containing different amounts of colored liquids, followed by instructions that manipulate their contents such as pouring the liquid from one beaker into another, or draining a beaker. given an input like (2), the model is expected to recognize that the first beaker has 4 units of brown liquid, the second beaker has 2 units of red liquid, and the third beaker is empty. the first beaker has 1 green, the second beaker has 2 red, the third beaker has 3 red. pour the last red beaker into beaker 1. mix. using such descriptions, @xcite found that a probing classifier that takes as input the encoding of these descriptions from t5 or bart is able to correctly predict the state of 75-76% of the entities, suggesting some degree of success on entity tracking. however, this conclusion becomes questionable when the datasets and the results are scrutinized further. specifically, we conducted a fine-grained analysis of the success cases of the alchemy experiment. in this experiment, the state of each beaker was probed after each state-changing instruction. because each instruction targets at most two beakers (e.g., pour x into y) and there are 7 beakers in total, there is a sparse representation of cases probing a beaker that actually underwent a change in the dataset. indeed, 62.7% of all beaker states probed were identical to the initial state, meaning that a simple baseline that always predicts the initial state already achieves 62.7% accuracy (this is also noted by @xcite . a second potential for shortcuts was the high rate of empty final states (32.4%). 2 for these cases, the initial state can often be entirely disregarded, due to the presence of an emptying instruction such as drain the fourth beaker. this instruction alone is sufficient to predict the fourth beaker's final state independent of its initial state. therefore, such examples are also not best suited to fully assess entity tracking. given the high prevalence of these two trivial scenarios (87.6% in total), only 12.4% of the datapoints can be considered as truly assessing state changes unfolding over a discourse context. if the accuracy is computed on the trivial and non-trivial cases separately, the probing classifier achieves 86.8% accuracy on trivial cases but only 3.1% accuracy on non-trivial cases, showing that most of the reported success derives from the trivial cases. in summary, our reanalysis suggests that the results of @xcite do not provide conclusive evidence for non-trivial state tracking abilities in language models. 3 however, it remains unclear whether this is due to issues with the setup or a true lack of entity tracking capacity. to this end, we propose a new behavioral evaluation. ## experiment 1: in-context demonstration in the first set of experiments, we evaluated pretrained lms using a small number of in-context 4 for example, the signature of an initial state in which the first box contains two objects and the rest contains 1 object each would be 2111111. ## experiment 2: finetuning we investigated whether smaller models at the scale of t5 can learn to track entity states through a series of experiments where we provide supervised ## conclusion we set out to investigate whether pretrained llms exhibit entity tracking behavior. we developed a task that allowed us to evaluate whether llms can predict the state of an entity based on an initial state description and operations that act upon it. in the first set of experiments, we found that gpt-3 davinci, a vanilla pretrained language model, and flan-t5, an instruction-finetuned language model, fail at this task and simply repeat the initial state description. the gpt-3.5 models (the core difference with the aforementioned models being code in pretraining corpora), on the other hand, exhibited non-trivial entity tracking behavior. even after many operations affecting an entity state, they consistently performed above a strong random baseline. in the second set of experiments, we showed that this behavior can also to a large extent be learned by smaller models such as t5. when we finetune the model on this task, it is also able to perform entity tracking on examples that differ along several dimensions from the training data. taken together, our results provide evidence that (a) vanilla language models that have mainly been trained on text data do not exhibit entity tracking abilities out of the box, (b) pretraining on code and text data considerably improves this ability, foot_6 and (c) finetuning on this task can make this behavior also surface in smaller models that have primarily been trained on text data, although it remains an open question how general this ability is in smaller finetuned models. what are reasons behind the efficacy of training both on text and code? for producing correct code, tracking the states of variables is important. therefore, to speculate, this kind of pretraining data may provide a stronger signal for the model to track entities compared to pure text data. it could also be that, as previously speculated @xcite , pretraining on code provides additional grounding. the present results also highlight the importance of the composition of the pretraining data. our results showed that models that are trained on the language modeling objective on large corpora show dramatically different behavior on entity tracking depending on whether the training data includes substantial amount of code or not. in this light, future work could investigate the effect of pretraining on code more widely, including its effect on the model representations and to what extent it facilitates the emergence of world models in lms. finally, we also laid out several principles that should be followed when evaluating state tracking abilities. apart from these specific principles, we make a more general point that in assessing abilities related to meaning, one needs to consider potential strategies that the model could use to solve the task and make sure that the test examples do not mimic the distributional patterns of the training or finetuning data. only then can we properly assess meaning-related capacities of lms.