SlowGuess/ACL_2020_OCR huggingface dataset:
Shubham Toshniwal<sup>1</sup>, Allyson Ettinger<sup>2</sup>, Kevin Gimpel<sup>1</sup>, Karen Livescu<sup>1</sup>
<sup>1</sup>Toyota Technological Institute at Chicago
$^{2}$ Department of Linguistics, University of Chicago
{shtoshni, kgimpel, klivescu}@ttic.edu, aettinger@uchicago.edu
We propose PeTra, a memory-augmented neural network designed to track entities in its memory slots. PeTra is trained using sparse annotation from the GAP pronoun resolution dataset and outperforms a prior memory model on the task while using a simpler architecture. We empirically compare key modeling choices, finding that we can simplify several aspects of the design of the memory module while retaining strong performance. To measure the people tracking capability of memory models, we (a) propose a new diagnostic evaluation based on counting the number of unique entities in text, and (b) conduct a small scale human evaluation to compare evidence of people tracking in the memory logs of PeTra relative to a previous approach. PeTra is highly effective in both evaluations, demonstrating its ability to track people in its memory despite being trained with limited annotation.
Understanding text narratives requires maintaining and resolving entity references over arbitrary-length spans. Current approaches for coreference resolution (Clark and Manning, 2016b; Lee et al., 2017, 2018; Wu et al., 2019) scale quadratically (without heuristics) with length of text, and hence are impractical for long narratives. These models are also cognitively implausible, lacking the incrementality of human language processing (Tanenhaus et al., 1995; Keller, 2010). Memory models with finite memory and online/quasi-online entity resolution have linear runtime complexity, offering more scalability, cognitive plausibility, and interpretability.
Memory models can be viewed as general problem solvers with external memory mimicking a Turing tape (Graves et al., 2014, 2016). Some of the earliest applications of memory networks
in language understanding were for question answering, where the external memory simply stored all of the word/sentence embeddings for a document (Sukhbaatar et al., 2015; Kumar et al., 2016). To endow more structure and interpretability to memory, key-value memory networks were introduced by Miller et al. (2016). The key-value architecture has since been used for narrative understanding and other tasks where the memory is intended to learn to track entities while being guided by varying degrees of supervision (Henaff et al., 2017; Liu et al., 2018a,b, 2019a).
We propose a new memory model, PeTra, for entity tracking and coreference resolution, inspired by the recent Referential Reader model (Liu et al., 2019a) but substantially simpler. Experiments on the GAP (Webster et al., 2018) pronoun resolution task show that PeTra outperforms the Referential Reader with fewer parameters and simpler architecture. Importantly, while Referential Reader performance degrades with larger memory, PeTra improves with increase in memory capacity (before saturation), which should enable tracking of a larger number of entities. We conduct experiments to assess various memory architecture decisions, such as learning of memory initialization and separation of memory slots into key/value pairs.
To test interpretability of memory models' entity tracking, we propose a new diagnostic evaluation based on entity counting—a task that the models are not explicitly trained for—using a small amount of annotated data. Additionally, we conduct a small scale human evaluation to assess quality of people tracking based on model memory logs. PeTra substantially outperforms Referential Reader on both measures, indicating better and more interpretable tracking of people.<sup>1</sup>
$^{1}$ Code available at https://github.com/shtoshni92/petra
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5415-5428 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
Figure 1: Illustration of memory cell updates in an example sentence where IG = ignore, OW = overwrite, CR = coref. Different patterns indicate the different entities, and an empty pattern indicates that the cell has not been used. The updated memory cells at each time step are highlighted.
Figure 2 depicts PeTra, which consists of three components: an input encoder that given the tokens generates the token embeddings, a memory module that tracks information about the entities present in the text, and a controller network that acts as an interface between the encoder and the memory.
Figure 2: Proposed model.
2.1 Input Encoder
Given a document consisting of a sequence of tokens  $\{w_{1},\dots ,w_{T}\}$ , we first pass the document through a fixed pretrained BERT model (Devlin et al., 2019) to extract contextual token embeddings. Next, the BERT-based token embeddings are fed into a single-layer unidirectional Gated Recurrent Unit (GRU) (Cho et al., 2014) running left-to-right to get task-specific token embeddings  $\{h_1,\dots ,h_T\}$ .
2.2 Memory
The memory  $\mathcal{M}_t$  consists of  $N$  memory cells. The  $i^{\mathrm{th}}$  memory cell state at time step  $t$  consists of a tuple  $(\boldsymbol{m}_t^i, u_t^i)$  where the vector  $\boldsymbol{m}_t^i$  represents the content of the memory cell, and the scalar  $u_t^i \in$
[0, 1] represents its recency of usage. A high value of  $u_{t}^{i}$  is intended to mean that the cell is tracking an entity that has been recently mentioned.
Initialization Memory cells are initialized to the null tuple, i.e.  $(0,0)$ ; thus, our memory is parameter-free. This is in contrast with previous entity tracking models such as EntNet (Henaff et al., 2017) and the Referential Reader (Liu et al., 2019a) where memory initialization is learned and the cells are represented with separate key and value vectors. We will later discuss variants of our memory with some of these changes.
2.3 Controller
At each time step  $t$  the controller network determines whether token  $t$  is part of an entity span and, if so, whether the token is coreferent with any of the entities already being tracked by the memory. Depending on these two variables, there are three possible actions:
(i) IgNORE: The token is not part of any entity span, in which case we simply ignore it.
(ii) OVERWRITE: The token is part of an entity span but is not already being tracked in the memory.
(iii) COREF: The token is part of an entity span and the entity is being tracked in the memory.
Therefore, the two ways of updating the memory are OVERWRITE and COREF. There is a strict ordering constraint to the two operations: OVERWRITE precedes COREF, because it is not possible to corefer with a memory cell that is not yet tracking anything. That is, the COREF operation cannot be applied to a previously unwritten memory cell, i.e. one with  $u_{t}^{i} = 0$ . Figure 1 illustrates an idealized version of this process.
Next we describe in detail the computation of the probabilities of the two operations for each memory cell at each time step  $t$ .
First, the entity mention probability  $e_t$ , which reflects the probability that the current token  $w_t$  is part of an entity mention, is computed by:
$$
e _ {t} = \sigma \left(\operatorname {M L P} _ {1} \left(\boldsymbol {h} _ {t}\right)\right) \tag {1}
$$
where  $\mathrm{MLP}_1$  is a multi-layer perceptron and  $\sigma$  is the logistic function.
Overwrite and Coref If the current token  $w_{t}$  is part of an entity mention, we need to determine whether it corresponds to an entity being currently tracked by the memory or not. For this we compute the similarity between the token embedding  $h_{t}$  and the contents of the memory cells currently tracking entities. For the  $i^{\text{th}}$  memory cell with memory vector  $m_{t-1}^{i}$  the similarity with  $h_{t}$  is given by:
$$
\operatorname {s i m} _ {t} ^ {i} = \operatorname {M L P} _ {2} \left(\left[ \boldsymbol {h} _ {t}; \boldsymbol {m} _ {t - 1} ^ {i}; \boldsymbol {h} _ {t} \odot \boldsymbol {m} _ {t - 1} ^ {i}; u _ {t - 1} ^ {i} \right]\right) \tag {2}
$$
where  $\mathrm{MLP}_2$  is a second MLP and  $\odot$  is the Hadamard (elementwise) product. The usage scalar  $u_{t - 1}^i$  in the above expression provides a notion of distance between the last mention of the entity in cell  $i$  and the potential current mention. The higher the value of  $u_{t - 1}^i$ , the more likely there was a recent mention of the entity being tracked by the cell. Thus  $u_{t - 1}^i$  provides an alternative to distance-based features commonly used in pairwise scores for spans (Lee et al., 2017).
Given the entity mention probability  $e_t$  and similarity score  $sim_t^i$ , we define the coref score  $c s_t^i$  as:
$$
c s _ {t} ^ {i} = s i m _ {t} ^ {i} - \infty \cdot \mathbb {1} [ u _ {t - 1} ^ {i} = 0 ] \qquad (3)
$$
where the second term ensures that the model does not predict coreference with a memory cell that has not been previously used, something not enforced by Liu et al. (2019a).<sup>2</sup> Assuming the coref score for a new entity to be  $0$ ,<sup>3</sup> we compute the coref probability  $c_t^i$  and new entity probability  $n_t$  as follows:
$$
\left( \begin{array}{c} c _ {t} ^ {1} \\ \vdots \\ c _ {t} ^ {N} \\ n _ {t} \end{array} \right) = e _ {t} \cdot \text {s o f t m a x} \left( \begin{array}{c} c s _ {t} ^ {1} \\ \vdots \\ c s _ {t} ^ {N} \\ 0 \end{array} \right) \tag {4}
$$
Based on the memory usage scalars  $u_{t}^{i}$  and the new entity probability  $n_t$ , the overwrite probability for
each memory cell is determined as follows:
$$
o _ {t} ^ {i} = n _ {t} \cdot \mathbb {1} _ {i = \arg \min  _ {j} u _ {t - 1} ^ {j}} \tag {5}
$$
Thus we pick the cell with the lowest usage scalar  $u_{t-1}^j$  to OVERWRITE. In case of a tie, a cell is picked randomly among the ones with the lowest usage scalar. The above operation is non-differentiable, so during training we instead use
$$
o _ {t} ^ {i} = n _ {t} \cdot \operatorname {G S} \left(\frac {1 - u _ {t - 1} ^ {i}}{\tau}\right) _ {i} \tag {6}
$$
where  $\mathrm{GS}(.)$  refers to Gumbel-Softmax (Jang et al., 2017), which makes overwrites differentiable.
For each memory cell, the memory vector is updated based on the three possibilities of ignoring the current token, being coreferent with the token, or considering the token to represent a new entity (causing an overwrite):
$$
\begin{array}{l} \boldsymbol {m} _ {t} ^ {i} = \overbrace {(1 - \left(o _ {t} ^ {i} + c _ {t} ^ {i}\right)) \boldsymbol {m} _ {t - 1} ^ {i}} ^ {\text {I G N O R E}} + \overbrace {o _ {t} ^ {i} \cdot \boldsymbol {h} _ {t}} ^ {\text {O V E R W R I T E}} \tag {7} \\ + \underbrace {c _ {t} ^ {i} \cdot \mathrm {M L P} _ {3} ([ \boldsymbol {h} _ {t} ; \boldsymbol {m} _ {t - 1} ^ {i} ])} _ {\text {C O R E F}} \\ \end{array}
$$
In this expression, the coreference term takes into account both the previous cell vector  $\boldsymbol{m}_{t-1}^{i}$  and the current token representation  $\boldsymbol{h}_t$ , while the overwrite term is based only on  $\boldsymbol{h}_t$ . In contrast to a similar memory update equation in the Referential Reader which employs a pair of GRUs and MLPs for each memory cell, our update parameter uses just  $\mathrm{MLP}_3$  which is memory cell-agnostic.
Finally, the memory usage scalar is updated as
$$
u _ {t} ^ {i} = \min  \left(1, o _ {t} ^ {i} + c _ {t} ^ {i} + \gamma \cdot u _ {t - 1} ^ {i}\right) \tag {8}
$$
where  $\gamma \in (0,1)$  is the decay rate for the usage scalar. Thus the usage scalar  $u_{t}^{i}$  keeps decaying with time unless the memory is updated via OVERWRITE or COREF in which case the value is increased to reflect the memory cell's recent use.
Memory Variants In vanilla PeTra, each memory cell is represented as a single vector and the memory is parameter-free, so the total number of model parameters is independent of memory size. This is a property that is shared with, for example, differentiable neural computers (Graves et al., 2016). On the other hand, recent models for entity tracking, such as the EntNet (Henaff et al., 2017) and the Referential Reader (Liu et al., 2019a), learn
2A threshold higher than 0 can also be used to limit coreference to only more recent mentions. The new entity coref score is a free variable that can be assigned any value, since only the relative value matters.
memory initialization parameters and separate the memory cell into key-value pairs. To compare these memory cell architectures, we investigate the following two variants of PeTra:
1. PeTra + Learned Initialization: memory cells are initialized at  $t = 0$  to learned parameter vectors.
2. PeTra + Fixed Key: a fixed dimensions of each memory cell are initialized with learned parameters and kept fixed throughout the document read, as in EntNet (Henaff et al., 2017).
Apart from initialization, the initial cell vectors are also used to break ties for overwrites in Eqs. (5) and (6) when deciding among unused cells (with  $u_{t}^{i} = 0$ ). The criterion for breaking the tie is the similarity score computed using Eq. (2).
2.4 Coreference Link Probability
The probability that the tokens  $w_{t_1}$  and  $w_{t_2}$  are coreferential according to, say, cell  $i$  of the memory depends on three things: (a)  $w_{t_1}$  is identified as part of an entity mention and is either overwritten to cell  $i$  or is part of an earlier coreference chain for an entity tracked by cell  $i$ , (b) Cell  $i$  is not overwritten by any other entity mention from  $t = t_1 + 1$  to  $t = t_2$ , and (c)  $w_{t_2}$  is also predicted to be part of an entity mention and is coreferential with cell  $i$ . Combining these factors and marginalizing over the cell index results in the following expression for the coreference link probability:
$$
\begin{array}{l} P _ {\mathrm {C L}} \left(w _ {t _ {1}}, w _ {t _ {2}}\right) \\ = \sum_ {i = 1} ^ {N} \left(o _ {t _ {1}} ^ {i} + c _ {t _ {1}} ^ {i}\right) \cdot \prod_ {j = t _ {1} + 1} ^ {t _ {2}} \left(1 - o _ {j} ^ {i}\right) \cdot c _ {t _ {2}} ^ {i} \tag {9} \\ \end{array}
$$
2.5 Losses
The GAP (Webster et al., 2018) training dataset is small and provides sparse supervision with labels for only two coreference links per instance. In order to compensate for this lack of supervision, we use a heuristic loss  $\mathcal{L}_{ent}$  over entity mention probabilities in combination with the end task loss  $\mathcal{L}_{coref}$  for coreference. The two losses are combined with a tunable hyperparameter  $\lambda$  resulting in the following total loss:  $\mathcal{L} = \mathcal{L}_{coref} + \lambda \mathcal{L}_{ent}$ .
2.5.1 Coreference Loss
The coreference loss is the binary cross entropy between the ground truth labels for mention pairs
and the coreference link probability  $P_{\mathrm{CL}}$  in Eq. (9). Eq. (9) expects a pair of tokens while the annotations are on pairs of spans, so we compute the loss for all ground truth token pairs:  $\mathcal{L}_{coref} =$
$$
\sum_ {\left(s _ {a}, s _ {b}, y _ {a b}\right) \in \mathrm {G}} \left(\sum_ {w _ {a} \in s _ {a}} \sum_ {w _ {b} \in s _ {b}} H \left(y _ {a b}, P _ {\mathrm {C L}} \left(w _ {a}, w _ {b}\right)\right)\right)
$$
where  $G$  is the set of annotated span pairs and  $H(p, q)$  represents the cross entropy of the distribution  $q$  relative to distribution  $p$ .
Apart from the ground truth labels, we use "implied labels" in the coreference loss calculation. For handling multi-token spans, we assume that all tokens following the head token are coreferential with the head token (self-links). We infer more supervision based on knowledge of the setup of the GAP task. Each GAP instance has two candidate names and a pronoun mention with supervision provided for the {name, pronoun} pairs. By design the two names are different, and therefore we use them as a negative coreference pair.
Even after the addition of this implied supervision, our coreference loss calculation is restricted to the three mention spans in each training instance; therefore, the running time is  $\mathcal{O}(T)$  for finite-sized mention spans. In contrast, Liu et al. (2019a) compute the above coreference loss for all token pairs (assuming a negative label for all pairs outside of the mentions), which results in a runtime of  $\mathcal{O}(T^3)$  due to the  $\mathcal{O}(T^2)$  pairs and  $\mathcal{O}(T)$  computation per pair, and thus will scale poorly to long documents.
2.5.2 Entity Mention Loss
We use the inductive bias that most tokens do not correspond to entities by imposing a loss on the average of the entity mention probabilities predicted across time steps, after masking out the labeled entity spans. For a training instance where spans  $s_A$  and  $s_B$  correspond to the person mentions and span  $s_P$  is a pronoun, the entity mention loss is
$$
\mathcal {L} _ {e n t} = \frac {\sum_ {t = 1} ^ {T} e _ {t} \cdot m _ {t}}{\sum_ {t = 1} ^ {T} m _ {t}}
$$
where  $m_t = 0$  if  $w_t \in s_A \cup s_B \cup s_P$  and  $m_t = 1$  otherwise.
Each GAP instance has only 3 labeled entity mention spans, but the text typically has other entity mentions that are not labeled. Unlabeled entity mentions will be inhibited by this loss. However, on average there are far more tokens outside entity spans than inside the spans. In experiments without
this loss, we observed that the model is susceptible to predicting a high entity probability for all tokens while still performing well on the end task of pronoun resolution. We are interested in tracking people beyond just the entities that are labeled in the GAP task, for which this loss is very helpful.
GAP is a gender-balanced pronoun resolution dataset introduced by Webster et al. (2018). Each instance consists of a small snippet of text from Wikipedia, two spans corresponding to candidate names along with a pronoun span, and two binary labels indicating the coreference relationship between the pronoun and the two candidate names. Relative to other popular coreference datasets (Pradhan et al., 2012; Chen et al., 2018), GAP is comparatively small and sparsely annotated. We choose GAP because its small size allows us to do extensive experiments.
3.2 Model Details
For the input BERT embeddings, we concatenate either the last four layers of  $\mathrm{BERT}_{\mathrm{BASE}}$ , or layers 19-22 of  $\mathrm{BERT}_{\mathrm{LARGE}}$  since those layers have been found to carry the most information related to coreference (Liu et al., 2019b). The BERT embeddings are fed to a 300-dimensional GRU model, which matches the dimensionality of the memory vectors.
We vary the number of memory cells  $N$  from 2 to 20. The decay rate for the memory usage scalar  $\gamma$  is 0.98. The MLPs used for predicting the entity probability and similarity score consist of two 300-dimensional ReLU hidden layers. For the Fixed Key variant of PeTra we use 20 dimensions for the learned key vector and the remaining 280 dimensions as the value vector.
All models are trained for a maximum of 100 epochs with the Adam optimizer (Kingma and Ba, 2015). The learning rate is initialized to  $10^{-3}$  and is reduced by half, until a minimum of  $10^{-4}$ , whenever there is no improvement on the validation performance for the last 5 epochs. Training stops when there is no improvement in validation performance for the last 15 epochs. The temperature  $\tau$  of the Gumbel-Softmax distribution used in the OVERWRITE operation is initialized to 1 and halved every 10 epochs. The coreference loss terms in
Section 2.5.1 are weighted differently for different coreference links: (a) self-link losses for multi-token spans are given a weight of 1, (b) positive coreference link losses are weighted by 5, and (c) negative coreference link losses are multiplied by 50. To prevent overfitting: (a) we use early stopping based on validation performance, and (b) apply dropout at a rate of 0.5 on the output of the GRU model. Finally, we choose  $\lambda = 0.1$  to weight the entity prediction loss described in Section 2.5.2.
3.4 People Tracking Evaluation
One of the goals of this work is to develop memory models that not only do well on the coreference resolution task, but also are interpretable in the sense that the memory cells actually track entities. Hence in addition to reporting the standard metrics on GAP, we consider two other ways to evaluate memory models.
As our first task, we propose an auxiliary entity-counting task. We take 100 examples from the GAP validation set and annotate them with the number of unique people mentioned in them.<sup>4</sup> We test the models by predicting the number of people from their memory logs as explained in Section 3.5. The motivation behind this exercise is that if a memory model is truly tracking entities, then its memory usage logs should allow us to recover this information.
To assess the people tracking performance more holistically, we conduct a human evaluation in which we ask annotators to assess the memory models on people tracking performance, defined as:(a) detecting references to people including pronouns, and (b) maintaining a 1-to-1 correspondence between people and memory cells. For this study, we pick the best run (among 5 runs) of PeTra and the Referential Reader for the 8-cell configuration using BERT<sub>BASE</sub> (PeTra: 81 F1; Referential Reader: 79 F1). Next we randomly pick 50 documents (without replacement) from the GAP dev set and split those into groups of 10 to get 5 evaluation sets. We shuffle the original 50 documents and follow the same steps to get another 5 evaluation sets. In the end, we have a total of 10 evaluation sets with 10 documents each, where each unique document belongs to exactly 2 evaluation sets.
We recruit 10 annotators for the 10 evaluation sets. The annotators are shown memory log visualizations as in Figure 5, and instructed to compare
4 In the GAP dataset, the only relevant entities are people.
(a) BERTBASE
(b) BERTLARGE
Figure 3: Mean F1 score on the GAP validation set as a function of the number of memory cells.
the models on their people tracking performance (detailed instructions in Appendix A.3). For each document the annotators are presented memory logs of the two models (ordered randomly) and asked whether they prefer the first model, prefer the second model, or have no preference (neutral).
GAP Given a pronoun span  $s_P$  and two candidate name spans  $s_A \& s_B$ , we have to predict binary labels for potential coreference links between  $(s_A, s_P)$  and  $(s_B, s_P)$ . Thus, for a pair of entity spans, say  $s_A$  and  $s_P$ , we predict the coreference link probability as:
$$
P_{\mathrm{CL}}(s_{A},s_{P}) = \max_{w_{A}\in s_{A},w_{P}\in s_{P}}P_{\mathrm{CL}}(w_{A},w_{P})
$$
where  $P_{\mathrm{CL}}(w_A, w_P)$  is calculated using the procedure described in Section 2.4<sup>5</sup>. The final binary prediction is made by comparing the probability against a threshold.
Counting unique people For the test of unique people counting, we discretize the overwrite operation, which corresponds to new entities, against a threshold  $\alpha$  and sum over all tokens and all memory cells to predict the count as follows:
$$
\# \text {u n i q u e p e o p l e} = \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {N} \mathbb {1} \left[ o _ {t} ^ {i} \geq \alpha \right]
$$
3.6 Evaluation Metrics
For GAP we evaluate models using F-score.6 First, we pick a threshold from the set  $\{0.01, 0.02, \dots\}$
$1.00\}$  which maximizes the validation F-score. This threshold is then used to evaluate performance on the GAP test set.
For the interpretability task of counting unique people, we choose a threshold that minimizes the absolute difference between ground truth count and predicted count summed over the 100 annotated examples. We select the best threshold from the set  $\{0.01, 0.02, \dots, 1.00\}$ . The metric is then the number of errors corresponding to the best threshold.
3.7 Baselines
The Referential Reader (Liu et al., 2019a) is the most relevant baseline in the literature, and the most similar to PeTra. The numbers reported by Liu et al. (2019a) are obtained by a version of the model using  $\mathrm{BERT}_{\mathrm{BASE}}$ , with only two memory cells. To compare against PeTra for other configurations, we retrain the Referential Reader using the code made available by the authors. $^8$