Value-aware transformers for 1.5d data

1 INTRODUCTION . Modelling the sequence of events in electronic health records ( EHRs ) poses fundamental challenges to contemporary sequence models owing to the complex , heterogeneous nature of the data . Even choosing a good input representation for such data presents difficult trade-offs . The variety of data shapes encountered within the records for different patients and across different sections of their visit is substantial . While there are some variables , such as patient demographics , that one might reasonably expect all patients to have an entry for , the heterogeneity of structure sets in as soon as the patient ’ s individual health status is more intricately explored . During an admission to hospital , a series of laboratory tests ( hereafter referred to as ‘ labs ’ for succinctness ) may be conducted on a patient : the number of possible types of test might be large ( in the hundreds ) , these may be measured repeatedly , once , or not at all . Hospital stays may vary significantly in length and this variation is not necessarily indicative of severity . Very ill patients may have short survival , or may be artificially sustained on life support in expectation of a recovery that fails to arrive . There are further complications : some tests may have been conducted in advance - attributed to their historical health record ; or recorded in a different modality of their health record ( e.g . notes , charts , etc. ) . The vast complexity and heterogeneity of EHR data presents representation learning with fertile ground for innovation . 1.1 OUR CONTRIBUTIONS . This paper seeks to address a particular data regime which is under-studied : that of sparse sequential and highly-multivariate data , where the total number of variables typically outstrips the total number of measurements for a given data instance . Table 6 illustrates this for the MIMIC-III lab-events dataset , where the mean number of completely unmeasured variables for a given visit is 697 out of a possible 753 . As we shall describe , conventional approaches to modelling this data make severe simplifying assumptions , or rely heavily on imputation . Many contemporary works restrict both the variables selected and the population studied , so as to reduce missingness . To address these issues , we : • Present a data encoding that effectively captures sparse sequential highly-multivariate data while avoiding considerable imputation . • Propose neat adaptations to make a standard transformer architecture value-aware , equipping it to learn and generate this data encoding . • Provide an empirical study on a real-world dataset demonstrating that a pre-trained valueaware transformer performs competitively on challenging downstream predictive tasks . 1.2 RELATED WORK . This work sits on the intersection of a number of exciting areas of research : representation learning for healthcare , transformer architectures , sparse sequence modelling . Here we try to do justice to some of the many relevant works . Representation learning for EHRs . Researchers seeking to make secondary use of the signal within EHRs have paid particular attention to the developments in deep learning ( Shickel et al. , 2017 ; Si et al. , 2020 ) , with applications ranging from unsupervised patient stratification Landi et al . ( 2020 ) to learning personalised comorbidity networks ( Qian et al. , 2020 ) . The wave of excitement set in motion by transformer architectures ( Vaswani et al. , 2017 ) quickly propagated to the healthcare world . Large language models such as BERT ( Devlin et al. , 2018 ) have been subjected to rounds of secondary pre-training on medical corpora ( Alsentzer et al. , 2019 ) , with the aim of harnessing signal locked up in the idiosyncratic natural language formats such as clinical notes . By encoding events as tokens , transformers can also be utilised for studying EHR event sequences . When arranged to constitute a sentence , these tokens capture the pattern of a patient ’ s stay . This approach can be pursued at various levels of granularity , for instance Li et al . ( 2020 ) use BERT to model sequences of diagnoses over multiple visits . Applying the same methodology at a finer grain presents a natural model for within-visit medical event sequences such as labs . But while knowledge of the performance of a lab test conveys information , it is the lab test together with its value that is the more complete indicator of the health status of a patient . Choi et al . ( 2020 ) made progress on integrating labs and their values using a graph convolutional transformer , modelling the graphical and hierarchical structure of EHRs in an exciting piece of work . However , their method was not designed to model repeat measurements of labs - a common phenomenon within the natural course of medical investigation and treatment , motivating our work . Healthcare benchmarks . In order to objectively compare systems such as those described above , EHR-specific benchmark tasks have been proposed , many based on the MIMIC dataset Johnson et al . ( 2016 ) . Of particular relevance here are the prediction tasks for in-hospital mortality and length-ofstay proposed in Harutyunyan et al . ( 2019 ) ; Purushotham et al . ( 2018 ) . In related work , and with a view to further increasing reproducibility , Wang et al . ( 2020 ) provided a pre-processing pipeline that addresses concerns such as unit conversion , outlier filtering and aggregation that are common in these tasks . The considerable heterogeneity of structure in EHRs means that the benchmarking papers need to make many design and preprocessing choices for the cohorts selected , the variables used , and the definitions of prediction targets . Not all of these choices are desirable for applications : e.g . in Harutyunyan et al . ( 2019 ) labels are defined in such a way that the events of interest could occur within the data visible to the model rather than being restricted to predicting future events , while in Purushotham et al . ( 2018 ) their largest feature set includes only 11 of the approx . 700 raw features available in the MIMIC-III lab tests , chosen due to their minimal missingness rate . Some of these considerations ( targets labels ) we would like to refine so as to provide more utility , others ( forced variable selection and aggregation ) do not align with the purpose of this work , which is to remove the need for cherry-picking of variables . This makes direct comparison with these works challenging - a point which other authors have highlighted , even when using largely similar cohorts , variables and identical prediction targets ( Che et al. , 2018 ) . Irregular times series . A closely related problem is that of studying irregularly-sampled multivariate time-series . Shukla & Marlin ( 2020 ) detailed three different types of input representations : series-based , vector-based and set-based , capturing time , variable and value for a range of scenarios . Horn et al . ( 2020 ) focus on the set-based approach considering each times series to be a bag-ofevents , proposing a model which applies a set function together with a single attention mechanism between elements of the set , topped with a classifier . In their experiments , an out-of-the-box transformer is state-of-the-art for the Physionet 2012 prediction task ( Goldberger et al. , 2000 ) , and performs competitively on a MIMIC-III mortality prediction task ( based on Harutyunyan et al . ( 2019 ) ) . Another approach is elicited in ( Che et al. , 2018 ) , where they introduce a GRU-based architecture , designed with inbuilt decay mechanisms . Their method is attuned to time-series variables such as flow data or regularly-sampled labs and relies on the imputation of missing values , which , for our data regime would constitute on average 98 % of the input ( see Table 6 ) . 2 DATA GREATER THAN 1D . Sequences of words such as those found in natural language are neatly represented as 1D data using tokens . Sequences of lab results such as those found early in EHRs are neatly represented as ( token , value ) pairs , which might suggest that the ideal representation of this data should a multivariate series with one dimension per lab type . To model it in this form would require considerable imputation due to the data ’ s sparsity . However , the presence or missingness of particular labs in a sequence are informative : they are requested by healthcare professionals guided by clinical intuition that the EHR may not capture . By imputing the data , this valuable signal is weakened . Perhaps we can make do with just a 1D form - and if so , how should that 1D be constituted ? These are the key questions that a practitioner hoping to model in our data regime faces . As we will discuss in this section , when the data is sparse and highly-multivariate , the choice is not clear a priori . For the motivational problem of lab-result sequences there may be only few or a large number of measurements with multiple repeat samples of the same lab - characteristics which may contain important information . We close the section by proposing a variation on the conventional approaches , which we refer to as “ 1.5−dimensional ” data : a form preserving the sequential nature of the lab orders , together with information about the values , but without the need for full multivariate or sparse vector-based sequence representations and the imputation issues that would ensue . 2.1 CHARACTERISTICS OF LAB-RESULT SEQUENCES . In-patient lab-result sequences are generated by a process of hospital investigation and treatment which imbue them with particular characteristics . Relevance : the labs are pertinent to the health investigation being conducted on the patient . Only a small proportion of the all of the labs available might be obtained during any given visit . Repeats : these relevant measurements may be taken once , or repeated many times . Pragmatism : difficult or inconvenient labs will not be taken unless absolutely necessarily . This leads to a rare-token problem also commonly encountered in natural language . If we believe that these characteristics deliver valuable information , then the way that we represent our data should not compromise them . 2.2 COLLAPSING DATA DOWN TO 1D . A common way to deal with this difficult data dimensionality is to ‘ collapse ’ it to 1D . There are two natural ways of doing this : the horizontal collapse and the vertical collapse , illustrated in Figure 1 . Horizontal collapse . The horizontal collapse retains the sequential structure of the data , and removes the value aspect . One form of the horizontal collapse is to ignore the value from the tokenvalue pair . We then have sequences of lab orders , and the premise is that there is sufficient information contained in the presence of a lab order to inform the underlying state of a patient . This capi- talises on recent advancements in NLP and has formed the basis for some of the medical sequencemodelling applications mentioned in Section 1.2 . Vertical collapse . An alternative form of 1D collapse is to collate repeated values of a variable using a summary measure ( such as the mean , median , max , most-recent ) and then use these summary measures – one per variable – as the input to a model . This representation respects the unique value scales on a per-variable basis , however it compromises the sequential nature of the data and collapses repeated measurements . 2.3 IMPUTING FOR MULTIVARIATE SEQUENCES . An enticing candidate for the representation of ( token , value ) pairs is as a multivariate sequence with one sequence dimension per token . This retains each variable ’ s independent value-scale and the innate temporally-ordered nature of the data . However , the sparsity of our specific data regime is problematic , affecting its suitability : a program of severe imputation is required in order to use this input form . This commonly takes the form of sequential forward- or backward-filling of data ; per-variable mean imputation ; or the use of sentinel values , the choice of which will necessarily depend on the variable in question ( so as not to inappropriately assign a sentinel value that is commonly attained by true data points ) . Imputation is not necessarily evil , but does require careful application , even in a relatively dense data regime . Not all missingness is made equal and even the most advanced general-purpose imputation schemes rely on assumptions ( MCAR or MAR : see Mohan & Pearl ( 2021 ) ) which are often unrealistic in a healthcare setting . Unrealistic imputation can have unexpected downstream implications ( Caruana , 2021 ) which can be dangerous for critical applications . Even a sparse vector-based multivariate representation such as that detailed in Shukla & Marlin ( 2020 ) would require on average 98 % of its input to be imputed for our data ( see Table 6 ) .

Summary: The paper proposes an approach for learning representations of spare sequential data, commonly found in healthcare settings. The work challenges issues with sequential data in electronic health records such as sparsity of events, unequal lengths, which also present imputation challenges. A transformer-based approach is presented which incorporates information about the different sequential events (labs) along with their values which is termed as a 1.5-dimensional representation. The transformer is trained to obtain the representation of sequences of events, labs conducted for in-hospital patients, along with the values of the labs which are categorized based on the quantile range obtained from the training data. The representations from the transformer embeddings are used for predicting two outcomes, 1) in-hospital mortality, and 2) length of stay at 48 hours after admission. The value- aware transformer performs marginally better than a value-unaware transformer with respect to $\mathcal{L_{\text{\[TOK\]}}}$. There is also an improvement in the balanced-accuracy and ROC-AUC for in-hospital mortality when compared with feed-froward neural network and value-unaware transformer. Similar performance characteristics are also observed for length-of-stay regression.