S pan P redict: Extraction of Predictive Document Spans with Neural Attention (2021)

Abstract: 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.

## introduction 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 embeddingstend 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 multipleword descriptions, not individual terms. morever, 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 @xcite ; and in many question answering tasks, including in the stanford question answering dataset @xcite , 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 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 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 @xcite or gradient estimation alternatives @xcite . 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: ## related work explaining neural network predictions is a wellknown problem, one that is particularly challenging in natural language processing, due to the presence of complex semantic structure and interdependencies @xcite . 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 @xcite . 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 @xcite ; 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 @xcite , for example, or candidate spans can be limited to individual words and sentences @xcite . 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 @xcite . 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 @xcite , and using subsequent notes, which are more likely to repeat earlier information if it is important @xcite . in contrast to these methods, our focus is on extracting predictive text in settings where span annotations are costly to obtain. @xcite 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 @xcite , 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 @xcite . ## learning the complete objective function we aim to minimize is thus given by: aside from the learning rate, our model consists of only three hyperparameters j, θ, and α, 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 α, we employ a method similar to that used in @xcite for choosing a learning rate. specifically, we slowly ramp up α 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 α; we found α = 0.1 to be appropriate for our datasets. parameter θ is selected via cross-validation (trading off performance for desired span length), and is a focus of our experiments, described below. 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± 166.1 tokens, and shorter for amazon at 84.3 ± 86.1 tokens. the three medical datasets were built by sampling the clinical progress notes of children visiting the duke university health system between october 1, 2013 and october 1, 2018. all analyses were approved by the duke university institutional review board. diagnosis codes (icd-9/10) were used to identify patients eventually diagnosed with autism spectrum disorder (asd), attention deficit hyperactivity disorder (adhd), or asthma. notes from each patient group were then selected at random and labeled as positive for the condition corresponding to that group. while many of these notes are not directly related to the condition of interest, a large proportion contain related information or risk factors. future work will focus on extracting predictive spans from all notes from a given patient; here we focus on individual notes to limit complexity and highlight span extraction performance. for each diagnosis prediction task, we then selected notes from age-matched controls not diagnosed with the condition as of october 1, 2018, and assigned them a negative label. each dataset contains an even number of positive and 2 https://www.kaggle.com/snap/ amazon-fine-food-reviews negative examples. descriptive statistics are shown in ## results and discussion in figure 2 , we describe trends in performance. baseline aucs are provided in the caption. note that lower aucs for diagnosis prediction reflect the comparative difficulty of these tasks. figure 2a shows performance relative to the baseline model for varying jsd penalties. performance decreases up to 6% as the penalty increases, with the exception of asd, on which the model performs about as well as or better than baseline for θ ∈ [0.4, 0.475]. thus, while some information may be lost during summarization, depending on the dataset, summarization may also serve to denoise the text, improving predictive performance. from figure 2b , we find that as the penalty is increased, spans become considerably shorter. inspecting the results when θ = 0.25, we found that the model tends to focus in on key words rather than phrases. from figure 2c , we see that overlap also shrinks with span size. the effect is more rapid for the medical datasets, likely because the non-medical passages contain text throughout that is relevant to the sentiment of the passage, whereas medical notes contain information not relevant to the prediction task. a notable exception is asthma, colors represent the different spans (@xmath0, @xmath1. solid lines denote r j (heights of r 1 and r 3 rescaled for visualization purposes). dashed lines denote pj and dotted lines qj , with shading to resolve overlap. the inset plot shows the scores z j , which are added to predict the log-odds of a positive label. which maintains a relatively constant span size and overlap, suggesting that diagnosing asthma requires identifying specific phrases (e.g., "shortness of breath") that cannot be decomposed into individual words. finally, we demonstrate in appendix b that, for @xmath2, auc is, on average, greater but at the cost of greater sensitivity to θ. figure 3 provides an illustration of individual spans inferred by spanpredict (θ = 0.5). in the imdb example (top), we see that the model captures two highly positive spans, each constituting 30-35% of the note, with words such as "professional," "laughed," and "appreciate" appearing in the red span. spanpredict is also able to capture meanings of complex positive phrases, such as "chock full", "sure handed," "none of the over the top," and "time has come." the blue and green spans each cover only a single word; however, these words -"flawless" and "beautifully" -have significant positive connotation. this is a feature our model shares with @xcite , which also picks out individual tokens. the amazon review (figure 3 , bottom) contains mixed sentiment. the green span contains the word "quality," which, akin to words such as "care" or "workmanship," is slightly positive. however, the blue span is filled with negative phrases. this is reflected in the z j scores in the inset plot, which are added to predict the log-odds of a positive label. we find that z j is negative for the blue span while positive for the green span. the orange span is most negative, suggesting that the model is able to synthesize information from the blue and green spans it overlaps to extract an overall meaning. figure 4 shows the human evaluation results. for each span, we computed the median rating among the 3 reviewers and performed a non-parametric anova (kruskal-wallis test) to assess agreement with model-predicted scores. statistically significant differences in means (@xmath3.001) were present in the imdb, amazon, and adhd datasets, but not the asd and asthma datasets. given our model's high agreement with human raters in the imdb and amazon tasks, the lower agreement observed on the medical diagnosis tasks may indicate that our model is identifying descriptive risk factors not familiar to our clinical raters. this hypothesis, which was suggested by our clinical collaborators, will be explored further in subsequent work. to measure inter-rater reliability, we computed cohen's kappa for each dataset -imdb: 0.73, amazon: 0.78, asd: 0.53, adhd: 0.64, asthma: 0.63. these values illustrate the difficulty of evaluating the clinical notes compared to the review datasets. ## conclusions we have introduced the task of predictive extraction, in which document labels are predicted from extracted contiguous segments of text called spans. we presented spanpredict, which constructs span representations nonparametrically from contextualized embeddings by predicting start and end positions using linear attention. our model is straightforward to tune, and assigns interpretable span scores that are added together to predict the logodds of a positive label. model performance and span quality are evaluated on two non-medical and three medical datasets. notably, we observe high correlation between human span ratings and model-predicted span scores, particularly in the non-medical datasets, illustrating that our model selects meaningful spans and scores them accurately. discrepancies between human ratings and model predictions in the medical datasets may suggest that our model is identifying condition-specific risk factors that are unfamiliar to trained clinicians. future work will consider prediction and span extraction from a collection of documents rather than individual documents, allowing descriptive risk factors to be extracted from patient medical histories. clinical findings consistently highlighted by span-predict will be analyzed as possible risk factors via standard statistical methods. additionally, whereas spanpredict identifies a set of spans sufficient to predict the label, future work will explore methods for ensuring that all predictive spans are identified.