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: This work proposes an approach for modeling sparse and multivariate sequential data which are commonly found in the healthcare domain. The paper presents and compares different data representations and introduces a 1.5d representation and outlines principles for modeling this type of data. The paper presents the value-aware transformer, a transformer decoder-only architecture for the 1.5d data representation. Based on the experiments on the MIMIC-III dataset, the proposed model and the data representation achieve better performance than several other standard data representations.