DeepJoint: Robust Survival Modelling Under Clinical Presence Shift
Abstract
Medical data arise from the complex interaction between patients and healthcare systems.
This data-generating process often constitutes an informative process.
Prediction models often ignore this process or only partially leverage it, potentially hampering performance and transportability when this interaction evolves.
This work explores how current models may suffer from shifts in this clinical presence process and proposes a multi-task recurrent neural network to tackle this issue.
The proposed joint modelling competes with state-of-the-art predictive models on a real-world prediction task.
More importantly, the approach appears more robust to change in the clinical presence setting.
This analysis emphasises the importance of modelling clinical presence to improve performance and transportability.
Introduction
Medical observations reflect the complex interaction between patients, clinical staff, and healthcare services.
In the context of a hospital stay, each observation from admission to discharge emerges from one of these interactions: observations’ times and types reflect the patient’s condition and the clinical staff ’s assessment of what was needed and when.
The data-generating process is, therefore, informative — we refer to this process as clinical presence .
While the machine learning community has leveraged individual aspects of clinical presence for improved performance, there are few predictive approaches that explicitly model this process.
More commonly, modelling pipelines consist of imputation and re-sampling that assume non-informativeness.
Although practitioners justify ignoring clinical presence to avoid leveraging a changeable process [1] in the hope of improved robustness, no evidence has demonstrated this property.
The challenge of robustness is at the core of clinical presence as changes in this process can occur due to new medical findings or insurance policy incentives.
Despite the importance of models’ transportability under clinical presence changes, no work has explored how different strategies for handling clinical presence impact this robustness
Our work studies (i) how state-of-the-art methodologies behave under clinical presence shifts, and (ii) how to leverage clinical presence for improved transportability.
Specifically, we propose a multitask architecture explicitly modelling clinical presence in addition to the survival outcome of interest as shown in Figure 1.
Three components aim to detect potential shifts in the clinical presence process while regularising the survival model to obtain a more robust embedding.
Related work
Clinical presence affects Electronic Health Records (EHRs) in multiple ways.
First, observation times potentially provide information on the clinical assessment of a patient’s condition.
For instance, the testing frequency decreases when clinical staff believe that the patient’s condition is improving.
Second, the sequence reflects the clinical staff’s understanding of the disease, e.g. if a laboratory test is inconclusive, a second would follow.
Third, missing values reflect irrelevance for diagnosis, or non-informativeness [2].
Leveraging these dimensions have shown predictive edge [3, 4, 5], and reduced biases [6].
For instance, [3, 7] show improved mortality prediction when including missingness indicators.
Methodologies to handle clinical presence [8] can be divided into four groups: (i) Assumed uninformative.
Assume the observations representative of the patient’s health status.
(ii) Leveraged for pre-processing.
Leverage the sampling process for improved imputation and interpolation [9, 10].
(iii) Featurised.
Use measures of clinical presence indicative of medical expertise as model’s inputs, e.g. missing indicators [3], observation times [7, 11, 12], inter-observation time [13, 14, 15, 16, 17], or frequency [18].
(iv) Jointly modelled.
Joint models [6, 19, 20], marked point process [21] and Markov models [22] and their machine learning extensions [11, 23, 24, 25, 26] model some aspects of the visit process and the outcome of interest.
Through shared effects, state modelling or decaying embedding, the data-generating process informs the modelling of the outcome of interest.
Closer to our work, this last category of methodologies either relies on strong parametric assumptions, or does not scale to large datasets.
For instance, [27] model irregular time series with a latent ordinary differential equation and incorporate the temporal process through a parametric
Poisson process.
Our proposed methodology aims to tackle this issue by non-parametric modelling both the clinical presence process and the outcome of interest.
A fundamental feature of clinical presence is its propensity to change as medical practices evolve or due to hospital-level, regional or national practice differences.
This is a crucial limitation of the application of machine learning for medical data as models may not generalise and may become outdated as practices and policies evolve [1].
Evidence suggests performance drops across regions and over time [28, 29, 30, 31].
Developing approaches to detect and potentially improve robustness to clinical presence changes is essential to machine learning applicability in clinical practice.
However, no previous work has explored how handling clinical presence may impact robustness. Our work aims to fill this gap.  Deep Joint 3.1 Setting formalisation The studied medical setting consists of a population of N individuals with a vector of laboratory tests lab i,j – with missing values – observed through time, in which the notation j ∈ [[0 , l i ]] denotes the j th observation for patient i at time t i,j ∈ R + ,
where t i, 0 = 0 is the patient’s admission to the
Intensive Care Unit (ICU) and l i is the last observation for patient i . To study clinical presence, we introduce two additional notations. First, we model the elapsed time between two observations: ϵ i,j = t i,j − t i,j − 1 , note the subscript i underlines that observations time may not align across patients. Second, the data-generating process is also characterised by which laboratory tests practitioners perform.
Consequently, we introduce the mask
m i,j = [ 1 Lab k observed at t i,j ] k ∈ labs indicating the observed laboratory tests at observation time t i,j .
This work models the survival outcome after the first 24 hours post-admission.
Each patient has an associated time of the end of follow-up T i and the type of observed event d i for which d i =signifies that the patient died, d i = 0 corresponds to right-censoring. 3.2 Components The proposed architecture decomposes clinical presence into three dimensions, modelled through three different neural networks. Each relies on the embedding h i,j outputted by a Long Short
Term Memory (LSTM) network [32] at the observation time t i,j .
Note that we input the series of observations at irregular times .
The embedding at the end of the observation time is then leveraged to model the survival outcome.
We now provide further details on the networks used for each of these tasks.
Longitudinal process
The latent state h i,j is leveraged through a multi-layer perceptron L to model the future value of each laboratory test to be observed after a period t from the prediction time t i,j . Assuming a Gaussian distribution for the observed data, the neural network L outputs the mean ( µ i ( t ) ) and variance ( σ i ( t ) ) of each future covariate given both the embedding ( h i,j ) and the longitudinal prediction time ( t ): [Equation]
Missingness process A second multi-layer perceptron M leverages h
i,j to model the likelihood of observing the different covariates, i.e. which tests are likely to be performed after the period t . One can assume a Bernoulli distribution parameterised by a neural network over the missingness patterns: [Equation] Temporal process A monotonic positive neural network I [33] leverages h i,j to model the recurrent patterns of observations.
A monotonic positive neural network is constrained to have positive weights and a final Softplus layer to guarantee the positivity and monotonicity of the outputted cumulative hazard I of observing any new laboratory test during the period t .
This approach expresses the probability of observing an event at any time ∆ after t minutes, while avoiding any parametric assumption: [Equation] Survival outcome Finally, to model the hazard function λ ( t ) , we used the DeepSurv model [34] that leverages a multi-layer perceptron S to extract a non-linear covariates shift used in a standard multiplicative proportional hazards Cox model. Given the baseline hazard λ 0 , the hazard of observing an event at horizon t is expressed as:
[Equation] The final model, therefore, combines state-of-the-art architectures in a novel way to jointly model clinical presence and survival. This results in a latent representation h i,j , which embeds the datagenerating process and the survival process.
Further description of the multi-task training strategy is provided in Appendix A.2.
3.3 Motivation
DeepJoint for improved performance.
The proposed method is motivated by joint models that incorporate informative processes into the outcome model through a shared effect [6, 19, 20].
Limited by poor scalability or strong parametric assumptions, we propose a multi-task neural network to tackle these challenges.
Evidence of improved performance [35] and generalisability [36, 37] have been shown in the multi-task literature.
However, theoretical foundations are still lacking [38] and rely on the intuition of regularisation of the shared embedding [39].
Our work builds a connection between joint modelling and multi-task learning by modelling the informative data-generating process and outcome of interest.
The aim is to detect if a new patient is not following the clinical presence patterns observed in the training set with the intent that this regularisation improves generalisability [36].
DeepJoint for robustness.
Under the assumption of negligible differences between clinical presence processes, i.e. similar medical practices, we provide intuition on why multi-task may improve robustness to small changes in clinical presence.
Specifically, we adapt a result from the adversarial attack literature [40] that shows the proportionality of the error to the inverse of the number of tasks modelled.
Assuming a bounded distance between marked point processes resulting from different clinical presence processes, one can adapt these theoretical results (see Appendix A.1 for proof and further discussion of its limitations).
Experiments
This paper studies the problem of in-hospital mortality prediction in the Medical Information Mart for Intensive Care III dataset (MIMIC III) [41] using laboratory tests after a 24-hour observation period in the intensive care unit.
The proposed implementations are available on Github 1 .
Data We choose to demonstrate the impact of clinical presence on laboratory tests only as other informative modalities might present different clinical presence patterns, e.g. semi-automatic vital signs collection.
After pre-processing [42] and subselection using an ECLAT algorithm [43], the resultant cohorts consist of 30,834 patients with 17 shared laboratory tests observed at least once in the observation period (See Tables 1 and 2 in Appendix).
Baselines We compare the proposed methodology ( DeepJoint ) against different methods for handling clinical presence.
All methods rely on the same normalised data imputed using last-observationscarried-forward with patient-mean imputation for the initial missing values.
All compared approaches extract a representation of the laboratory tests time series, then leveraged by a DeepSurv [34] architecture.
First, we compared against two non-sequential approaches: (i)
Last : Extract the last observation l i as representation for each patient.
(ii) Count :
Add the count of each test performed in the first 24 hours to the previous representation.
This assumes the informativeness of the counting process but ignores its temporal evolution.
Then, we used RNN-based approaches to take advantage of the longitudinal evolution of the laboratory tests: (iii) Ignore : An LSTM is trained on the imputed data leveraging the inputs’ temporal order but ignoring their irregularity and missingness patterns.
(iv) Resample : Data are re-sampled every hour.
(v) GRU-D : Data concatenated with missingness indicators serve as inputs to a GRU-D model [11] which leverages inter-observation times to decay the embedding.
(vi) Feature :
An LSTM leverages missingness indicators and time elapsed since previous observation [3] in addition to the lab results as inputs to model survival.
Finally, our proposed method (vii) DeepJoint leverages the same input as (vi) but explicitly models clinical presence.
Robustness to shift in clinical presence Motivated by the difference in the counts of tests performed on patients admitted on weekends and weekdays ( 2 . vs 1 . 91(0 . 99) tests for weekend admission, see Figure 4 in Appendix), we hypothesise that practice might differ between weekends and weekdays due to physicians and laboratory availability.
While providing an opportunity to study survival models’ transfer, the weekend effect [44] is marked by different outcome distributions (Overall mortality: 15 . 0% if admitted on weekend, vs. 13 . 5% if admitted on weekdays).
We adopt the evaluation methodology used in the robustness literature.
We stratify patients given days of admission (from Monday 8 am to Saturday 8 am ( n = 23 , 359 ) vs the rest ( n = 7 , 475 )).
Each group is further divided between training and testing.
A first model uses the training set of patients admitted on weekends and tested on the testing set of weekdays-admitted patients.
Then, a second model uses the training set of patients admitted on weekdays and tested on the testing set of weekdays-admitted patients.
The two resulting models are comparable as evaluated on the same testing set, but trained under two different clinical presence processes (See Figure 3 in Appendix for visualisation of the training strategy).
All methods use the same 90%-10% train-test patients split to train the network.
Their training relies on gradient back-propagation with an Adam optimiser [45] over 1000 epochs with early stopping on the survival loss computed on a left aside 10% of the training set, i.e. validation set.
The entire network is optimised for 500 epochs.
The remaining iterations are for fine-tuning the survival component.
We perform hyperparameter tuning on the validation set on the grid presented in Appendix Table 3.
We compute survival prediction at the last observation in the 24-hour postadmission period.
Following the survival literature [46], models are compared using time-dependent C-index [47] and Brier score [48] evaluated at time horizons of 1, 7 and 14 days after the observation period.
95% confidence intervals were obtained using 100 bootstrapped iterations on the testing set predicted values.
Results and Conclusions
Figure 2 presents the performance robustness when a model is transferred from one clinical presence setting to another.
A robust model has similar performances when transferred from another setting and when trained under the same clinical presence.
Practitioners should select a model close to the diagonal with the best discriminative performance ( upper right corner ).
Current clinical presence handling under-perform under shift.
Ignoring the clinical presence process leads to the worst performance and transportability with the naive re-sampling strategy underperforming in both settings.
The performance gap between Last and Count underlines that simple featurization may not be an adequate approach, as shown by the detrimental effect of the additional count of events on transportability.
Leveraging clinical presence improves performance with explicit modelling outperforming all other approaches.
Clinical presence reflects short-term instability.
Note that performances decay with longer horizons with shrinking differences between the different methodologies.
This observation shows both the task complexity in intensive care settings and the short-term relevance of clinical presence.
In conclusion, [1, 3] underline how taking advantage of missing data might lead to a mismatch between the training and deployment settings, but ignoring this information might not be possible.
Our work shows how leveraging clinical presence through modelling improves robustness to changes.
Nonetheless, practitioners should remain wary of informing medical practice with clinical presence as this process may embed not only informative patterns but also potential historical biases [49].
In future work, we aim to explore how our proposed method can detect changes in clinical presence, and under which changes to expect improved robustness.
Additionally, this work focuses on RNNbased architecture, additional state-of-the-art approaches could be considered.
Acknowledgments and Disclosure of Funding
This work has been partially funded by UKRI Medical Research Council (MC_UU_00002/5 and
MC_UU_00002/2)
A Appendix A.1 Multi task robustness to clinical presence shift
For completeness, we repeat and slightly adapt the results presented in [40] (see original paper for detailed proof).
These results motivate the use of multi-task modelling for robust modelling under shift.
First, we introduce the clinical presence vulnerability defined as the expected change in loss under two different data-generating processes.
Assuming the time series x ∈ f ( X ) , the target y , a loss
L , the data-generating process observed in the original dataset o s and the target one o t such that the difference between the marked point processes resulting from these processes lies in a p-norm bounded ball with radius r , i.e. ||
o s ( x ) − o t ( x ) || p < r , then the clinical presence vulnerability over the target dataset is:
[Equation]
Assuming infinitesimal changes in the data-generating processes, i.e. r → 0 , one can develop this expression using Taylor expansion:
[Equation]
Leveraging the property of the dual norm q , one can show that:
[Equation]
Theorem 1 (Theorem 2 from [40])
If the tasks are correlated with each other such that the covariance between the gradient of task i and task j is Cov ( r i , r j ) , and the gradient for each task is i.i.d.
with zero mean (because the model has converged), then clinical presence vulnerability of the given model is proportional to:
[Equation]
M where M is the number of output tasks selected.
Remark 1 Note that we assume a bounded distance between the marked point process from the source and the target processes, not between the observed and the underlying distributions.
This result may therefore suffer when the data-generating processes present strong dissimilarities but is robust to consistent mechanisms in the sampling processes.
From an application point of view, this means, for instance, that the theorem holds when practitioners have similar training and incentives, but may break if the healthcare recommendations differ, e.g. public vs. private systems.
A.2 DeepJoint Training
Our proposed methodology leverages the multi-task literature to maximise the clinical presence and survival likelihoods.
Each objective is back-propagated simultaneously by averaging the loss of the different tasks: survival, longitudinal, timing and missingness processes.
We ensure that no objective is over-represented in the average by using a dynamic weighting average scheme [50], as imbalances can greatly impact performance [51].
Each loss is weighted by its relative change at iteration e as follows:
[Equation]
which is then normalised between the three clinical presence tasks using a Softmax.
l task ( e ) is the validation log-likelihood at iteration e for the given task and θ is a temperature hyperparameter that controls softness, i.e. larger values would lead to equal weights.
Longitudinal process This first neural network aims to maximise the likelihood of the next observed laboratory test values at time
ϵ i,j .
The assumption of normality leads to the use of the
Gaussian log-likelihood loss:
[Equation]
with µ i,j and σ i,j the means, variances vectors outputted by L ( h i,j , ϵ i,j +1 ) and ∗ the element-wise multiplication.
Note the filtering through m i,j +1 to avoid a penalty on unobserved data.
Missingness process Similarly, the loss for the missingness process results in the binary cross entropy loss:
[Equation]
Temporal process Monontonic networks allow the exact computation of the likelihood without the need for distributional assumptions on the inter-observation distribution.
The model outputs the cumulative intensity at horizon t , and automatic differentiation results in the instantaneous intensity at no extra computational cost in the context of the neural networks’ training procedure.
[Equation]
Survival outcome DeepSurv relies on the proportional hazards assumption made by the Cox model resulting in the optimisation of the partial log-likelihood:
[Equation]
Finally, the population baseline hazard, λ 0 ( t ) , is estimated as a piece-wise constant based on the training population.
The final loss is defined at iteration e as
[Equation]
With α , a hyperparameter balancing between survival and clinical presence.
This loss is computed on a validation set for early stopping of the multi-task training.
Then, we perform a fine-tuning of the network S with all other weights fixed with early stopping.
A.3 MIMIC III - Experiments A.3.1 Data characteristics
Table 1 presents the demographic characteristics of the studied population and Table 2 summarises the set of tests selected with the mean number of tests performed during the 24 hours post-admission and their mean values.
The results are presented at the population level and differentiated by the subgroups used to study the impact of clinical presence shift.
[Table] [Table]
A.3.2 Hyperparameters tuning
All models hyper-parameters were selected over the following grid of hyperparameters (if appropriate) following 100 iterations of random search.
All experiments were run on a A100 GPU over 100 hours for each experiment.
[Table]
A.3.4 Modelling clinical presence
Table 4 presents the discriminative results on the testing set of a random split of the population for which patients are randomly assigned to train and test sets.
The proposed methodology presents similar performances to Feature which is based on the same inputs.
However, as shown Deepjoint presents better transportability.
Similarly, Count seems to achieve better than Last in this setting but presents strong overfitting when transferred.
[Table]
Figure 5 and Table 5 present the Brier score obtained on a random split of patients, showing little difference in the calibration of the different models at shorter time horizons.
Evaluation
[Table]
A.3.5 Transfer performances
Figure 6 shows the performance on patients admitted on weekends.
This echoes the conclusion made in Section 5 that the proposed model is more robust to change in the clinical presence at shorter time horizons.
However, this advantage fades as clinical presence reflects short-term changes.
Tables 6, 7 and 8 present the discriminative difference on the weekends-admissions testing set between the model trained on weekends-admissions and the model transferred from weekdays to weekends
At short time horizons, methodologies that leverage clinical presence present more robust results.
Note that the proposed methodology presents better performance than Feature based on the same inputs.
Additionally, for weekends-admissions evaluation, GRU-D presents state-of-the-art predictive performance and robustness but this result didn’t generalise to the other settings.
Interestingly, ignoring the clinical process ( Resample or Ignore ) seems to be less robust to a change in the data-generating process, echoing the remark made by [3] about the difficulty of ignoring it
[Table]
[Table] [Table] [Table]