End-to-end clinical temporal information extraction with multi-head attention (2023)

Abstract: Understanding temporal relationships in text from electronic health records can be valuable for many important downstream clinical applications. Since Clinical TempEval 2017, there has been little work on end-to-end systems for temporal relation extraction, with most work focused on the setting where gold standard events and time expressions are given. In this work, we make use of a novel multi-headed attention mechanism on top of a pre-trained transformer encoder to allow the learning process to attend to multiple aspects of the contextualized embeddings. Our system achieves state of the art results on the THYME corpus by a wide margin, in both the in-domain and cross-domain settings.

## introduction temporal information extraction is the task of discovering event mentions, time expressions, and relations between them that indicate their relative temporal ordering. the goal of these tasks is to place events on a timeline with as specific granularity as possible, to better understand topics that have strong temporal components. in the clinical setting, temporal information extraction can be applied to the text in electronic health records to create a timeline for a patient, enabling downstream applications that are heavily time-dependent (e.g., modeling the temporal ordering of problems and medication use to extract causally plausible candidates for adverse drug events). recent datasets have been released as part of community shared tasks that enable the research community to make progress on this challenging suite of tasks @xcite @xcite @xcite . the thyme (1 and 2) and i2b2 datasets have led to progress in this area but these problems are far from solved. both the clinical tempevals @xcite @xcite and the 2012 i2b2 challenge @xcite evaluated systems on an end-to-end basis -that is, given only raw text, participants needed to both extract relation arguments and relations between them. however, most systems attacked this setting with a pipeline approach, combining the best event and time extraction systems with a relation system. in recent general domain work, models for doing end-to-end relation extraction have been obtaining some success recently by building on pre-trained transformers in the bert family @xcite , using additional task-specific components that aggregate encoder information into relation predictions @xcite . in this work, we extend liu et al. (2020)'s model for the task of clinical temporal information extraction. instead of using multi-headed attention for different relation categories, we use multiple heads to capture different aspects of the token representations. the additional layer allows the model to then integrate the signals to make the final relation predictions. our results show that this architecture can obtain state of the art performance for relation extraction on the clinical tempeval 2016 and 2017 tasks, which evaluated in-domain (colon cancer) and cross-domain (colon @xmath0, respectively. our analysis shows that, for this task, multi-headed attention is far superior to single-headed attention with the same number of model parameters, suggesting that the multi-headed approach is valuable for representing different aspects of the input. in addition, we compare the trade-offs in modeling this as a multi-task learning setup, as opposed to using multiple fine-tuned classifiers that work separately, where the time expressions and events are found separately and used as the anchors for the relation classifier. finally, we report results of this model on the thyme2 data, for which no current baseline results exist in the end-to-end task. ## background and related work this work takes advantage of the thyme (temporal history of your medical events) corpus @xcite , and specifically the data released as part of the clinical tempeval shared tasks hosted at semeval @xcite @xcite . this english-language dataset uses colorectal and brain cancer notes from patients at the mayo clinic. each patient is represented with 3-4 notes written around the time of cancer diagnosis, typically with two clinical notes and one additional radiology and/or pathology note. the first two clinical tempeval's evaluated on colon cancer only. the third clinical tempeval was a domain adaptation task where participants were given the colon cancer data and evaluated on the brain cancer test set. there were two settings, one with unsupervised domain adaptation (no access to brain training data), and another with a small amount of labeled brain cancer data. in practice, there was little unsupervised adaptation attempted, with one top system obtaining negative results with the technique of freezing embedding weights before training (to prevent some of the embeddings from updating and drifting away from target-domain-specific terms that might not be seen during training) @xcite , and another also obtaining negative results by replacing rare terms in the training and test data with an unk token @xcite . in this work we use the thyme colon cancer data for training, the development set for tuning and model selection, and evaluate on both the colon cancer test set and brain cancer test set, in the unsupervised setting. more recent work on thyme has focused on the easier relation classification setting, where gold standard time expressions and events are provided, and the system task is to decide which pairs of temporal entities should be linked with temporal relations. the current best-performing system is from @xcite . that paper's main contribution was a continued pre-training technique for clinical text, where the masked language modeling task was modified to preferentially mask tokens that were part of temporal entities. the temporal relation extraction was modeled as a sentence classification task, where the gold standard arguments were marked by special tokens in the input, and a softmax layer off of the pre-trained transformer's sentence representation token was used to predict the relation between the two marked tokens (including the none relation). such an approach cannot be fairly compared (or even deployed) without building additional systems to do the event and time expression detection, so we do not directly compare to that result. other work in end-to-end relation extraction is also relevant to this work. the most relevant related work @xcite used a bilstm encoder to get token representations, which fed into a multiheaded attention layer, where each attention head was used to score a different relation type. that work differs from ours in that it is applied only to "general-domain" tasks including the nyt @xcite and webnlg datasets @xcite , and treats outputs values from different attention heads as prediction scores for individual relation categories. our work, in contrast, uses attention scores as features for a single downstream relation classifier (the contains relation), allowing the model to use different aspects of the representation to make classification decisions. the current work also has connections to techniques for end-to-end relation extraction where the task was treated as a table filling task @xcite , where each cell (i, j) in the table represents the relation between token i and token j. ## methods our model uses a multi-task learning setup built on top of a pre-trained transformer encoder, and specifically the pubmedbert-base-uncasedabstract model @xcite . time expressions and events are modeled as two independent begin-inside-outside (bio) tagging tasks, so that each contextualized token embedding in the sequence is used to make a classification decision for that token's categorization as an event or time expression. this bio tagging is implemented with softmax layer at every token index over the b-or i-version of each of the time expression or event types. relation classification is treated as a binary task applied to n 2 token pairs to find relations, using a mechanism very similar to the self-attention mechanism already used extensively in transformers. our attention-inspired classification model is an additional layer on top of the pre-trained transformer that takes in the contextualized token representations from the last layer of the transformer encoder. attention first came to wide use in encoderdecoder architectures (e.g. @xcite ), as a way for a decoder to pick out relevant sub-components of the input (e.g., individual token representations for indeterminately long input sentences) during generation of outputs. typically, it is implemented as a fullyconnected neural layer that learns a score for each sub-component @xcite and uses the computed weighted average over the sub-components at each decoding step. self-attention is a simple extension where the attention mechanism is used to update weights over a discrete set of elements at a higher layer l by applying attention over the same elements at a lower layer l -1. @xcite , in introducing the transformer architecture, showed that attention alone could be used both in encoding and decoding to learn hidden representations. further, they introduced the idea of multi-headed attention, that is, the idea that several different attention distributions can be computed, as a way to focus on and combine multiple different aspects of the input representations. the first stage in attention is projecting the input with query, key, and value matrices, and then the attention is computed as: multi-headed attention multiplies this out with distinct q, k, and v projections for h attention heads (indexed by h), and then concatenates the outputs, with the sizes of the projection matrices scaled such that the total output dimensionality matches the input. for self-attention, the inputs to the q, k, and v projections are all the same -the output from the previous layer. the most important aspect of attention for our purposes is that the matrix product s @xmath1. in the work by @xcite , each relation type is aligned to one attention head, so the value at each position in the matrix represents the affinity score for the relation aligned to head h: in this work, we use h attention matrices to compute a single h × n × n tensor, where a softmax along the attention head dimension leads to a classification for each token pair into a single relation, the temporal narrative container relation: the overall loss function combines time expression extraction, event extraction, and relation classification, all of which use versions of cross-entropy loss. the overall loss function simply sums the three, with a weight on the relation loss that we tune on the development set to maximize f1 score for the relation classification task. ## evaluation we evaluate on the thyme colon and brain corpora, tuning hyperparameters on the colon dev set and reporting results of the narrative container extraction task on the colon and brain test sets. 1 hyperparameters we tuned over included fine tuning learning rate, batch size, classifier layer (which layer the classification head is connected to), relation task weight, number of classifier attention heads, and dimensionality of attention heads. for all reported results we use the official semeval scoring tool, which performs transitive closure for the relations (uzzaman and allen, 2011). while the system we describe is multitask in nature, it does not have explicit interactions between the sequence tagging models and the relation predictor. we first find predicted relations between tokens and work backwards to find the argument span that exists at each token to make a relation prediction. this independence allows us to investigate the importance of the quality of the extraction of those entities, by using alternative methods, including an oracle, in the second step of finding the arguments for predicted relations. the row labeled + gold args shows the performance with oracle-based arguments taken from the gold standard, to show us an upper bound on how well relation extraction would work with perfect argument extraction. the row labeled -mtl shows the performance if we use separately tuned event and time expression extraction systems, since multi-task learning sometimes can cause degraded performance of independent components. the e2e-mha system obtains the best performance on both corpora. for the colon cancer corpus, precision is 14 points higher than the best performing clinical tempeval system, while recall is within one point, giving a 5-point increase in f1. for the brain cancer corpus, precision decreases by 11 but the increase in recall results in a small 2-point increase in f1. the next result we report is on the importance of multi-headed attention. our hypothesis is that using multiple attention heads allows the model to more easily learn to focus on multiple facets of an instance. an alternative hypothesis is that the gain comes from using a lot more parameters than an architecture that uses one attention head per relation (e.g., @xcite ). to test this, we can compare against a similar model that has fewer attention heads but more feature dimensions per head, holding the total number of model parameters constant. for this experiment, we report colon development set validation scores directly as measured by scikit-learn (i.e., not using clinical tempeval scorer, so it does not include temporal closure), as we tuned this parameter before running final tests. finally, we report the results when the same system is tuned, trained, and applied to the thyme2 corpus @xcite ## discussion and conclusion while this task remains challenging, with the highest obtained f1 scores 0.53 (in-domain) and 0.36 (out-of-domain), the combination of biomedically pre-trained transformers and a novel multi-headed attention classifier obtains results that improve the state of the art. even in-domain, the task is fairly far from being solved. additional future work should also explore explicit domain adaptation, as out-of-domain perfor-mance still suffers despite obtaining better performance than previous methods. simple unsupervised methods like target domain continued pretraining @xcite are likely to provide benefits on the brain cancer corpus, and future work should investigate this and other approaches. we did see performance improvements in preliminary work when switching from general domain pre-trained models (roberta-base @xcite ) to a biomedically trained model, so it is possible that some of those potential gains are already incorporated by the more general language model. the trained models developed for this work are available in the huggingface models repository. 2 the code for implementing and evaluating these models is available on github. 3 the code repository also contains a fastapi-based demonstration server that sets up a rest server to convert text arguments into temporal relations in a simple json schema.