Entity Tracking in Language Models
Year: 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 nontrivial 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.

heading: 1 Introduction
text: A key prerequisite to long-context understanding and generating coherent text is the ability to accurately represent entities as the discourse unfolds (Karttunen, 1976; Groenendijk and Stokhof, 1991; Heim, 2002; Nieuwland and Van Berkum, 2006; Kamp et al., 2011, i.a.). For example, consider the following example in the context of a recipe:
(1) Put the eggs, sugar, flour, and baking powder in a bowl and mix to form a light batter. Make sure that the final batter does not contain any lumps of flour or sugar.
∗Equal contribution. Author order was determined by a fair process of letting Cookie the cat pick a colored ball.

heading: 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 ____
text: 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., Walker et al., 2006; Pradhan et al., 2012; Rahman and Ng, 2012; Weston et al., 2015; Chen et al., 2018; Bamman et al., 2020; Uryupina et al., 2020) and many NLP models that aim to solve these tasks (e.g., Haghighi and Klein, 2010; Lee et al., 2011; Hill et al., 2016; Henaff et al., 2017; Ji et al., 2017; Lee et al., 2017; Bosselut et al., 2018; Gupta and Durrett, 2019a,b; Aina et al., 2019; Toshniwal et al., 2020; Wu et al., 2020). In the context of large language models (LLMs), Tenney et al. (2019), Clark et al. (2019), and Sorodoc et al. (2020) found that representations of LSTMs and Transformer-based models such as BERT (Devlin et al., 2019) do capture coreference relations. Loáiciga et al. (2022) and Schuster and Linzen (2022) found that pretrained models are able to
3835
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. Toshniwal et al. (2022) showed that GPT-2 (Radford et al., 2019) can learn to predict valid chess moves based on a compact, nonlinguistic description of previous moves. Similarly, Li et al. (2023) 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 Li et al. (2021), 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 (Raffel et al., 2020) and BART (Lewis et al., 2020) 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 GPT-3 (Brown et al., 2020), GPT-3.5, and Flan-T5 (Chung et al., 2022) 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., Li et al., 2023; Tsai et al., 2023).1

heading: 2 Reanalysis of Li et al. (2021)
text: We start from examining Li et al. (2021), the most relevant work to ours. They adapted two exist-
1Code and data are available at https://github.com/ sebschu/entity-tracking-lms.
ing datasets, Alchemy (Long et al., 2016) and TextWorld (Côté et al., 2019), 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.
(2) 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, Li et al. (2021) 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 Li et al.). 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
2This percentage also includes cases where a beaker was already empty in the initial world state description. The percentage of cases where a beaker was empty but contained some liquid in the initial description is 24.9%.
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 Li et al. (2021) 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.

heading: 3 Task Design and Dataset

heading: 3.1 Desiderata
text: The ability to track entities should be largely independent of specific linguistic forms. For a model that can properly track entities, it should not matter whether one talks about beakers or recipes or which specific syntactic constructions are used. This makes it an interesting ability to evaluate in the context of assessing whether and how meaning is represented, since at least classic language models are only trained on forms (Bender and Koller, 2020). At the same time, this independence of form and entity states that should hold for true entity tracking poses a challenge for evaluation, since one needs to ensure that the state of entities cannot be predicted from individual lexical items or phrases (such as the word drain in the Alchemy dataset, as discussed in Section 2). Furthermore, language models pretrained on large corpora may have learned common states of entities; for instance, that eggs often end up in a bowl. For these reasons, any task that evaluates entity tracking abilities should conform to the following four desiderata:
1. The probed states of entities should not follow similar distributional patterns to those that are likely to be present in the pretraining data (see also Linzen, 2020).
2. Individual words or phrases should not predict by themselves the state of an entity without
3Li et al. (2021) also presented two other sets of experiments. See Appendix A for details on how the other experiments exhibit similar issues.
considering the previous discourse in order.
3. If any data is used for demonstration, finetuning or training, the training and evaluation data should have little lexical overlap.
4. If any data is used for demonstration, finetuning or training, the task should not be solvable by slot-filling based on observed datapoints.
These properties cannot be guaranteed with naturalistic datasets such as recipes (Kiddon et al., 2015), science texts (Dalvi et al., 2019), or the Alchemy and TextWorld datasets, which have been previously used to evaluate entity tracking abilities. We therefore programmatically generated datasets for which these properties hold.

heading: 3.2 Dataset
text: We take inspiration from Winograd (1971) and Li et al. (2021) in designing our data. Our datasets consist of text descriptions of a particular state of the world followed by a sequence of changes. The worlds contain boxes that can be filled with objects. The objects can be placed inside the box, taken out of the box, or moved from one box to another. We define a worldW asW = (O,n,m, e) where O is a set of objects, n is the number of boxes, m is the maximum number of objects one box can contain, and e is the expected number of objects in each box in the initial world states. For our datasets, we used n = 7, m = 3, e = 2, and used a set of nouns denoting items that can plausibly fit inside a box (e.g., book, rock, brain; |O| = 100), selected from a list of words with frequency greater than 27 in the British National Corpus (BNC; Leech et al. 2001).
A dataset consists of multiple distinct scenarios. A scenario consists of an initial state and a set of operations applied to this initial state. We fixed the number of operations (NumOps) in each scenario to 12. We randomly sampled 2200 scenarios, where the initial state and the 12 operations were both randomly sampled. The sampling process is designed such that only valid operations given the current world state can be sampled. The initial state and the operations were converted into naturalistic descriptions using predefined templates. Relation to Desiderata We selected the task of moving objects across boxes because this is a domain where lexical contents of the entities do not offer cues to predict the outcome of state changes (Desideratum 1). We did not include an operation that empties a box that allows the previous
discourse to be discarded (Desideratum 2). For Desideratum 2, we furthermore considered experiments using operation descriptions with greater context dependence. Specifically, we tested scenarios with an additional Move contents of Box N to Box M operation that does not explicitly mention the object names, and scenarios where object descriptions in the operations can only be fully disambiguated by knowing the current state of a box. For Desideratum 3, we considered experiments where the phrasing of the states and operations differ entirely between demonstration/finetuning and evaluation (see Table 2). Finally, for all experiments, we computed a “signature” of every initial state that indicates the number of objects contained in each box.4 Then, we ensured that there were no two examples with identical initial descriptions modulo the object names where one appeared in the training split and the other one in the evaluation split. This prevents this task from being solvable via slot-filling (Desideratum 4).5

heading: 3.3 Task
text: We define the entity tracking task as follows. Given a natural language description of the initial state of the world followed by 0–12 state-changing operations, the content of each box at the end of the description must be correctly identified. To evaluate this, we created a test example for each box after each operation. This corresponds to n× (NumOps + 1) examples per scenario (91 exs. in our datasets). Each example is formulated in the style of a cloze test. That is, the input describes the initial state followed by a sequence of operations, ending in Box N contains __. The expected output is the correct set of objects in Box N based on the prefix description. See Appendix B for an example.

heading: 4 Experiment 1: In-context Demonstration
text: In the first set of experiments, we evaluated pretrained LMs using a small number of in-context
4For 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.
5Additionally, compared to the Alchemy setup, our setup also has a benefit of requiring fewer additional reasoning abilities. The beaker domain in Alchemy requires the model to count and perform simple arithmetic (e.g., adding one unit of liquid followed by adding two units of liquid results in three units of liquid). Moreover, some of the operations in Alchemy require knowledge about how colors are combined (e.g., mixing red and green liquids results in a brown liquid). The boxes domain removes these requirements.
demonstrations of the entity tracking task. This provides a way to probe the model without requiring substantial supervision for the task, as well as guiding the model to output the final state in a consistent format that can be automatically assessed.6

heading: 4.1 Models
text: We used models that are known to support task adaptation via in-context demonstrations: GPT3 175B (davinci: Brown et al. 2020), GPT-3.5 (text-davinci-0037), and Flan-T5 (base and XL: Chung et al. 2022). The little information that OpenAI revealed about their models8 suggests that davinci is an autoregressive language model primarily trained on text corpora. text-davinci-003 was trained on the language modeling objective on a mix of text and code, and additionally tuned with human feedback using reinforcement learning. Flan-T5 is based on T5, a sequence-to-sequence model trained on a denoising objective, that has been further instruction-finetuned on a battery of tasks. This has been shown to promote better responses to instructions, both with and without demonstrations (Chung et al., 2022). We evaluated the GPT models through the OpenAI API and the Flan-T5 using the HuggingFace library (Wolf et al., 2020). See Table 1 for a summary of the models, and Appendix C for implementation details.
We compared these models against a baseline that randomly outputs 0 to m = 3 objects from the set of objects that appeared in the same clauses as the box in question. Note that this baseline is much stronger than a fully random baseline that selects outputs from all mentioned objects.

heading: 4.2 Prompting and Demonstrations
text: Our prompts consist of: (a) a general instruction for the task, (b) two examples of the task to demonstrate the expected format, (c) an initial state description followed by a series of operations, and (d) an incomplete sentence Box N contains ___ that the model should complete (see Appendix E for full prompts). We used demonstrations that output the state of all boxes at once. However, in early experiments, Flan-T5 frequently only output the state of the first box even when the in-context demonstrations contained final descriptions of all box states.
6For GPT-3.5, we were able to instruct the model to output predictions in a consistent format without any in-context demonstrations, thus making it a true zero-shot experiment. We discuss the results of this experiment in Appendix D.
7 https://beta.openai.com/docs/models/gpt-3 8 https://beta.openai.com/docs/model-index-for-researchers
Therefore, for Flan-T5, we adjusted the prompts to output each box individually.9
Demonstration/Test Mismatch for Form-meaning Disentanglement (AltForms) As discussed in Sections 3.1–3.2, we additionally experimented with a setup where the demonstration and test examples were mismatched in the form of the descriptions of the states and operations. Under this setup, the models were evaluated with set of object names and phrasings of the state and operation descriptions that were different from the demonstration examples (see Table 2). Except for the determiner the and the preposition into, the two sets share no words (although subwords may be shared depending on tokenization). Scenarios with Greater Context Dependence (MoveContents, AmbiRef) As also discussed in Sections 3.1–3.2, we experimented with two additional scenarios with greater context dependence. The first scenario (MoveContents) introduces an additional operation Move contents of Box N to Box M that moves all objects in Box N to Box M, that does not provide an explicit enumeration of objects being moved. This requires the model to rely on the preceding description to identify the set of objects being moved, further removing room for heuristics that allow the prediction of the final state without composing the initial description and the operations in temporal order. The second scenario (AmbiRef) adds adjectival modification to object names (e.g., the big brain and the small brain),
9To verify that this difference in the task format does not underestimate the accuracy of the GPT models, we also conducted an experiment in which we prompted GPT to only output the state of one box at a time. We found that, contrary to the Flan-T5 models, the accuracy of GPT was lower when we prompted it to output individual boxes than prompting it to output all boxes, so we can rule out that this difference in the task format is underestimating GPT’s performance.
where the adjective can be omitted in some of the operations, depending on the state of the box being described. Specifically, the modifier is dropped if there is only one object of a specific type in a box (e.g., Move the brain from Box 1 to Box 2 if there is only one brain in Box 1). When predicting the contents of a box, the model is instructed to output the object type with the correct adjective to fully disambiguate the referring expression for the object. This again requires composition of the initial description and the operations in temporal order to correctly interpret the otherwise ambiguous object mentions. (See Appendix B for example scenarios and operations.)

heading: 4.3 Evaluation
text: We estimated the entity tracking capacity of the models by computing the accuracy of predicting the contents of each box after each operation. Given that we rely on arbitrary-length cloze completion to predict the contents, we had to score unconstrained generations. While we only considered instances as correct where the output mentions all objects (and no additional objects) in a given box, our evaluation allowed for minor deviations from the exact form of the response. Namely, we allowed the objects to appear in any order, the object names may be separated by commas or and, and both complete noun phrases with a determiner (e.g., the furby) and bare nouns (e.g., furby) are considered correct.
The task becomes intrinsically harder as more operations are applied to a specific box, since the initial state description needs to be combined sequentially with all subsequent operations. Further, similarly to our observation in Section 2, every operation changes the state of at most two boxes.
This implies that the number of datapoints corresponding to fewer operations is much greater than the number of datapoints for more operations. For these reasons, we report accuracy as a function of the number of operations affecting a particular box rather than reporting aggregate accuracy, and show all results with 95% confidence intervals.

heading: 4.4 Results
text: Figure 2 shows the prediction accuracy for different number of operations that affected a box (e.g., 3 indicates that three operations changed the content of the box after the initial state). The left panel shows the instances where the probed state differed from the initial state; the right panel shows the instances where the probed state was the same as the initial state. As the left panel shows, only GPT-3.5 text-davinci-003 consistently outperformed the (strong) random baseline. While, not surprisingly, the accuracy of this model also decreases as the number of operations increases, it still correctly predicted all contents of a box after 7 operations in more than 25% of the cases. The Flan-T5 models, on the other hand, seemed to ignore the operations and primarily predicted the initial state description, as indicated by the consistently high accuracy when the final state matches the initial state (right panel), as well as the consistently low accuracy when the final state deviates from the initial state (left panel). GPT-3 davinci also primarily repeated the initial state, but as indicated by the steep decrease in the right panel, it was distracted by intervening operations even when repeating the initial state. Form-meaning Disentanglement We additionally evaluated text-davinci-003, the only model that exhibited a non-trivial entity tracking capacity in the first set of results, under the AltForms setup as described in Section 4.2 where the demonstration examples have low lexical overlap with the test examples. Figure 3 shows the prediction accuracy of text-davinci-003 on a representative subsample10 of our data. The blue line represents the performance when the descriptions in the demonstra-
10For each number of operations n affecting a box, we sampled 100 states with at least one example with n operations.
tion and the test examples have low lexical overlap. As the comparison to the original results (red line) shows, the disjoint demonstrations did lead to a small drop in performance when there were more than two operations affecting a box. Nevertheless, text-davinci-003 was able to predict the correct state of entities in many cases, further adding support for its non-trivial entity tracking capacity.
Greater Context Dependence Finally, we evaluated text-davinci-003 on two additional datasets described in Section 4.2: the AmbiRef dataset where adjectival modification can be omitted depending on the current context, and the MoveContents dataset where a new operation Move contents of Box N to M that requires the contents of Box N to be contextually identified. Figure 4 shows the performance of text-davinci-003 on these two datasets compared to our random baseline. As these plots show, even in these more challenging
cases, we observe non-trivial entity tracking abilities. However, as the number of operations grows, performance rapidly approaches the random baseline, suggesting that entity tracking becomes increasingly more brittle as more state changes need to be considered jointly.

heading: 4.5 Discussion
text: Our results show that among the models we evaluated, only GPT-3.5 text-davinci-003 exhibit non-trivial entity tracking behavior. While its performance does decrease as the number of operations increases, the model still produced many accurate predictions even after six or seven sequences of operations. Furthermore, through the AltForms experiment with low lexical overlap between demonstration/test, we also ruled out the possibility that the demonstrations are teaching the model this task or that the model is primarily relying on superficial slot-filling heuristics. Therefore, we conclude that text-davinci-003 does have some capacity to track discourse entities in linguistically expressed contexts. To a lesser extent, we also observed this capacity in the highly context-dependent AmbiRef and MoveContents experiments. This provides further evidence against superficial heuristics in the performance we observe; at the same time, this highlights scenarios in which even the most recent models exhibit difficulties.
On the other hand, entity tracking behavior did not surface in GPT-3 davinci (likely of similar size as GPT-3.5 text-davinci-003), a model pretrained primarily on text corpora on the next word prediction objective. This was also true for denoising models that have been finetuned on many tasks combined with instructions and demonstrations: the Flan-T5 models also showed near-zero accuracy on non-trivial examples.
Our results overall show that there exists a language model that can perform entity tracking to some degree, but this capacity does not necessarily surface in all sufficiently large models trained on large corpora. Then, which factors are responsible for this difference? Given that davinci and text-davinci-003 differ along at least two dimensions (text-davinci-003 is based on a model that was trained on code and it was trained with additional human feedback (Ouyang et al., 2022); see Table 1), our initial results do not shed light on what exactly contributes to this difference. We therefore conducted a follow-up experiment where we compared a range of GPT-3 and GPT-3.5 models to
identify factors that contribute to the stark difference between davinci and text-davinci-003.11
Training on Code Encourages Entity Tracking Behavior As Table 1 shows, two key dimensions of variation across models are additional training on human feedback and pretraining on code. If additional training on human feedback imbues language models with the ability to track entities, all models except for GPT-3 davinci and GPT-3.5 code-davinci-002 should be able to track entities. If, on the other hand, pretraining on code leads to better entity tracking, we expect all GPT3.5 models to outperform GPT-3 on our task. As Figure 5 shows, GPT-3.5 models that have been trained on code systematically outperformed GPT3 models, including code-davinci-002 that was not trained on human feedback. This suggests that a substantial representation of code in the pretraining data is beneficial for a language model’s entity tracking capacity to surface.
A further question that our results so far do not answer is to what extent model size matters and whether models at the scale of Flan-T5 can also exhibit non-trivial entity tracking behavior. Since there exist no smaller models that have been trained with the same objective and training data as the GPT-3.5 models, we explore this question through finetuning experiments with T5.

heading: 5 Experiment 2: Finetuning
text: 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
11To limit inference costs, we used the same subsample of data as in the AltForms experiment.
training to the models.

heading: 5.1 Train/test splits
text: As discussed in Section 3.1, one challenge of evaluating entity tracking abilities is distinguishing this capacity from simple heuristics such as templatic slot-filling. We therefore designed various types of training/evaluation mismatches that block several possible shortcuts as described below. Base Split Here, we used the same format for training and evaluation examples. All initial states differed across training and evaluation to block simple slot-filling heuristics, as discussed in Section 3.2. NumOps Split The NumOps split restricts the maximum number of operations within a single example in the training set to 2, but includes up to 12 operations in the evaluation set. This split was intended to test whether a finetuned model is able to generalize to longer sequences of operations than it has seen during finetuning. Vocab/AltForms Splits The vocab split tests whether objects that are not part of the set of objects used during training can also be adequately tracked. We compiled a list of comparatively infrequent object names (e.g., pomelo, furby, Flav-RStraw; not in BNC) and sampled the training and test sets using two completely disjoint sets of object names. The training set used the infrequent object list and the test set used the original object list. The AltForm split follows the design described in Section 4.2. These splits aim to tease apart whether the model learns to associate specific words/phrases with the operations or whether finetuning leads to more generalizable entity tracking behavior. AmbiRef/MoveContents Splits The AmbiRef and MoveContents splits follow the design described in Section 4.2. These splits aim to test whether the model can learn to interpret operations that are underspecified without considering the current state of the affected boxes. For these splits, the
training and test examples share the same format.

heading: 5.2 Models
text: We evaluated T5-base, the best-performing model in Li et al. (2021), by finetuning it on each of the datasets described above. As an additional baseline, we compared against T5 with randomly initialized parameters trained directly on our datasets.

heading: 5.3 Results and Discussion
text: Pretrained T5 can Learn to Perform Entity Tracking As shown in Figure 6 (left), finetuning T5 leads to near-perfect accuracy on the base split. This suggests that the model is capable of learning this task. Training a randomly initialized T5 did not yield the same result: the accuracy of a model trained from random weights is considerably lower, due to the model almost exclusively predicting that a box is empty. These two results suggest that pretraining is crucial for the model to be able to learn this task. Furthermore, the model’s entity tracking capacity is robust to novel object names at test time, with only minor degradation on accuracy (Figure 6, middle). Training only on operation sequences with a maximum length of 2 (NumOps split) leads to a larger degradation in performance for longer operation sequences, but even for longer operation sequences, the model is able to infer the correct final state in more than 45% of the cases. Finally, the model performance does degrade substantially when the training examples have low lexical overlap with test examples (Figure 6, right). Nevertheless, model performance remains above the random baseline when the model is trained on up to 12 operations (pink line). If we trained only on up to two operations (blue line), however, the performance degradation was compounded and performance no longer exceeded the random baseline. These results suggest that finetuning on an entity tracking task does lead to entity tracking abilities that generalize
to many challenging scenarios. At the same time, however, performance rapidly degrades as the finetuning and evaluation splits become increasingly dissimilar, and it remains an open question to what extent finetuning on a limited domain such as the boxes environment transfers to more general entity tracking abilities in more naturalistic discourse. Interpreting Context-Dependent Operations Remains Challenging If we include operations that either contain ambiguous referring expressions (Figure 7, left) or a Move contents operation (Figure 7, right) in the finetuning and evaluation data, we see a slight degradation in performance and the model no longer achieves near-perfect accuracy, despite training and evaluating on examples of the same format. In both cases, the model almost exclusively made mistakes with examples that require the interpretation of context-dependent operations. This suggests that the context-dependent operations do compound the difficulty of entity tracking even when the models receive direct training signal.

heading: 6 Conclusion
text: 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,12 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 (Potts, 2020; Merrill et al., 2021, i.a.), 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.
Limitations
One limitation of this work is that we are only considering behavioral data which makes it difficult to
12This is also compatible with the observation by Sap et al. (2022) that GPT-3.5 performs considerably better than GPT-3 in answering factual questions about object location changes. Further, see Madaan et al. (2022) for other tasks for which pretraining on code seems to be beneficial.
establish a fully causal link between entity tracking capacities and high performance on our task. Entity tracking is a high-level linguistic behavior and many other capacities are necessary for achieving high accuracy on our task. Therefore, we cannot rule out that differences in some other capacity, such as interpreting sentences compositionally (see Bogin 2022 and Bogin et al. 2022 for evidence that GPT-3 and GPT-3.5 models differ in their compositional generalization behavior), are the main driver for the differences in behavior we see across models.
A possible criticism of our setup is that it requires short-term memory capacities that exceed the memory capacities of most, if not all, humans. That is, if we presented humans with the same input as the model, we would not expect them to be able to keep track of the contents of all 7 boxes due to memory limitations. Therefore we are potentially expecting models to do super-human entity tracking, a setup that has been criticized for model evaluations of other linguistic abilities (Lampinen, 2022). We nevertheless believe that our task is justified given the architecture of the evaluated models. Transformer-based models can look back to any token in the entire input sequence within their context window, so a proper comparison between humans and models would be to present humans with the full description in written form and let them re-read the description after being prompted to state the contents of a box. While we did not formally evaluate whether humans have this ability on a larger population, we personally did not have any trouble tracking the contents of boxes when we had access to the written description.
Relatedly, we designed our task such that the entire description fits within the context window of pretrained language models. However, as we mentioned in the introduction, entity tracking is an important ability for understanding long contexts and given the limited context window, our results do not apply to texts whose length exceeds a model’s context window, and likely different model architectures will be necessary to perform proper entity tracking for longer texts.
Further, while we found that GPT-3.5 models as well as finetuned T5 models can track entities in our task with higher accuracy than a strong random baseline, our results also indicate that this behavior is not very stable once several operations act on an entity. Our results should therefore not be taken
as justification for using these models for critical applications where high accuracy is needed.
Lastly, we only evaluated English models in this work. Given that we showed that even without high lexical overlap between the training and evaluation examples, models can keep track of entities to some extent, it seems likely that our results also apply to other languages. However, whether this actually the case remains an open question.

heading: Acknowledgements
text: We thank Jacob Andreas, Ellie Pavlick, Allyson Ettinger, Tal Linzen, Will Merrill, and the members of the NYU Computation and Psycholinguistics lab for discussions, and Belinda Li for sharing model outputs and details about their data preparation procedures and experiments. We thank Cookie for contributing to the authorship decision and for emotional support. This research was conducted in part through the NYU IT High Performance Computing resources, services, and staff expertise, and it was supported by the NSF under Grant #2030859 to the Computing Research Association for the CIFellows Project and the European Research Council (ERC) under the European Union’s Horizon 2020 Research and Innovation Program (Grant Agreement #948878). Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation nor the Computing Research Association.

heading: B Example Input-Output Pairs
text: (3) shows an example input-output pair from our base dataset (NumOps on Box 6 = 2).
(3) a. INPUT: Box 0 contains the painting, Box 1 contains the bell, Box 2 contains the guitar, Box 3 contains the egg and the mirror and the sheet, Box 4 contains the chemical, Box 5 contains the disk and the wire, Box 6 contains the glass and the knife. Move the glass from Box 6 to Box 4. Put the gift into Box 5. Move the guitar from Box 2 to Box 6. Put the milk into Box 4. Remove the mirror and the sheet from
Box 3. Box 6 __ b. OUTPUT: contains the guitar and the
knife.
(4) shows an example input-output pair from the dataset with operations that contain referring expressions that are ambiguous without considering the initial state and the previous operations (NumOps on Box 6 = 2). Note that in order to correctly interpret Move the guitar from Box 2 to Box 6, the model has to consider the information that the blue guitar (as opposed to the red guitar) was in Box 2 prior to the operation. Hence, this operation is ambiguous out of context.
(4) a. INPUT: Box 0 contains the yellow book and the green flower and the red guitar, Box 1 contains the small bomb and the small book and the blue bone, Box 2 contains the blue guitar, Box 3 contains the blue bell, Box 4 contains the green paper and the yellow note and the yellow television, Box 5 contains the yellow bell, Box 6 is empty. Move the guitar from Box 2 to Box 6. Put the blue wire and the big television into Box 5. Move the flower from Box 0 to Box 6. Box 6 __
b. OUTPUT: contains the blue guitar and the green flower.
(5) shows an example input-output pair from the dataset with the Move contents of operation (NumOps on Box 6 = 2). Note that in order to correctly interpret Move the contents of Box 2 to Box 6, the model has to consider the information that the tea (as opposed to the red guitar) was in Box 2 prior to the operation.
(5) a. INPUT: Box 0 contains the fan and the gift and the letter, Box 1 contains the beer and the mirror and the tie, Box 2 contains the tea, Box 3 contains the boot, Box 4 contains the coat and the plate and the shirt, Box 5 contains the bottle, Box 6 is empty. Move the contents of Box 2 to Box 6. Put the dress and the painting into Box 5. Move the letter from Box 0 to Box 6. Box 6 __
b. OUTPUT: contains the letter and the tea.
C Implementation Details
In-context For GPT-3, we used the OpenAI API. We used greedy decoding with the temperature parameter set to 0, used a maximum target generation length of 150, and used the newline character as an additional stop token. Inference time and model size is not available for GPT-3, but we generated about 16 million tokens in total. All GPT-3 experiments that we report here in the paper were conducted in January 2023.
For Flan-T5, we used a beam size of 3 and a maximum target generation length of 256. Other hyperparameters were kept as the default values of T5ConditionalGeneration in the HuggingFace library. Inference for T5-XL took about 5 hours on a single A100 GPU, and for T5 base, about 2 hours. Finetuning We finetuned T5 for a single epoch using a batch size of 8 and a learning rate of 1×10−4. In initial explorations, increasing the number of finetuning epochs did not yield substantial gains on development set performance, and was sometimes even harmful. Training and inference took around 3 hours on a single RTX8000 GPU.

heading: D Additional Results
text: D.1 GPT-3.5 Zero-shot Results
As mentioned in the main text, GPT-3.5 was able to to output the contents of boxes in the correct format without any in-context demonstrations (see Table 5 for the prompt template).
Figure 9 compares the performance of text-davinci-003 in the 2-shot setting to the zero-shot setting. Model performance degraded slightly without demonstration examples, which suggests that the examples are indeed helpful in guiding the model to correctly perform the task. Nevertheless, we observed non-negligible performance even in the zero-shot setting, which corroborates the conclusion that GPT 3.5 exhibits non-trivial entity tracking capacities.
D.2 GPT-4 Results
We also probed the more recently released GPT-4 model (OpenAI, 2023) through the OpenAI API. Figure 8 compares the performance of GPT-3.5 text-davinci-003 and GPT-4 on the base dataset (left) as well as the more challenging MoveContents (middle) and AmbiRef (right) datasets discussed in Section 4.2. Across all datasets, GPT-4
performed considerably better than GPT-3.5. But similarly to GPT-3.5, performance degraded as more operations affected a box. Considering that no public information about the training data, training procedure or architectural details about GPT-4 exists, we cannot draw any conclusions about the reason behind the improved performance of GPT-4 compared to previous models.

heading: E Prompts
text: Table 3 shows the 2-shot prompts used for incontext experiments. Table 4 shows the 2-shot prompts used for the AltForms experiments, where the demonstrations contain descriptions that have lower lexical overlap with the test examples. Table 5 shows the zero-shot prompts discussed in Appendix D. All prompts are available at https:// github.com/sebschu/entity-tracking-lms.

heading: F Dataset Statistics and License Information
text: See Tables 6 and 7 for descriptive statistics of our datasets. All datasets are released under the GNU General Public License v3.0.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
Limitations on page 10.
A2. Did you discuss any potential risks of your work? Not applicable. We evaluate existing models, so we don’t see any potential risks other than misinterpreting our results. We hope we managed to limit the potential for misunderstanding by discussing in detail how to interpret the results.
3 A3. Do the abstract and introduction summarize the paper’s main claims? Abstract, Section 1 ("Contributions" paragraph on p2)
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? We developed a novel procedure for programmatically generating datasets and we will release that code and the generated datasets upon publication (described in Section 3). We also discuss an analysis of the probing datasets from Li et al. (2021) in Section 2 and Appendix A.
3 B1. Did you cite the creators of artifacts you used? We cited the pretrained models that we used (Sections 4 and 5). We cite Li et al. and relevant source datasets in Section 2 and Appendix A.
3 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? See Appendix E.
7 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Our data mostly concerns moving objects from one box to another, so do not envision a risk scenario of use outside of research contexts.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. We generated the data from a set lexicon and templates and therefore, we can rule out that any identifying information is contained in the data.
3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? See Section 3.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. See Appendix E.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
C 3 Did you run computational experiments? See Sections 4 and 5.
3 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? See Section 4 and Appendix C.
3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found hyperparameter values? See Appendices B, C, and D.
3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? See Sections 4 and 5.
3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? See Section 4 and Appendix C.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: 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. Limitations
One limitation of this work is that we are only considering behavioral data which makes it difficult to
12This is also compatible with the observation by Sap et al. (2022) that GPT-3. 5 performs considerably better than GPT-3 in answering factual questions about object location changes. Further, see Madaan et al. (2022) for other tasks for which pretraining on code seems to be beneficial. establish a fully causal link between entity tracking capacities and high performance on our task. Entity tracking is a high-level linguistic behavior and many other capacities are necessary for achieving high accuracy on our task. Therefore, we cannot rule out that differences in some other capacity, such as interpreting sentences compositionally (see Bogin 2022 and Bogin et al. 2022 for evidence that GPT-3 and GPT-3. 5 models differ in their compositional generalization behavior), are the main driver for the differences in behavior we see across models. A possible criticism of our setup is that it requires short-term memory capacities that exceed the memory capacities of most, if not all, humans. That is, if we presented humans with the same input as the model, we would not expect them to be able to keep track of the contents of all 7 boxes due to memory limitations. Therefore we are potentially expecting models to do super-human entity tracking, a setup that has been criticized for model evaluations of other linguistic abilities (Lampinen, 2022). We nevertheless believe that our task is justified given the architecture of the evaluated models. Transformer-based models can look back to any token in the entire input sequence within their context window, so a proper comparison between humans and models would be to present humans with the full description in written form and let them re-read the description after being prompted to state the contents of a box. While we did not formally evaluate whether humans have this ability on a larger population, we personally did not have any trouble tracking the contents of boxes when we had access to the written description. Relatedly, we designed our task such that the entire description fits within the context window of pretrained language models. However, as we mentioned in the introduction, entity tracking is an important ability for understanding long contexts and given the limited context window, our results do not apply to texts whose length exceeds a model’s context window, and likely different model architectures will be necessary to perform proper entity tracking for longer texts. Further, while we found that GPT-3. 5 models as well as finetuned T5 models can track entities in our task with higher accuracy than a strong random baseline, our results also indicate that this behavior is not very stable once several operations act on an entity. Our results should therefore not be taken
as justification for using these models for critical applications where high accuracy is needed. Lastly, we only evaluated English models in this work. Given that we showed that even without high lexical overlap between the training and evaluation examples, models can keep track of entities to some extent, it seems likely that our results also apply to other languages. However, whether this actually the case remains an open question