SlowGuess/NAACL_2021_OCR huggingface dataset:
Vivek Subramanian, Matthew Engelhard, Samuel Berchuck, Liquin Chen, Ricardo Henao, and Lawrence Carin
{vivek.subramanian, matthew.engelhard, samuel.berchuck, liquor.chen, ricardo.henao, lcarin}@duke.edu
In many natural language processing applications, identifying predictive text can be as important as the predictions themselves. When predicting medical diagnoses, for example, identifying predictive content in clinical notes not only enhances interpretability, but also allows unknown, descriptive (i.e., text-based) risk factors to be identified. We here formalize this problem as predictive extraction and address it using a simple mechanism based on linear attention. Our method preserves differentiability, allowing scalable inference via stochastic gradient descent. Further, the model decomposes predictions into a sum of contributions of distinct text spans. Importantly, we require only document labels, not ground-truth spans. Results show that our model identifies semantically-cohesive spans and assigns them scores that agree with human ratings, while preserving classification performance.
Attention-based neural network architectures achieve human-level performance in many document classification tasks. However, understanding model predictions remains challenging. Common feature attribution methods are often inadequate, because the "features" of a document classification model – individual words or their embeddings – tend to have limited or ambiguous meaning in isolation, and must instead be interpreted in context. Rather than examining the importance of individual words and passing the contextualization task to the end-user, we may wish to extract distinct spans of text, such as sentences or paragraphs, and quantify the effect of each span on model predictions. However, the appropriate span boundaries depend on the document type, and processing all possible spans individually is computationally prohibitive.
In some settings, understanding model predictions can be as important as the predictions themselves. When predicting medical diagnoses from
clinical notes, for example, attributing predictions to specific note content assures clinicians that the model is not relying on data artifacts that are not clinically meaningful or generalizable. Moreover, this process may illuminate previously unknown risk factors that are described in clinical notes but not captured in a structured manner. Our work is motivated by the problem of autism spectrum disorder (ASD) diagnosis, in which many early symptoms are behavioral rather than physiologic, and are documented in clinical notes using multiple-word descriptions, not individual terms. Moreover, extended and nuanced descriptions are important in many common document classification tasks, for instance, the scoring of movie or food reviews.
Identifying important spans of text is a recurring theme in natural language processing. In extractive summarization, a document summary is created by selecting and concatenating important spans within a document (Narayan et al., 2018); and in many question answering tasks, including in the Stanford Question Answering Dataset (Rajpurkar et al., 2018), the goal is to identify a span within a paragraph of text that answers a given question. In both cases, training typically relies on ground truth spans, i.e., correct start and end positions are available during training, which the model learns to predict.
In contrast, our goal is to identify distinct spans within a document that, taken together, are sufficient to predict its associated label. In this task, which we call predictive extraction, ground truth spans are not available; instead, training is based on document labels alone, and without predefined spans, e.g., sentences or paragraphs. Moreover, similar to feature attribution methods, we wish to assign scores to each span such that predictions are effectively decomposed into the contributions of individual spans. In the current work, which for simplicity focuses on binary classification, we achieve this by summing individual span scores to
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5234-5258 June 6-11, 2021. ©2021 Association for Computational Linguistics
obtain the log-odds of a positive label.
Since correct start and end positions are not known, they are represented as latent variables that must be learned to  $(a)$  optimize classification performance, and  $(b)$  satisfy additional span constraints; in particular, we wish to ensure that spans are concise, and do not significantly overlap. A brute-force approach – in which all sets of spans satisfying these constraints are evaluated – is computationally intractable, as the number of possibilities is  $\mathcal{O}(n^k)$ , where  $n$  is the length of the document and  $k$  is the number of spans. Alternatively, predicting discrete start and end positions would introduce categorical latent variables, necessitating the use of a continuous relaxation (Jang et al., 2016; Maddison et al., 2016) or gradient estimation alternatives (Tucker et al., 2017). Instead, we formulate a simple but effective approach in which span representations are derived directly from a continuous (probabilistic) representation of the start and end positions, avoiding more computationally expensive gradient estimation; and the positions themselves, are predicted using linear attention. Our contributions are as follows:
- We define predictive extraction and describe its importance particularly for prediction tasks in which model performance exceeds human performance.
- We formulate SpanPredict, a neural network model for predictive extraction in which predicted log-odds are formulated as the sum of contributions of distinct spans.
- We quantify prediction and span selection performance on five binary classification tasks, including three real-world medical diagnosis prediction tasks.
- In the context of these studies, we quantify the effect of span constraints on performance.
Explaining neural network predictions is a well-known problem, one that is particularly challenging in natural language processing, due to the presence of complex semantic structure and interdependencies (Belinkov and Glass, 2019). The importance of individual words, or their embeddings, can be quantified using word-pooling strategies in which some words contribute to predictions, and others do not (Shen et al., 2018). In many settings, however, examining individual words in isolation provides limited insight. One solution is to ask the
model to generate an explanation along with each prediction (Zhang et al., 2016); inconveniently, explanations must be available during training.
Alternatively, explanations may be selected from within the document itself. This strategy is closely related to question answering and extractive summarization, in which text spans are selected to answer a given question or summarize a document, respectively. If correct spans are known during training, representations of candidate spans can be generated and used to evaluate each span as the possible answer to a question, or for inclusion in a document summary. Representations for all short spans can be generated via bidirectional recurrent neural networks (Lee et al., 2016), for example, or candidate spans can be limited to individual words and sentences (Cheng and Lapata, 2016).
Clinical notes contain redundant information as well as medical jargon and abbreviations, making meaningful text extraction more useful but also more challenging. Concept recognition and relation detection have been used to identify salient note content, which is then used to create a summary (Liang et al., 2019). Alternatively, the importance of specific content can be evaluated based on its presence or absence in subsequent notes; this concept has been used to train extractive summarization models using discharge summaries, which distill information collected during a clinical encounter (Alsentzer and Kim, 2018), and using subsequent notes, which are more likely to repeat earlier information if it is important (Liu et al., 2018).
In contrast to these methods, our focus is on extracting predictive text in settings where span annotations are costly to obtain. (Lei et al., 2016) tackle this by introducing two networks, a generator and an encoder, which, respectively, filter for important words before making a prediction. However, theirs is a sampling-based method that must be trained via REINFORCE. Moreover, unlike our approach, they are unable to score individual phrases, limiting interpretability. Our work is perhaps most closely related to (Bastings et al., 2019), which defines candidate spans using a modified Kumaraswamy distribution and then selects spans that are predictive via fused LASSO. Instead, our approach uses an attention mechanism to identify promising start and end positions, which are then used to construct spans nonparametrically. Lastly, another approach is the prediction-constrained topic model, which provides interpretable topics that are useful for pre
dicting labels of interest (Ren et al., 2019; Hughes et al., 2017).
3.1 Predictive Extraction
We define predictive extraction as follows. Given a document  $\mathbf{X}$  and its associated binary label  $y$ , the goal of predictive extraction is to select contiguous sequences of text called spans that, jointly, are sufficient to predict the label  $y$  effectively. One wishes to also assign each span a score reflecting its contribution to the prediction  $\hat{y}$ . In this work, span selection is regularized by quantifying span size and overlap among spans, and performance is evaluated via human rating of randomly selected spans.
3.2 Proposed Model: SpanPredict
The architecture for the proposed SpanPredict model is given in Figure 1. For a given passage of text, let  $t = 1,\dots ,T$  index token  $s_t$ , and let  $\mathbf{e}_t\in \mathbb{R}^D$  denote an embedding of token  $s_t$ . Note that the  $\mathbf{e}_t$  may be linear token embeddings, but may also be contextualized embeddings generated by BERT (Devlin et al., 2018), for example. For each embedding  $\mathbf{e}_t$ , two probability vectors  $\tilde{\mathbf{p}} = \mathrm{softmax}\left(\mathbf{E}^\top \mathbf{w}_p\right)$  and  $\tilde{\mathbf{q}} = \mathrm{softmax}\left(\mathbf{E}^\top \mathbf{w}_q\right)$ , where  $\mathbf{E} = [\mathbf{e}_1,\ldots ,\mathbf{e}_T]$ , are computed using a pair of trainable, sentinel attention vectors  $\mathbf{w}_p,\mathbf{w}_q\in \mathbb{R}^D$ . Vectors  $\tilde{\mathbf{p}} = [\tilde{p}_1,\dots,\tilde{p}_T]\in \Delta^{T - 1}$  and  $\tilde{\mathbf{q}} = [\tilde{q}_1,\dots,\tilde{q}_T]\in \Delta^{T - 1}$ , where  $\Delta^{T - 1}$  is the  $T - 1$  simplex, represent the set of probabilities of each token in the sequence being the start and end of a span of text, respectively. While it is tempting to create a span by choosing the start and end positions with highest probabilities, i.e.,  $\arg \max_{t}\tilde{\mathbf{p}}$  and  $\arg \max_{t}\tilde{\mathbf{q}}$ , respectively, this is problematic since the arg max function is not differentiable, precluding training by standard backpropagation.
To produce a span representation  $\mathbf{r}$  that is amenable to backpropagation, we employ the cumulative sum function  $\mathrm{cumsum}(\mathbf{x}) : \mathbf{x} \in \mathbb{R}^T \mapsto \mathbf{c} \in \mathbb{R}^T$ , where  $c_t = \sum_{t' \leq t} x_{t'}$  is an element of  $\mathbf{c}$ . Using this function, we define  $\mathbf{p} = \mathrm{cumsum}(\tilde{\mathbf{p}})$  and  $\mathbf{q} = \mathrm{cumsum}(\tilde{\mathbf{q}}_{::-1})$ , where  $\mathbf{x}_{::-1}$  is the vector  $\mathbf{x}$  with its elements reversed. Intuitively,  $p_t$  (element of  $\mathbf{p}$ ) represents the probability that the start of a span has occurred by token  $t$  when coming from the left of the sequence and  $q_t$  (element of  $\mathbf{q}$ ) represents the probability that the end has occurred by token  $t$  when coming from the right. We
then calculate a set of weights  $\tilde{\mathbf{r}} = \mathbf{p} \odot \mathbf{q}$ , where  $\odot$  denotes the element-wise product. The product  $\tilde{\mathbf{r}}$  therefore assigns large weights to tokens which have high mass under both  $\mathbf{p}$  and  $\mathbf{q}$ , i.e., those that are identified as falling between the start and end points of a span.
Rather than directly using  $\tilde{\mathbf{r}}$  to compute a span representation, we first normalize  $\tilde{\mathbf{r}} = [\tilde{r}_1,\dots ,\tilde{r}_T]$  such that its elements sum to 1. We define the elements of  $\mathbf{r}$  as  $r_t = \tilde{r}_t / (\sum_t\tilde{r}_t + \epsilon)$  and  $\epsilon \approx 10^{-8}$  is included for numeric stability, since  $\tilde{\mathbf{r}}$  is zero everywhere if the support of  $\mathbf{p}$  and  $\mathbf{q}$  do not overlap, indicating a null span. Importantly, normalization allows us to compute a score that reflects each word's contribution to the span as a whole, regardless of the length of the overall sequence. We then construct a span representation  $\mathbf{m} = \mathbf{Er}\in \mathbb{R}^{D}$ , by taking an average of the embeddings  $\mathbf{E}$  weighted by  $\mathbf{r}$ . This method of constructing spans is a key feature of our model as it allows for span location and length to be dictated nonparametrically, driven only by the content within the identified spans and the quality of the predictions.
We repeat this procedure  $J$  times to identify  $J$  spans  $\mathbf{m}_j$ ,  $j = 1, \dots, J$ , using unique pairs of sentinel vectors  $\{\mathbf{w}_{pj}, \mathbf{w}_{qj}\}$  for each span. Finally, we employ attention over the  $J$  span representations to generate span scores  $z_j = \mathbf{w}_z^\top \mathbf{m}_j$ . These scores are effectively logits, which can be interpreted as the log-odds of a positive label associated with the span. The output of the model,  $\hat{y} = \sigma(\sum_j z_j)$ , where  $\sigma(\cdot)$  is the sigmoid function, is compared against the truth  $y$ , and the model is trained via backpropagation with binary cross-entropy loss.
In this work, we pad or truncate documents, as appropriate, to have fixed length  $\tilde{T}$ . Tokens are mapped to dense vectors using 100-dimensional GloVe embeddings, which are then contextualized with three parallel convolutional layers with filters of kernel sizes  $K\in \{2,3,5\}$  prior to span selection (see Section 5.1 for details). We chose this simple approach over more complex embeddings, e.g., BERT, to focus on the quality of span extraction and its effect on classification performance rather than on maximizing performance per se. However, our approach is agnostic to the choice of embedding, and alternative embeddings may be used if desired.
Figure 1: Model architecture. We begin with tokenization followed by an embedding lookup. Three convolutions with kernel sizes  $K \in \{2,3,5\}$  (shades of blue) are performed in parallel, and outputs are concatenated to form contextual embeddings. The span detection module then identifies  $J = 3$  (in this example) spans denoted by green, yellow, and red. Word scores from the span detectors are used to compute  $J$  weighted average span representations, each denoted by  $\mathbf{m}$ . These are stacked to form  $\mathbf{M}$ . Note that the red span weights are all 0, indicating a null span representation. Finally, we perform attention over the span representations to obtain scores  $z_{j}$ , which are added and passed through a sigmoid to predict  $\hat{y} \in (0,1)$ .
3.3 Constraining span uniqueness and size
Our model already contains an implicit penalty for span size - specifically, the greater the number of tokens over which the model averages to compute a span representations, the smaller the contribution of influential words to the span logits. Hence, the model should implicitly prefer to have spans that are concise and not overwhelmed with "filler" words. Further, our model naturally encourages sparsity of number of spans. Spans that do not carry meaning are biased towards generating weights  $z_{j}$  of zero since, otherwise, they would inadvertently reduce the predictive performance. This also means that the model implicitly learns the number of spans required to make predictions on an individual document basis.
In practice, we observed that spans identified by our model tend to be rather long and suffer from significant overlap, which suggests the need for an additional explicit penalty to make the spans more concise and distinct. Methods involving  $L_{2}$ -regularization on the magnitudes of  $\mathbf{r}_j$  or  $z_j$  may shrink the spans or encourage sparsity, but they do not directly address the overlap issue. Thus, we seek a regularization method that directly compares spans  $\mathbf{r}_j$  with one another.
Since vectors  $\{\mathbf{r}_j\}_{j = 1}^J$  each constitute a discrete probability distribution, a natural choice is to consider divergences between them. Among these, the generalized Jensen-Shannon divergence (JSD) (Lin, 1991), a symmetric measure of similarity among a
set of  $J$  probability distributions, is appealing for several reasons. The JSD is defined as
$$
\underbrace {H \left(\sum_ {j = 1} ^ {J} \pi_ {j} \mathbf {r} _ {j}\right)} _ {\text {s p a n o v e r l a p}} - \underbrace {\sum_ {j = 1} ^ {J} \pi_ {j} H (\mathbf {r} _ {j})} _ {\text {s p a n c o n c i s e n e s s}}, \tag {1}
$$
where  $H(\cdot)$  denotes the entropy and  $\pi = [\pi_1, \dots, \pi_J] \in \Delta^{J-1}$  is a distribution of mixing coefficients among the  $J$  distributions  $\{\mathbf{r}_j\}_{j=1}^J$  (Lin, 1991). While the JSD is commonly expressed as a weighted average of Kullback-Leibler divergences (Manning et al., 1999), in this form, we emphasize that the JSD can be decomposed into two terms: the entropy of the (weighted) average of the  $\mathbf{r}_j$ s and the (weighted) average of the entropies of each  $\mathbf{r}_j$ . Thus, by maximizing the JSD, we simultaneously maximize the entropy of the average distribution (i.e., minimize overlap between the  $\mathbf{r}_j$ s) while minimizing the entropy of each  $\mathbf{r}_j$  (i.e., maximize conciseness of each  $\mathbf{r}_j$ ). In addition, the JSD is bounded below and above by 0 and  $\log(J)$ , respectively, allowing one to monitor convergence during training (see Appendix C) (Lin, 1991).
We can modify the JSD formulation by introducing a tunable parameter  $\theta \in [0, 0.5]$  as follows:
Dataset	Num. Notes	Num. tokens	Age: patients (years)	Age: controls (years)
ASD	44458	560.1 ± 515.3	2.1 ± 2.4	2.2 ± 2.4
ADHD	45160	480.0 ± 437.3	5.6 ± 2.2	5.5 ± 1.6
Asthma	46588	505.7 ± 441.3	1.7 ± 2.0	1.7 ± 0.2
Table 1: Diagnosis prediction statistics. Each dataset was divided into training, validation, and testing subsets with a 45:10:45 split. Positive and negative examples are balanced in each dataset.
$$
\operatorname {J S D} _ {\boldsymbol {\pi}} \left(\mathbf {r} _ {1}, \dots , \mathbf {r} _ {J}; \theta\right) =
$$
$$
2 \left\{\theta \left[ H \left(\sum_ {j = 1} ^ {J} \pi_ {j} \mathbf {r} _ {j}\right) \right] - \right. \tag {2}
$$
$$
(1 - \theta) \left[ \sum_ {j = 1} ^ {J} \pi_ {j} H (\mathbf {r} _ {j}) \right] \Bigg \},
$$
where we recover (1) when  $\theta = 0.5$ . As we slide  $\theta$  closer to 0, the contribution of the second term increases; hence, the smaller the value of  $\theta$ , the smaller we can expect the entropies of the individual distributions to be. This implies that the span sizes can be made smaller by reducing  $\theta$ .
Lemma 3.1. The modified JSD is bounded above by a constant, independent of the entropies of the individual  $\{\mathbf{r}_j\}_{j=1}^J$ .
Proof. Defining  $H_{1} = H\left(\sum_{j=1}^{J}\pi_{j}\mathbf{r}_{j}\right)$  and  $H_{2} = \sum_{j=1}^{J}\pi_{j}H(\mathbf{r}_{j})$ , we have:
$$
\begin{array}{l} \operatorname {J S D} _ {\pi} \left(\mathbf {r} _ {1}, \dots , \mathbf {r} _ {J}; \theta\right) = \\ = 2 \left\{\theta H _ {1} - (1 - \theta) H _ {2} \right\} \\ = 2 \left\{\theta H _ {1} - \theta H _ {2} \right\} - 2 (1 - 2 \theta) H _ {2} \tag {3} \\ = 2 \theta \mathrm {J S D} _ {\pi} (\mathbf {r} _ {1}, \dots , \mathbf {r} _ {J}) - \\ 2 (1 - 2 \theta) H _ {2} \\ \leq 2 \theta \mathrm {J S D} _ {\pi} (\mathbf {r} _ {1}, \dots , \mathbf {r} _ {J}), \\ \end{array}
$$
where the last line follows from the fact that  $1 - 2\theta \geq 0 \forall \theta \in [0, 0.5]$  and  $H_2 \geq 0$ .
This result provides a lower bound on our JSD objective, useful for monitoring convergence during training, i.e.,  $-\mathrm{JSD}_{\pi}(\mathbf{r}_1,\dots ,\mathbf{r}_J;\theta)\geq -2\theta \log (J)$
4 Learning
The complete objective function we aim to minimize is thus given by:
$$
\begin{array}{l} \mathcal {L} = - \mathbb {E} _ {\mathcal {D}} \left[ (1 - \alpha) \left(y \log \hat {y} + \right. \right. \\ (1 - y) \log (1 - \hat {y})) + \tag {4} \\ \left. \alpha \operatorname {J S D} \left(\mathbf {r} _ {1}, \mathbf {r} _ {2}, \dots , \mathbf {r} _ {J}; \theta\right) \right] \\ \end{array}
$$
where  $\mathcal{D}$  is our dataset, and  $\alpha \in [0,1)$  is a hyperparameter denoting the weight of the modified JSD penalty relative to the classification loss. For simplicity, we choose to take  $\pi_{j} = 1 / J$  in (2) and have therefore omitted  $\pi$  from the expression for  $\mathrm{JSD}_{\pi}(\mathbf{r}_1,\dots ,\mathbf{r}_J;\theta)$  in (4).
Aside from the learning rate, our model consists of only three hyperparameters  $J$ ,  $\theta$ , and  $\alpha$ , making it highly attractive for experimentation. Predictive performance is not very sensitive to the choice of  $J$ ; here we select  $J$  to be proportional to the average document length in each dataset, but we investigate the impact of a fixed larger value of  $J$  in Appendix B. To choose  $\alpha$ , we employ a method similar to that used in (Smith, 2017) for choosing a learning rate. Specifically, we slowly ramp up  $\alpha$  from a minimum value of 0 in increments of  $10^{-5}$  batch by batch and monitor validation accuracy. When the accuracy starts to level off or drop, we mark the value of  $\alpha$ ; we found  $\alpha = 0.1$  to be appropriate for our datasets. Parameter  $\theta$  is selected via cross-validation (trading off performance for desired span length), and is a focus of our experiments, described below.
Datasets We perform experiments on five datasets: two publicly available non-medical datasets, and three constructed from clinical notes from the Duke University Health System. We consider the IMDb movie reviews dataset<sup>1</sup> (Maas et al., 2011), which contains 25,000 training and testing
1https://www.tensorflow.org/datasets/catalog/imdbreviews
(a) AUC relative to baseline
(b) Average span size
Figure 2: Trends in performance. Baseline AUCs - IMDb: 0.938, Amazon: 0.931, ASD: 0.702, ADHD: 0.804, Asthma: 0.630. Performance tends to drop slightly as  $\theta$  is decreased, but spans become more concise and distinct.
(c) Average intersection-over-union
examples of movie reviews and a binary viewer rating; and the Amazon Fine Food Reviews dataset $^2$  (McAuley and Leskovec, 2013), which contains  $>500,000$  reviews of food items, which we subsample to 25,000 training and testing examples and 5000 validation examples for consistency. Positive and negative examples are balanced in each subset. Reviews are on a 5-point scale, but we binarize by labeling ratings of 3 or higher as positive. Average document length for IMDb is  $225.4 \pm 166.1$  tokens, and shorter for Amazon at  $84.3 \pm 86.1$  tokens.